import gradio as gr
import asyncio
import websockets
import json
import uuid
import argparse
from datetime import datetime
import logging
import sys
import os
import random

# Configure logging - the gossip queen 📢👑 sets up the chatterbox!
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger("chat-node")

# Fun usernames with emojis - the VIP list of quirky characters! 🎉😜
FUN_USERNAMES = [
    "CosmicJester 🌌", "PixelPanda 🐼", "QuantumQuack 🦆", "StellarSquirrel 🐿️",
    "GizmoGuru ⚙️", "NebulaNinja 🌠", "ByteBuster 💾", "GalacticGopher 🌍",
    "RocketRaccoon 🚀", "EchoElf 🧝", "PhantomFox 🦊", "WittyWizard 🧙",
    "LunarLlama 🌙", "SolarSloth ☀️", "AstroAlpaca 🦙", "CyberCoyote 🐺",
    "MysticMoose 🦌", "GlitchGnome 🧚", "VortexViper 🐍", "ChronoChimp 🐒"
]

# Directory for logs - the secret vault where tales are stashed! 🗄️🔒
LOG_DIR = "user_logs"
os.makedirs(LOG_DIR, exist_ok=True)

# Node name - the app’s codename generator, sneaky and slick! 🕵️‍♂️💾
def get_node_name():
    """🎲 Spins the wheel of fate to name our node - a random alias or user pick! 🏷️"""
    parser = argparse.ArgumentParser(description='Start a chat node with a specific name')
    parser.add_argument('--node-name', type=str, default=None, help='Name for this chat node')
    parser.add_argument('--port', type=int, default=7860, help='Port to run the Gradio interface on')
    args = parser.parse_args()
    return args.node_name or f"node-{uuid.uuid4().hex[:8]}", args.port

# Log file generator - crafts a shiny new logbook with a timestamp twist! 📅✍️
def get_log_file(username):
    """⏰ Stamps the clock and names the log after our witty hero - fresh and ready! 📜"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    safe_username = username.replace(" ", "_").encode('ascii', 'ignore').decode('ascii')
    return os.path.join(LOG_DIR, f"{safe_username}_{timestamp}.md")

# Log saver - the scribe that pens epic log tales into Markdown glory! 🖋️📚
def save_log_entry(username, entry):
    """🖌️ Scribbles a log entry into a Markdown masterpiece - tales of triumph and woe! 🎭"""
    log_file = get_log_file(username)
    try:
        with open(log_file, 'a') as f:
            f.write(f"```log\n[{entry['timestamp']}] {entry['level']}: {entry['message']}\n```\n")
        logger.info(f"Saved log entry to {log_file}")
        return log_file
    except Exception as e:
        logger.error(f"Error saving log to {log_file}: {e}")
        return None

# Log loader - the treasure hunter digging up the latest log loot! 🏴‍☠️💰
def load_latest_log(username):
    """🔍 Hunts down the freshest log file and dresses it up with line numbers - snazzy! ✨"""
    safe_username = username.replace(" ", "_").encode('ascii', 'ignore').decode('ascii')
    log_files = [f for f in os.listdir(LOG_DIR) if f.startswith(safe_username) and f.endswith(".md")]
    if not log_files:
        return "# No logs yet\nStart interacting to generate logs!"
    latest_file = max(log_files, key=lambda f: os.path.getmtime(os.path.join(LOG_DIR, f)))
    try:
        with open(os.path.join(LOG_DIR, latest_file), 'r') as f:
            content = f.read()
        lines = content.strip().split('\n')
        numbered_content = "\n".join(f"{i+1}. {line}" for i, line in enumerate(lines))
        return f"# Log for {username} ({latest_file})\n\n{numbered_content}"
    except Exception as e:
        logger.error(f"Error loading log {latest_file}: {e}")
        return "# Error loading log\nCheck server logs for details."

active_connections = {}

# WebSocket handler - the bouncer at the log party, keeping the vibe alive! 🎉🚪
async def websocket_handler(websocket, path):
    """🎧 Listens at the door, letting log fans in and kicking out crashers! 🚨"""
    try:
        client_id = str(uuid.uuid4())
        room_id = "logs"
        active_connections.setdefault(room_id, {})[client_id] = websocket
        logger.info(f"Client {client_id} connected to logs")

        async for message in websocket:
            try:
                data = json.loads(message)
                data["timestamp"] = datetime.now().isoformat()
                await broadcast_message(data, room_id)
            except json.JSONDecodeError:
                await websocket.send(json.dumps({"type": "error", "content": "Invalid JSON", "timestamp": datetime.now().isoformat(), "sender": "system"}))
    except websockets.ConnectionClosed:
        logger.info(f"Client {client_id} disconnected from logs")
    finally:
        if room_id in active_connections and client_id in active_connections[room_id]:
            del active_connections[room_id][client_id]
            if not active_connections[room_id]:
                del active_connections[room_id]

# Broadcaster - the town crier shouting log updates to all who’ll listen! 📣🌍
async def broadcast_message(message, room_id):
    """📢 Blasts the latest log scoop to every ear in the room - no one misses out! 🎙️"""
    if room_id in active_connections:
        disconnected = []
        for client_id, ws in active_connections[room_id].items():
            try:
                await ws.send(json.dumps(message))
            except websockets.ConnectionClosed:
                disconnected.append(client_id)
        for client_id in disconnected:
            del active_connections[room_id][client_id]

# WebSocket server starter - lights the fuse on the log-streaming rocket! 🚀🔥
async def start_websocket_server(host='0.0.0.0', port=8765):
    """🌐 Fires up the WebSocket engine, ready to zip logs at warp speed! ⚡"""
    server = await websockets.serve(websocket_handler, host, port)
    logger.info(f"WebSocket server started on ws://{host}:{port}")
    return server

# Log handler - the sneaky spy logging every move with a witty twist! 🕵️‍♀️📝
class LogBroadcastHandler(logging.Handler):
    def __init__(self, username):
        super().__init__()
        self.username = username

    def emit(self, record):
        """🕵️‍♂️ Catches every whisper and scribbles it into the logbook with flair! ✍️"""
        entry = {
            "timestamp": datetime.now().isoformat(),
            "level": record.levelname,
            "message": self.format(record),
            "name": record.name
        }
        log_file = save_log_entry(self.username, entry)
        if log_file:
            asyncio.create_task(broadcast_message({
                "type": "log",
                "content": f"New log entry saved to {log_file}",
                "timestamp": entry["timestamp"],
                "sender": "system",
                "username": self.username
            }, "logs"))

# Interface maker - the artist painting a live log masterpiece! 🎨🖼️
def create_gradio_interface(username):
    """🖌️ Whips up a snazzy UI canvas that updates live with log magic! 🌟"""
    logger.handlers = [LogBroadcastHandler(username)]
    def update_log_display(dummy_input):
        """🎭 The puppet master pulling fresh log strings for the show! 🎬"""
        return load_latest_log(username)

    with gr.Blocks(title=f"Log Viewer: {NODE_NAME}", css=".code-container { font-family: monospace; background: #1e1e1e; color: #d4d4d4; padding: 10px; border-radius: 5px; }") as interface:
        gr.Markdown(f"# Log Viewer for {username} on Node: {NODE_NAME}")
        gr.Markdown("Logs update live as new entries are written - watch the action unfold! 🎥")
        dummy_input = gr.Textbox(visible=False, value="trigger", interactive=True)
        log_display = gr.Code(label="Your Latest Log", language="markdown", lines=20, elem_classes=["code-container"])

        # Live interface FTW! 🚀
        interface = gr.Interface(
            fn=update_log_display,
            inputs=[dummy_input],
            outputs=[log_display],
            live=True,
            title=f"Log Viewer: {NODE_NAME}"
        )
    return interface

# Main event - the ringmaster kicking off the log circus! 🎪🤡
async def main():
    """🎤 Drops the mic and starts the log party - it’s showtime, folks! 🎉"""
    global NODE_NAME
    NODE_NAME, port = get_node_name()
    await start_websocket_server()

    username = random.choice(FUN_USERNAMES)
    logger.info(f"Assigned username: {username}")

    interface = create_gradio_interface(username)
    import uvicorn
    await uvicorn.Server(uvicorn.Config(interface, host="0.0.0.0", port=port)).serve()

if __name__ == "__main__":
    asyncio.run(main())