from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from typing import Dict, Set, Optional
from jose import JWTError
import json
import uuid

from app.utils.security import decode_token

router = APIRouter(tags=["websocket"])

# Store active connections by connection_id (unique per connection)
active_connections: Dict[str, WebSocket] = {}

# Store connections subscribed to specific events
event_subscriptions: Dict[int, Set[str]] = {}  # event_id -> set of connection_ids

# Store connections subscribed to specific matches (bet_id)
match_subscriptions: Dict[int, Set[str]] = {}  # bet_id -> set of connection_ids

# Map connection_id to websocket
connection_websockets: Dict[str, WebSocket] = {}


@router.websocket("/api/v1/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    token: str = Query(...),
    event_id: Optional[int] = Query(None),
    match_id: Optional[int] = Query(None)
):
    await websocket.accept()

    # Generate unique connection ID
    connection_id = str(uuid.uuid4())

    # Try to decode token to get user_id (for logging purposes)
    user_id = None
    if token and token != 'guest':
        try:
            payload = decode_token(token)
            user_id = payload.get("sub")
        except (JWTError, Exception):
            pass  # Invalid token, treat as guest

    print(f"[WebSocket] New connection: {connection_id}, user_id: {user_id}, event_id: {event_id}")

    # Store connection
    active_connections[connection_id] = websocket
    connection_websockets[connection_id] = websocket

    # Subscribe to event if specified
    if event_id:
        if event_id not in event_subscriptions:
            event_subscriptions[event_id] = set()
        event_subscriptions[event_id].add(connection_id)
        print(f"[WebSocket] Subscribed {connection_id} to event {event_id}. Total subscribers: {len(event_subscriptions[event_id])}")

    # Subscribe to match if specified
    if match_id:
        if match_id not in match_subscriptions:
            match_subscriptions[match_id] = set()
        match_subscriptions[match_id].add(connection_id)
        print(f"[WebSocket] Subscribed {connection_id} to match {match_id}. Total subscribers: {len(match_subscriptions[match_id])}")

    try:
        while True:
            data = await websocket.receive_text()
            # Handle incoming messages - could be used to subscribe/unsubscribe
            try:
                msg = json.loads(data)
                if msg.get('action') == 'subscribe' and msg.get('event_id'):
                    eid = msg['event_id']
                    if eid not in event_subscriptions:
                        event_subscriptions[eid] = set()
                    event_subscriptions[eid].add(connection_id)
                    print(f"[WebSocket] {connection_id} subscribed to event {eid}")
                elif msg.get('action') == 'unsubscribe' and msg.get('event_id'):
                    eid = msg['event_id']
                    if eid in event_subscriptions:
                        event_subscriptions[eid].discard(connection_id)
                        print(f"[WebSocket] {connection_id} unsubscribed from event {eid}")
                elif msg.get('action') == 'subscribe_match' and msg.get('match_id'):
                    mid = msg['match_id']
                    if mid not in match_subscriptions:
                        match_subscriptions[mid] = set()
                    match_subscriptions[mid].add(connection_id)
                    print(f"[WebSocket] {connection_id} subscribed to match {mid}")
                elif msg.get('action') == 'unsubscribe_match' and msg.get('match_id'):
                    mid = msg['match_id']
                    if mid in match_subscriptions:
                        match_subscriptions[mid].discard(connection_id)
                        print(f"[WebSocket] {connection_id} unsubscribed from match {mid}")
            except json.JSONDecodeError:
                pass
    except WebSocketDisconnect:
        print(f"[WebSocket] Disconnected: {connection_id}")
        # Clean up connection
        if connection_id in active_connections:
            del active_connections[connection_id]
        if connection_id in connection_websockets:
            del connection_websockets[connection_id]

        # Clean up event subscriptions
        for eid, subs in event_subscriptions.items():
            if connection_id in subs:
                subs.discard(connection_id)
                print(f"[WebSocket] Removed {connection_id} from event {eid} subscriptions")

        # Clean up match subscriptions
        for mid, subs in match_subscriptions.items():
            if connection_id in subs:
                subs.discard(connection_id)
                print(f"[WebSocket] Removed {connection_id} from match {mid} subscriptions")


async def broadcast_to_event(event_id: int, event_type: str, data: dict):
    """Broadcast a message to all connections subscribed to an event"""
    message = json.dumps({
        "type": event_type,
        "data": {"event_id": event_id, **data}
    })

    print(f"[WebSocket] Broadcasting {event_type} to event {event_id}")

    if event_id not in event_subscriptions:
        print(f"[WebSocket] No subscribers for event {event_id}")
        return

    subscribers = event_subscriptions[event_id].copy()
    print(f"[WebSocket] Found {len(subscribers)} subscribers for event {event_id}")

    disconnected = set()
    for conn_id in subscribers:
        ws = connection_websockets.get(conn_id)
        if ws:
            try:
                await ws.send_text(message)
                print(f"[WebSocket] Sent message to {conn_id}")
            except Exception as e:
                print(f"[WebSocket] Failed to send to {conn_id}: {e}")
                disconnected.add(conn_id)
        else:
            print(f"[WebSocket] Connection {conn_id} not found in websockets map")
            disconnected.add(conn_id)

    # Clean up disconnected connections
    for conn_id in disconnected:
        event_subscriptions[event_id].discard(conn_id)
        if conn_id in active_connections:
            del active_connections[conn_id]
        if conn_id in connection_websockets:
            del connection_websockets[conn_id]


async def broadcast_event(event_type: str, data: dict, user_ids: list[int] = None):
    """Broadcast an event to all connected users"""
    message = json.dumps({
        "type": event_type,
        "data": data
    })

    for conn_id, ws in list(connection_websockets.items()):
        try:
            await ws.send_text(message)
        except Exception:
            pass


async def broadcast_to_match(match_id: int, event_type: str, data: dict):
    """Broadcast a message to all connections subscribed to a match"""
    message = json.dumps({
        "type": event_type,
        "data": {"match_id": match_id, **data}
    })

    print(f"[WebSocket] Broadcasting {event_type} to match {match_id}")

    if match_id not in match_subscriptions:
        print(f"[WebSocket] No subscribers for match {match_id}")
        return

    subscribers = match_subscriptions[match_id].copy()
    print(f"[WebSocket] Found {len(subscribers)} subscribers for match {match_id}")

    disconnected = set()
    for conn_id in subscribers:
        ws = connection_websockets.get(conn_id)
        if ws:
            try:
                await ws.send_text(message)
                print(f"[WebSocket] Sent message to {conn_id}")
            except Exception as e:
                print(f"[WebSocket] Failed to send to {conn_id}: {e}")
                disconnected.add(conn_id)
        else:
            print(f"[WebSocket] Connection {conn_id} not found in websockets map")
            disconnected.add(conn_id)

    # Clean up disconnected connections
    for conn_id in disconnected:
        match_subscriptions[match_id].discard(conn_id)
        if conn_id in active_connections:
            del active_connections[conn_id]
        if conn_id in connection_websockets:
            del connection_websockets[conn_id]
