routers/websocket.py

"""
Real-time WebSocket Manager for PopCorn
Handles: Watch Parties, Chat, Notifications, and Live Events
"""
import asyncio
import json
import time
import uuid
from typing import Any
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Request, HTTPException
from database import fetch_one, fetch_all, execute
from auth import verify_init_data, get_current_user

router = APIRouter()

# ─── Connection Manager ───────────────────────────────────────────────────────

class ConnectionManager:
    def __init__(self):
        self._rooms: dict[str, set[WebSocket]] = {}       # room_code -> {ws}
        self._user_ws: dict[int, set[WebSocket]] = {}     # user_id -> {ws}
        self._room_owners: dict[str, int] = {}             # room_code -> host_user_id
        self._room_media: dict[str, dict] = {}             # room_code -> media info
        self._room_participants: dict[str, dict] = {}      # room_code -> {user_id: participant_info}

    async def connect_to_room(self, room_code: str, ws: WebSocket, user_id: int, user_data: dict, role: str = "follower"):
        await ws.accept()
        if room_code not in self._rooms:
            self._rooms[room_code] = set()
        self._rooms[room_code].add(ws)

        if user_id not in self._user_ws:
            self._user_ws[user_id] = set()
        self._user_ws[user_id].add(ws)

        # Store participant info
        if room_code not in self._room_participants:
            self._room_participants[room_code] = {}
        self._room_participants[room_code][user_id] = {
            "user_id": user_id,
            "name": user_data.get("first_name", "Unknown"),
            "username": user_data.get("username", ""),
            "avatar_url": user_data.get("photo_url", ""),
            "role": role,
            "is_ready": False,
            "joined_at": time.time(),
            "is_online": True,
        }

        # Notify others
        await self.broadcast_to_room(room_code, {
            "type": "user_joined",
            "participant": self._room_participants[room_code][user_id],
        }, exclude=ws)

        # Send current state
        await ws.send_json({
            "type": "room_state",
            "participants": list(self._room_participants[room_code].values()),
            "media": self._room_media.get(room_code),
            "room_code": room_code,
            "your_role": role,
        })

    async def disconnect(self, room_code: str, ws: WebSocket, user_id: int):
        if room_code in self._rooms:
            self._rooms[room_code].discard(ws)
            if not self._rooms[room_code]:
                del self._rooms[room_code]
                self._room_participants.pop(room_code, None)
                self._room_media.pop(room_code, None)

        if user_id in self._user_ws:
            self._user_ws[user_id].discard(ws)
            if not self._user_ws[user_id]:
                del self._user_ws[user_id]

        # Remove participant and notify
        if room_code in self._room_participants:
            self._room_participants[room_code].pop(user_id, None)
            await self.broadcast_to_room(room_code, {
                "type": "user_left",
                "user_id": user_id,
            })

    async def broadcast_to_room(self, room_code: str, message: dict, exclude: WebSocket | None = None):
        if room_code not in self._rooms:
            return
        dead = set()
        for ws in self._rooms[room_code]:
            if ws == exclude:
                continue
            try:
                await ws.send_json(message)
            except Exception:
                dead.add(ws)
        for ws in dead:
            self._rooms[room_code].discard(ws)

    async def send_to_user(self, user_id: int, message: dict):
        if user_id not in self._user_ws:
            return
        dead = set()
        for ws in self._user_ws[user_id]:
            try:
                await ws.send_json(message)
            except Exception:
                dead.add(ws)
        for ws in dead:
            self._user_ws[user_id].discard(ws)

    def get_room_participants_count(self, room_code: str) -> int:
        return len(self._room_participants.get(room_code, {}))


manager = ConnectionManager()


# ─── WebSocket Endpoint ───────────────────────────────────────────────────────

@router.websocket("/party/{room_code}")
async def party_websocket(websocket: WebSocket, room_code: str):
    user_id = None
    user_data = None

    try:
        # Wait for auth message first
        raw = await websocket.receive_text()
        auth_data = json.loads(raw)

        if auth_data.get("type") != "auth":
            await websocket.close(code=4001, reason="Auth required")
            return

        init_data = auth_data.get("initData", "")
        user_data = verify_init_data(init_data)
        if not user_data or not user_data.get("id"):
            await websocket.close(code=4001, reason="Invalid auth")
            return

        user_id = user_data["id"]
        role = auth_data.get("role", "follower")

        # Verify room exists
        room = await fetch_one(
            "SELECT * FROM party_rooms WHERE room_code = ? AND status != 'ended'",
            [room_code],
        )
        if not room:
            await websocket.send_json({"type": "error", "code": "ROOM_NOT_FOUND", "message": "Party room not found"})
            await websocket.close(code=4004, reason="Room not found")
            return

        await manager.connect_to_room(room_code, websocket, user_id, user_data, role)

        # Main message loop
        while True:
            raw = await websocket.receive_text()
            msg = json.loads(raw)
            msg_type = msg.get("type")

            if msg_type == "ping":
                await websocket.send_json({"type": "pong", "timestamp": time.time()})

            elif msg_type == "sync":
                # Leader sends sync; broadcast to all others
                if manager._room_participants.get(room_code, {}).get(user_id, {}).get("role") == "leader":
                    await manager.broadcast_to_room(room_code, {
                        "type": "sync",
                        "is_playing": msg.get("is_playing", False),
                        "position": msg.get("position", 0),
                        "timestamp": time.time(),
                        "playback_rate": msg.get("playback_rate", 1.0),
                    }, exclude=websocket)

            elif msg_type == "chat":
                message_data = {
                    "type": "chat_message",
                    "message": {
                        "id": str(uuid.uuid4()),
                        "user_id": str(user_id),
                        "user_name": user_data.get("first_name", "Unknown"),
                        "content": msg.get("content", ""),
                        "type": "chat",
                        "reply_to_id": msg.get("reply_to_id"),
                        "created_at": time.time(),
                    },
                }
                await manager.broadcast_to_room(room_code, message_data)

            elif msg_type == "reaction":
                await manager.broadcast_to_room(room_code, {
                    "type": "reaction",
                    "emoji": msg.get("emoji", "👍"),
                    "target_id": msg.get("target_id"),
                    "target_type": msg.get("target_type", "message"),
                    "user_id": str(user_id),
                })

            elif msg_type == "ready":
                is_ready = msg.get("is_ready", False)
                if room_code in manager._room_participants and user_id in manager._room_participants[room_code]:
                    manager._room_participants[room_code][user_id]["is_ready"] = is_ready

                await manager.broadcast_to_room(room_code, {
                    "type": "ready",
                    "user_id": str(user_id),
                    "is_ready": is_ready,
                })

            elif msg_type == "buffering":
                await manager.broadcast_to_room(room_code, {
                    "type": "buffering",
                    "user_id": str(user_id),
                    "is_buffering": msg.get("is_buffering", False),
                })

            elif msg_type == "leave":
                break

    except WebSocketDisconnect:
        pass
    except Exception as e:
        print(f"[ws] error: {e}")
    finally:
        if room_code and user_id:
            await manager.disconnect(room_code, websocket, user_id)


# ─── Chat WebSocket ────────────────────────────────────────────────────────────

@router.websocket("/chat/{conversation_id}")
async def chat_websocket(websocket: WebSocket, conversation_id: str):
    user_id = None
    try:
        raw = await websocket.receive_text()
        auth_data = json.loads(raw)
        init_data = auth_data.get("initData", "")
        user_data = verify_init_data(init_data)
        if not user_data or not user_data.get("id"):
            await websocket.close(code=4001)
            return

        user_id = user_data["id"]
        await websocket.accept()

        while True:
            raw = await websocket.receive_text()
            msg = json.loads(raw)
            if msg.get("type") == "ping":
                await websocket.send_json({"type": "pong"})
            elif msg.get("type") == "message":
                content = msg.get("content", "").strip()
                if content:
                    msg_id = str(uuid.uuid4())
                    await execute(
                        "INSERT INTO chat_messages (id, conversation_id, sender_id, content, created_at) VALUES (?, ?, ?, ?, datetime('now'))",
                        [msg_id, conversation_id, str(user_id), content],
                    )
                    # Update conversation last_message
                    await execute(
                        "UPDATE conversations SET last_message = ?, last_message_at = datetime('now') WHERE id = ?",
                        [content, conversation_id],
                    )
                    await websocket.send_json({
                        "type": "message_sent",
                        "id": msg_id,
                        "content": content,
                        "sender_id": str(user_id),
                        "created_at": time.time(),
                    })
    except WebSocketDisconnect:
        pass
    except Exception as e:
        print(f"[ws-chat] error: {e}")


# ─── REST Endpoints for Room Management ───────────────────────────────────────

@router.get("/room/{room_code}/participants")
async def get_room_participants(room_code: str):
    participants = manager._room_participants.get(room_code, {})
    return {"ok": True, "data": {"participants": list(participants.values()), "count": len(participants)}}


@router.post("/room/{room_code}/media")
async def set_room_media(room_code: str, request: Request):
    body = await request.json()
    manager._room_media[room_code] = body
    await manager.broadcast_to_room(room_code, {"type": "media_changed", "media": body})
    return {"ok": True}


@router.post("/room/{room_code}/kick")
async def kick_participant(room_code: str, request: Request):
    body = await request.json()
    target_user_id = body.get("user_id")
    if target_user_id:
        await manager.send_to_user(target_user_id, {"type": "kicked", "user_id": str(target_user_id)})
        return {"ok": True}
    raise HTTPException(status_code=400, detail="Missing user_id")