import json
import time
from typing import List, Dict
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
import redis.asyncio as redis
import asyncio
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# Redis config (injected via Vercel env)
REDIS_URL="redis://default:S8PlPlmKsUO7aqCYxQnjLuHAHbRMMesm@redis-16085.c239.us-east-1-2.ec2.redns.redis-cloud.com:16085"
MESSAGES_KEY = "messages"
CHAT_CHANNEL = "chat"
TTL_SECONDS = 86400  # 24 hours

# Global Redis client
r = None

@app.on_event("startup")
async def startup():
    global r
    r = redis.from_url(REDIS_URL)
    try:
        await r.ping()
        logger.info("Connected to Redis")
        # Set expiration on messages key
        await r.expire(MESSAGES_KEY, TTL_SECONDS)
    except Exception as e:
        logger.error(f"Redis connection failed: {e}")
        raise

@app.on_event("shutdown")
async def shutdown():
    if r:
        await r.close()
        logger.info("Redis connection closed")

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        # Receive username
        username_data = await websocket.receive_text()
        try:
            username = json.loads(username_data).get("username", "Anonymous")
            if not username:
                username = "Anonymous"
        except json.JSONDecodeError:
            username = "Anonymous"
            logger.warning("Invalid username JSON, defaulting to Anonymous")

        # Fetch history: last 24h, limit 100 messages
        now = int(time.time())
        cutoff = now - TTL_SECONDS
        history: List[Dict] = await r.zrangebyscore(MESSAGES_KEY, cutoff, now, withscores=False)
        history = [json.loads(msg) for msg in history[-100:]]  # Last 100

        # Send history
        await websocket.send_text(json.dumps({"type": "history", "messages": history}))
        logger.info(f"User {username} connected, sent {len(history)} messages")

        # Start pub/sub listener task
        listener_task = asyncio.create_task(forward_messages(websocket))

        # Receive loop
        while True:
            data = await websocket.receive_text()
            msg_data = json.loads(data)
            if msg_data.get("type") == "quit":
                logger.info(f"User {username} quit")
                break

            message = {
                "username": username,
                "message": msg_data.get("message", ""),
                "timestamp": int(time.time())
            }

            # Store and trim
            await r.zadd(MESSAGES_KEY, {json.dumps(message): message["timestamp"]})
            await r.zremrangebyscore(MESSAGES_KEY, 0, cutoff)  # Remove old messages
            await r.expire(MESSAGES_KEY, TTL_SECONDS)  # Ensure key expires

            # Publish for broadcast
            await r.publish(CHAT_CHANNEL, json.dumps(message))
            logger.info(f"Message from {username}: {message['message']}")

    except WebSocketDisconnect:
        logger.info(f"User {username} disconnected")
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
    finally:
        if 'listener_task' in locals():
            listener_task.cancel()

async def forward_messages(websocket: WebSocket):
    """Forward pub/sub messages to this WebSocket."""
    pubsub = r.pubsub()
    await pubsub.subscribe(CHAT_CHANNEL)
    try:
        async for message in pubsub.listen():
            if message["type"] == "message":
                data = json.loads(message["data"])
                await websocket.send_text(json.dumps({"type": "message", "data": data}))
    except asyncio.CancelledError:
        pass
    except Exception as e:
        logger.error(f"Pub/sub error: {e}")
    finally:
        await pubsub.unsubscribe(CHAT_CHANNEL)
        await pubsub.close()

@app.get("/")
async def root():
    return {"status": "Server running", "domain": "msg.ap1.cursedbots.xyz"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)