import gradio as gr
import asyncio
import websockets
import json
import uuid
import argparse
import urllib.parse
from datetime import datetime
import logging
import sys
import os
import shutil
from pathlib import Path
import time

# 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")

# Dictionary to store active connections
active_connections = {}
# Dictionary to store message history for each chat room (in-memory cache)
chat_history = {}
# Dictionary to store message history for logs
log_history = []

# Dictionary to track file modification times
file_modification_times = {}

# Dictionary to track users in each room/sector
sector_users = {}

# Grid dimensions for 2D sector map
GRID_WIDTH = 10
GRID_HEIGHT = 10

# Directories for persistent storage
HISTORY_DIR = "chat_history"
LOG_DIR = "server_logs"
os.makedirs(HISTORY_DIR, exist_ok=True)
os.makedirs(LOG_DIR, exist_ok=True)

# README files
README_PATH = os.path.join(HISTORY_DIR, "README.md")
if not os.path.exists(README_PATH):
    with open(README_PATH, "w") as f:
        f.write("# Chat History\n\nThis directory contains persistent chat history files.\n")

LOG_README_PATH = os.path.join(LOG_DIR, "README.md")
if not os.path.exists(LOG_README_PATH):
    with open(LOG_README_PATH, "w") as f:
        f.write("# Server Logs\n\nThis directory contains server log files.\n")

# Get node name from URL or command line
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()
    node_name = args.node_name
    port = args.port
    
    if not node_name:
        node_name = f"node-{uuid.uuid4().hex[:8]}"
        
    return node_name, port

def get_room_history_file(room_id):
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    return os.path.join(HISTORY_DIR, f"{room_id}_{timestamp}.jsonl")

def get_log_file():
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    return os.path.join(LOG_DIR, f"log_{timestamp}.jsonl")

def get_all_room_history_files(room_id):
    files = []
    for file in os.listdir(HISTORY_DIR):
        if file.startswith(f"{room_id}_") and file.endswith(".jsonl"):
            files.append(os.path.join(HISTORY_DIR, file))
    files.sort(key=lambda x: os.path.getmtime(x), reverse=True)
    return files

def list_all_history_files():
    """List all history files with room IDs and modification times."""
    files = []
    for file in os.listdir(HISTORY_DIR):
        if file.endswith(".jsonl") and file != "README.md":
            parts = file.split('_', 1)
            if len(parts) > 0:
                room_id = parts[0]
                file_path = os.path.join(HISTORY_DIR, file)
                mod_time = os.path.getmtime(file_path)
                files.append((room_id, file_path, mod_time))
    files.sort(key=lambda x: x[2], reverse=True)  # Sort by modification time
    return files

def load_room_history(room_id):
    if room_id not in chat_history:
        chat_history[room_id] = []
        history_files = get_all_room_history_files(room_id)
        
        for file in history_files:
            if file not in file_modification_times:
                file_modification_times[file] = os.path.getmtime(file)
        
        messages = []
        for history_file in history_files:
            try:
                with open(history_file, 'r') as f:
                    for line in f:
                        line = line.strip()
                        if line:
                            try:
                                data = json.loads(line)
                                messages.append(data)
                            except json.JSONDecodeError:
                                logger.error(f"Error parsing JSON line in {history_file}")
            except Exception as e:
                logger.error(f"Error loading history from {history_file}: {e}")
        
        messages.sort(key=lambda x: x.get("timestamp", ""), reverse=False)
        chat_history[room_id] = messages
        logger.info(f"Loaded {len(messages)} messages from {len(history_files)} files for room {room_id}")
    
    if room_id not in sector_users:
        sector_users[room_id] = set()
    
    return chat_history[room_id]

def save_message_to_history(room_id, message):
    history_files = get_all_room_history_files(room_id)
    if not history_files:
        history_file = get_room_history_file(room_id)
    else:
        newest_file = history_files[0]
        if os.path.getsize(newest_file) > 1024 * 1024:  # 1 MB
            history_file = get_room_history_file(room_id)
        else:
            history_file = newest_file
    
    try:
        with open(history_file, 'a') as f:
            f.write(json.dumps(message) + '\n')
        file_modification_times[history_file] = os.path.getmtime(history_file)
        logger.debug(f"Saved message to {history_file}")
    except Exception as e:
        logger.error(f"Error saving message to {history_file}: {e}")

def save_log_entry(entry):
    """Save a log entry to the newest log file."""
    log_files = [f for f in os.listdir(LOG_DIR) if f.startswith("log_") and f.endswith(".jsonl")]
    log_files.sort(key=lambda x: os.path.getmtime(os.path.join(LOG_DIR, x)), reverse=True)
    
    if not log_files:
        log_file = get_log_file()
    else:
        newest_log = os.path.join(LOG_DIR, log_files[0])
        if os.path.getsize(newest_log) > 1024 * 1024:  # 1 MB
            log_file = get_log_file()
        else:
            log_file = newest_log
    
    try:
        with open(log_file, 'a') as f:
            f.write(json.dumps(entry) + '\n')
        file_modification_times[log_file] = os.path.getmtime(log_file)
    except Exception as e:
        logger.error(f"Error saving log to {log_file}: {e}")

def check_for_new_messages():
    updated_rooms = set()
    for file in os.listdir(HISTORY_DIR):
        if file.endswith(".jsonl"):
            file_path = os.path.join(HISTORY_DIR, file)
            current_mtime = os.path.getmtime(file_path)
            if file_path not in file_modification_times or current_mtime > file_modification_times[file_path]:
                parts = file.split('_', 1)
                if len(parts) > 0:
                    room_id = parts[0]
                    updated_rooms.add(room_id)
                file_modification_times[file_path] = current_mtime
    
    for room_id in updated_rooms:
        if room_id in chat_history:
            old_history_len = len(chat_history[room_id])
            chat_history[room_id] = []
            load_room_history(room_id)
            new_history_len = len(chat_history[room_id])
            if new_history_len > old_history_len:
                logger.info(f"Found {new_history_len - old_history_len} new messages for room {room_id}")
    
    return updated_rooms

def check_for_new_logs():
    """Check for new log entries and broadcast them."""
    global log_history
    updated = False
    
    for file in os.listdir(LOG_DIR):
        if file.endswith(".jsonl"):
            file_path = os.path.join(LOG_DIR, file)
            current_mtime = os.path.getmtime(file_path)
            if file_path not in file_modification_times or current_mtime > file_modification_times[file_path]:
                updated = True
                file_modification_times[file_path] = current_mtime
    
    if updated:
        log_history = []
        for file in sorted([f for f in os.listdir(LOG_DIR) if f.endswith(".jsonl")], key=lambda x: os.path.getmtime(os.path.join(LOG_DIR, x))):
            file_path = os.path.join(LOG_DIR, file)
            try:
                with open(file_path, 'r') as f:
                    for line in f:
                        line = line.strip()
                        if line:
                            try:
                                log_history.append(json.loads(line))
                            except json.JSONDecodeError:
                                logger.error(f"Error parsing log line in {file_path}")
            except Exception as e:
                logger.error(f"Error loading logs from {file_path}: {e}")
        
        # Broadcast logs to all connected clients in a "logs" room
        log_msg = {
            "type": "log",
            "content": "\n".join([entry["message"] for entry in log_history[-10:]]),  # Last 10 entries
            "timestamp": datetime.now().isoformat(),
            "sender": "system",
            "room_id": "logs"
        }
        asyncio.create_task(broadcast_message(log_msg, "logs"))

def get_sector_coordinates(room_id):
    try:
        if ',' in room_id:
            x, y = map(int, room_id.split(','))
            return max(0, min(x, GRID_WIDTH-1)), max(0, min(y, GRID_HEIGHT-1))
    except:
        pass
    hash_val = hash(room_id)
    x = abs(hash_val) % GRID_WIDTH
    y = abs(hash_val >> 8) % GRID_HEIGHT
    return x, y

def generate_sector_map():
    grid = [[' ' for _ in range(GRID_WIDTH)] for _ in range(GRID_HEIGHT)]
    for room_id, users in sector_users.items():
        if users:
            x, y = get_sector_coordinates(room_id)
            user_count = len(users)
            grid[y][x] = str(min(user_count, 9)) if user_count < 10 else '+'
    
    header = '  ' + ''.join([str(i % 10) for i in range(GRID_WIDTH)])
    map_str = header + '\n'
    for y in range(GRID_HEIGHT):
        row = f"{y % 10}|" + ''.join(grid[y]) + '|'
        map_str += row + '\n'
    map_str += header
    return f"```\n{map_str}\n```\n\nLegend: Number indicates users in sector. '+' means 10+ users."

async def clear_all_history():
    global chat_history
    chat_history = {}
    for file in os.listdir(HISTORY_DIR):
        if file.endswith(".jsonl"):
            os.remove(os.path.join(HISTORY_DIR, file))
    
    clear_msg = {
        "type": "system",
        "content": "🧹 All chat history has been cleared by a user",
        "timestamp": datetime.now().isoformat(),
        "sender": "system"
    }
    
    for room_id in list(active_connections.keys()):
        clear_msg["room_id"] = room_id
        await broadcast_message(clear_msg, room_id)
    
    logger.info("All chat history cleared")
    return "All chat history cleared"

async def websocket_handler(websocket, path):
    try:
        path_parts = path.strip('/').split('/')
        room_id = path_parts[0] if path_parts else "default"
        
        client_id = str(uuid.uuid4())
        if room_id not in active_connections:
            active_connections[room_id] = {}
        active_connections[room_id][client_id] = websocket
        
        if room_id not in sector_users:
            sector_users[room_id] = set()
        sector_users[room_id].add(client_id)
        
        x, y = get_sector_coordinates(room_id)
        room_history = load_room_history(room_id)
        
        welcome_msg = {
            "type": "system",
            "content": f"Welcome to room '{room_id}' (Sector {x},{y})! Connected from node '{NODE_NAME}'",
            "timestamp": datetime.now().isoformat(),
            "sender": "system",
            "room_id": room_id
        }
        await websocket.send(json.dumps(welcome_msg))
        
        map_msg = {
            "type": "system",
            "content": f"Sector Map:\n{generate_sector_map()}",
            "timestamp": datetime.now().isoformat(),
            "sender": "system",
            "room_id": room_id
        }
        await websocket.send(json.dumps(map_msg))
        
        for msg in room_history:
            await websocket.send(json.dumps(msg))
        
        join_msg = {
            "type": "system",
            "content": f"User joined the room (Sector {x},{y}) - {len(sector_users[room_id])} users now present",
            "timestamp": datetime.now().isoformat(),
            "sender": "system",
            "room_id": room_id
        }
        await broadcast_message(join_msg, room_id)
        save_message_to_history(room_id, join_msg)
        
        logger.info(f"New client {client_id} connected to room {room_id} (Sector {x},{y})")
        
        if room_id == "logs":
            for log_entry in log_history[-10:]:  # Send last 10 log entries
                log_msg = {
                    "type": "log",
                    "content": log_entry["message"],
                    "timestamp": log_entry["timestamp"],
                    "sender": "system",
                    "room_id": "logs"
                }
                await websocket.send(json.dumps(log_msg))
        
        async for message in websocket:
            try:
                data = json.loads(message)
                
                if data.get("type") == "command" and data.get("command") == "clear_history":
                    result = await clear_all_history()
                    continue
                
                if data.get("type") == "command" and data.get("command") == "show_map":
                    map_msg = {
                        "type": "system",
                        "content": f"Sector Map:\n{generate_sector_map()}",
                        "timestamp": datetime.now().isoformat(),
                        "sender": "system",
                        "room_id": room_id
                    }
                    await websocket.send(json.dumps(map_msg))
                    continue
                
                data["timestamp"] = datetime.now().isoformat()
                data["sender_node"] = NODE_NAME
                data["room_id"] = room_id
                
                chat_history[room_id].append(data)
                if len(chat_history[room_id]) > 500:
                    chat_history[room_id] = chat_history[room_id][-500:]
                
                save_message_to_history(room_id, data)
                await broadcast_message(data, room_id)
                
            except json.JSONDecodeError:
                error_msg = {
                    "type": "error",
                    "content": "Invalid JSON format",
                    "timestamp": datetime.now().isoformat(),
                    "sender": "system",
                    "room_id": room_id
                }
                await websocket.send(json.dumps(error_msg))
    
    except websockets.exceptions.ConnectionClosed:
        logger.info(f"Client {client_id} disconnected from room {room_id}")
    finally:
        if room_id in active_connections and client_id in active_connections[room_id]:
            del active_connections[room_id][client_id]
            if room_id in sector_users and client_id in sector_users[room_id]:
                sector_users[room_id].remove(client_id)
            
            x, y = get_sector_coordinates(room_id)
            leave_msg = {
                "type": "system",
                "content": f"User left the room (Sector {x},{y}) - {len(sector_users.get(room_id, set()))} users remaining",
                "timestamp": datetime.now().isoformat(),
                "sender": "system",
                "room_id": room_id
            }
            await broadcast_message(leave_msg, room_id)
            save_message_to_history(room_id, leave_msg)
            
            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_clients = []
        for client_id, websocket in active_connections[room_id].items():
            try:
                await websocket.send(json.dumps(message))
            except websockets.exceptions.ConnectionClosed:
                disconnected_clients.append(client_id)
        
        for client_id in disconnected_clients:
            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

main_event_loop = None
message_queue = []

def send_message(message, username, room_id):
    if not message.strip():
        return None
    
    global message_queue
    msg_data = {
        "type": "chat",
        "content": message,
        "username": username,
        "room_id": room_id
    }
    message_queue.append(msg_data)
    formatted_msg = f"{username}: {message}"
    return formatted_msg

def join_room(room_id, chat_history_output):
    if not room_id.strip():
        return "Please enter a valid room ID", chat_history_output
    
    room_id = urllib.parse.quote(room_id.strip())
    history = load_room_history(room_id)
    
    formatted_history = []
    for msg in history:
        if msg.get("type") == "chat":
            sender_node = f" [{msg.get('sender_node', 'unknown')}]" if "sender_node" in msg else ""
            time_str = ""
            if "timestamp" in msg:
                try:
                    dt = datetime.fromisoformat(msg["timestamp"])
                    time_str = f"[{dt.strftime('%H:%M:%S')}] "
                except:
                    pass
            formatted_history.append(f"{time_str}{msg.get('username', 'Anonymous')}{sender_node}: {msg.get('content', '')}")
        elif msg.get("type") == "system":
            formatted_history.append(f"System: {msg.get('content', '')}")
    
    return f"Joined room: {room_id}", formatted_history

def send_clear_command():
    global message_queue
    msg_data = {
        "type": "command",
        "command": "clear_history",
        "username": "System"
    }
    message_queue.append(msg_data)
    return "🧹 Clearing all chat history..."

def list_available_rooms():
    history_files = list_all_history_files()
    
    if not history_files:
        return "No chat rooms available yet. Create one by joining a room!"
    
    room_list = "### Available Chat Rooms\n\n"
    for room_id, file_path, mod_time in history_files:
        last_activity = datetime.fromtimestamp(mod_time).strftime("%Y-%m-%d %H:%M:%S")
        room_list += f"- **{room_id}**: Last activity {last_activity}\n"
    
    return room_list

def create_gradio_interface():
    with gr.Blocks(title=f"Chat Node: {NODE_NAME}") as interface:
        gr.Markdown(f"# Chat Node: {NODE_NAME}")
        gr.Markdown("Join a room by entering a room ID below or create a new one.")
        
        with gr.Row():
            with gr.Column(scale=3):
                room_list = gr.Markdown(value="Loading available rooms...")
                refresh_button = gr.Button("🔄 Refresh Room List")
            with gr.Column(scale=1):
                clear_button = gr.Button("🧹 Clear All Chat History", variant="stop")
        
        with gr.Row():
            with gr.Column(scale=2):
                room_id_input = gr.Textbox(label="Room ID", placeholder="Enter room ID or use x,y coordinates")
                join_button = gr.Button("Join Room")
            with gr.Column(scale=1):
                with gr.Row():
                    x_coord = gr.Number(label="X", value=0, minimum=0, maximum=GRID_WIDTH-1, step=1)
                    y_coord = gr.Number(label="Y", value=0, minimum=0, maximum=GRID_HEIGHT-1, step=1)
                grid_join_button = gr.Button("Join by Coordinates")
        
        chat_history_output = gr.Textbox(label="Chat History", lines=20, max_lines=20)
        
        with gr.Row():
            username_input = gr.Textbox(label="Username", placeholder="Enter your username", value="User")
            with gr.Column(scale=3):
                message_input = gr.Textbox(
                    label="Message",
                    placeholder="Type your message here. Press Shift+Enter for new line, Enter to send.",
                    lines=3
                )
            with gr.Column(scale=1):
                send_button = gr.Button("Send")
                map_button = gr.Button("🗺️ Show Map")
        
        current_room_display = gr.Textbox(label="Current Room", value="Not joined any room yet")
        
        refresh_button.click(list_available_rooms, inputs=[], outputs=[room_list])
        clear_button.click(send_clear_command, inputs=[], outputs=[room_list])
        
        def join_by_coordinates(x, y):
            return f"{int(x)},{int(y)}"
        
        grid_join_button.click(join_by_coordinates, inputs=[x_coord, y_coord], outputs=[room_id_input]).then(
            join_room, inputs=[room_id_input, chat_history_output], outputs=[current_room_display, chat_history_output]
        )
        
        join_button.click(join_room, inputs=[room_id_input, chat_history_output], outputs=[current_room_display, chat_history_output])
        
        def send_and_clear(message, username, room_id):
            if not room_id.startswith("Joined room:"):
                return "Please join a room first", message
            
            actual_room_id = room_id.replace("Joined room: ", "").strip()
            message_lines = message.strip().split("\n")
            formatted_msg = ""
            
            for line in message_lines:
                if line.strip():
                    sent_msg = send_message(line.strip(), username, actual_room_id)
                    if sent_msg:
                        formatted_msg += sent_msg + "\n"
            
            return "", formatted_msg if formatted_msg else message, None
        
        send_button.click(
            send_and_clear,
            inputs=[message_input, username_input, current_room_display],
            outputs=[message_input, chat_history_output]
        )
        
        def show_sector_map(room_id):
            if not room_id.startswith("Joined room:"):
                return "Please join a room first to view the map"
            return generate_sector_map()
        
        map_button.click(show_sector_map, inputs=[current_room_display], outputs=[chat_history_output])
        
        def on_message_submit(message, username, room_id):
            return send_and_clear(message, username, room_id)
        
        message_input.submit(
            on_message_submit,
            inputs=[message_input, username_input, current_room_display],
            outputs=[message_input, chat_history_output]
        )
        
        interface.load(list_available_rooms, inputs=[], outputs=[room_list])
    
    return interface

async def process_message_queue():
    global message_queue
    while True:
        if message_queue:
            msg_data = message_queue.pop(0)
            await broadcast_message(msg_data, msg_data["room_id"])
        await asyncio.sleep(0.1)

async def process_logs():
    """Periodically check and broadcast new log entries."""
    while True:
        check_for_new_logs()
        await asyncio.sleep(1)  # Check every second

# Custom logging handler to save logs and broadcast them
class LogBroadcastHandler(logging.Handler):
    def emit(self, record):
        log_entry = {
            "timestamp": datetime.now().isoformat(),
            "level": record.levelname,
            "message": self.format(record),
            "name": record.name
        }
        save_log_entry(log_entry)
        log_history.append(log_entry)

logger.addHandler(LogBroadcastHandler())

async def main():
    global NODE_NAME, main_event_loop
    NODE_NAME, port = get_node_name()
    main_event_loop = asyncio.get_running_loop()
    
    server = await start_websocket_server()
    asyncio.create_task(process_message_queue())
    asyncio.create_task(process_logs())  # Start log processor
    
    interface = create_gradio_interface()
    
    from starlette.middleware.base import BaseHTTPMiddleware
    
    class NodeNameMiddleware(BaseHTTPMiddleware):
        async def dispatch(self, request, call_next):
            global NODE_NAME
            query_params = dict(request.query_params)
            if "node_name" in query_params:
                NODE_NAME = query_params["node_name"]
                logger.info(f"Node name set to {NODE_NAME} from URL parameter")
            response = await call_next(request)
            return response
    
    app = gr.routes.App.create_app(interface)
    app.add_middleware(NodeNameMiddleware)
    
    import uvicorn
    config = uvicorn.Config(app, host="0.0.0.0", port=port)
    server = uvicorn.Server(config)
    
    logger.info(f"Starting Gradio interface on http://0.0.0.0:{port} with node name '{NODE_NAME}'")
    await server.serve()

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