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
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
FUN_USERNAMES = [
    "CosmicJester 🌌", "PixelPanda 🐼", "QuantumQuack 🦆", "StellarSquirrel 🐿️",
    "GizmoGuru ⚙️", "NebulaNinja 🌠", "ByteBuster 💾", "GalacticGopher 🌍",
    "RocketRaccoon 🚀", "EchoElf 🧝", "PhantomFox 🦊", "WittyWizard 🧙",
    "LunarLlama 🌙", "SolarSloth ☀️", "AstroAlpaca 🦙", "CyberCoyote 🐺",
    "MysticMoose 🦌", "GlitchGnome 🧚", "VortexViper 🐍", "ChronoChimp 🐒"
]

# Directories
LOG_DIR = "user_logs"
os.makedirs(LOG_DIR, exist_ok=True)

# Node name
def get_node_name():
    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

def get_log_file(username):
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    # Replace spaces and emojis with underscores for valid filenames
    safe_username = username.replace(" ", "_").encode('ascii', 'ignore').decode('ascii')
    return os.path.join(LOG_DIR, f"{safe_username}_{timestamp}.md")

def save_log_entry(username, entry):
    log_file = get_log_file(username)
    try:
        with open(log_file, 'a') as f:
            # Format as Markdown with timestamp
            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

def load_latest_log(username):
    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()
        # Add line numbers to the Markdown content
        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 = {}

async def websocket_handler(websocket, path):
    try:
        client_id = str(uuid.uuid4())
        room_id = "logs"  # Simplified to a single logs room
        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]

async def broadcast_message(message, room_id):
    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]

async def start_websocket_server(host='0.0.0.0', port=8765):
    server = await websockets.serve(websocket_handler, host, port)
    logger.info(f"WebSocket server started on ws://{host}:{port}")
    return server

class LogBroadcastHandler(logging.Handler):
    def __init__(self, username):
        super().__init__()
        self.username = username

    def emit(self, record):
        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:
            # Notify WebSocket clients of new 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"))

def create_gradio_interface(username):
    logger.handlers = [LogBroadcastHandler(username)]  # Replace handlers with user-specific one
    with gr.Blocks(title=f"Log Viewer: {NODE_NAME}") as interface:
        interface.css = """
            .code-container { font-family: monospace; background: #1e1e1e; color: #d4d4d4; padding: 10px; border-radius: 5px; }
        """
        interface.js = f"""
        () => {{
            setInterval(async () => {{
                const response = await fetch('/get_log?username={username}');
                const logContent = await response.text();
                document.getElementById('log-display').innerText = logContent;
            }}, 900); // Refresh every 0.9 seconds
        }}
        """

        gr.Markdown(f"# Log Viewer for {username} on Node: {NODE_NAME}")
        gr.Markdown("Your logs are displayed below, auto-refreshing every 0.9 seconds.")
        log_display = gr.Code(label="Your Latest Log", language="markdown", lines=20, elem_classes=["code-container"], elem_id="log-display")

        # Custom endpoint to fetch log
        def get_log():
            return load_latest_log(username)
        interface.load(get_log, [], [log_display])

    return interface

async def main():
    global NODE_NAME
    NODE_NAME, port = get_node_name()
    await start_websocket_server()

    # Assign a random username on startup
    username = random.choice(FUN_USERNAMES)
    logger.info(f"Assigned username: {username}")

    # Create and run Gradio interface
    interface = create_gradio_interface(username)
    
    from starlette.routing import Route
    from starlette.responses import PlainTextResponse
    
    async def get_log_endpoint(request):
        return PlainTextResponse(load_latest_log(username))
    
    app = gr.routes.App.create_app(interface)
    app.routes.append(Route("/get_log", get_log_endpoint))
    
    import uvicorn
    await uvicorn.Server(uvicorn.Config(app, host="0.0.0.0", port=port)).serve()

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