app.py

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

# Dictionaries for state
active_connections = {}
chat_history = {}
log_history = []
file_modification_times = {}
sector_users = {}

# Grid dimensions
GRID_WIDTH = 10
GRID_HEIGHT = 10

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

# README files
for dir_path, content in [
    (HISTORY_DIR, "# Chat History\n\nThis directory contains persistent chat history files.\n"),
    (LOG_DIR, "# Server Logs\n\nThis directory contains server log files.\n")
]:
    readme_path = os.path.join(dir_path, "README.md")
    if not os.path.exists(readme_path):
        with open(readme_path, "w") as f:
            f.write(content)

# 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 🐒"
]

# 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()
    node_name = args.node_name or f"node-{uuid.uuid4().hex[:8]}"
    return node_name, args.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 = [os.path.join(HISTORY_DIR, f) for f in os.listdir(HISTORY_DIR) if f.startswith(f"{room_id}_") and f.endswith(".jsonl")]
    files.sort(key=os.path.getmtime, reverse=True)
    return files

def list_all_history_files():
    files = []
    for file in os.listdir(HISTORY_DIR):
        if file.endswith(".jsonl") and file != "README.md":
            room_id = file.split('_', 1)[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)
    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:
            file_modification_times[file] = file_modification_times.get(file, os.path.getmtime(file))
            try:
                with open(file, 'r') as f:
                    for line in f:
                        if line.strip():
                            try:
                                chat_history[room_id].append(json.loads(line))
                            except json.JSONDecodeError:
                                logger.error(f"Error parsing JSON line in {file}")
            except Exception as e:
                logger.error(f"Error loading history from {file}: {e}")
        chat_history[room_id].sort(key=lambda x: x.get("timestamp", ""))
        logger.info(f"Loaded {len(chat_history[room_id])} messages from {len(history_files)} files for room {room_id}")
    sector_users[room_id] = sector_users.get(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)
    history_file = get_room_history_file(room_id) if not history_files or os.path.getsize(history_files[0]) > 1024 * 1024 else history_files[0]
    try:
        with open(history_file, 'a') as f:
            f.write(json.dumps(message) + '\n')
        file_modification_times[history_file] = os.path.getmtime(history_file)
    except Exception as e:
        logger.error(f"Error saving message to {history_file}: {e}")

def save_log_entry(entry):
    log_files = [os.path.join(LOG_DIR, 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(x), reverse=True)
    log_file = get_log_file() if not log_files or os.path.getsize(log_files[0]) > 1024 * 1024 else log_files[0]
    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 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:
        hash_val = hash(room_id)
        return abs(hash_val) % GRID_WIDTH, abs(hash_val >> 8) % GRID_HEIGHT

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' + '\n'.join(f"{y % 10}|{''.join(grid[y])}|" for y in range(GRID_HEIGHT)) + '\n' + 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 cleared", "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:
        room_id = path.strip('/').split('/')[0] or "default"
        client_id = str(uuid.uuid4())
        active_connections.setdefault(room_id, {})[client_id] = websocket
        sector_users.setdefault(room_id, set()).add(client_id)
        x, y = get_sector_coordinates(room_id)
        room_history = load_room_history(room_id)

        for msg in [
            {"type": "system", "content": f"Welcome to room '{room_id}' (Sector {x},{y})! Node: '{NODE_NAME}'", "timestamp": datetime.now().isoformat(), "sender": "system", "room_id": room_id},
            {"type": "system", "content": f"Sector Map:\n{generate_sector_map()}", "timestamp": datetime.now().isoformat(), "sender": "system", "room_id": room_id}
        ] + room_history:
            await websocket.send(json.dumps(msg))

        join_msg = {"type": "system", "content": f"User joined (Sector {x},{y}) - {len(sector_users[room_id])} users", "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"Client {client_id} connected to room {room_id} (Sector {x},{y})")

        if room_id == "logs":
            for entry in log_history[-10:]:
                await websocket.send(json.dumps({"type": "log", "content": entry["message"], "timestamp": entry["timestamp"], "sender": "system", "room_id": "logs"}))

        async for message in websocket:
            try:
                data = json.loads(message)
                if data.get("type") == "command":
                    if data.get("command") == "clear_history":
                        await clear_all_history()
                        continue
                    if data.get("command") == "show_map":
                        await websocket.send(json.dumps({"type": "system", "content": f"Sector Map:\n{generate_sector_map()}", "timestamp": datetime.now().isoformat(), "sender": "system", "room_id": room_id}))
                        continue
                data.update({"timestamp": datetime.now().isoformat(), "sender_node": NODE_NAME, "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