from __future__ import annotations

import asyncio
import json
from datetime import datetime, timezone
from typing import Any
from weakref import WeakSet

from fastapi import WebSocket, WebSocketDisconnect

from prompt_auto_lab.api.services import user_service


class ConnectionManager:
    """Manages WebSocket connections for multi-user session sharing."""

    def __init__(self) -> None:
        # session_id -> set of WebSocket connections
        self._connections: dict[str, WeakSet[WebSocket]] = {}
        self._lock = asyncio.Lock()

    async def connect(self, websocket: WebSocket, session_id: str) -> None:
        """Accept connection and add to session room."""
        await websocket.accept()
        async with self._lock:
            if session_id not in self._connections:
                self._connections[session_id] = WeakSet()
            self._connections[session_id].add(websocket)

    async def disconnect(self, websocket: WebSocket, session_id: str) -> None:
        """Remove connection from session room."""
        async with self._lock:
            if session_id in self._connections:
                self._connections[session_id].discard(websocket)
                if not self._connections[session_id]:
                    del self._connections[session_id]

    async def broadcast_to_session(
        self, session_id: str, message: dict[str, Any]
    ) -> None:
        """Broadcast message to all connections in a session."""
        async with self._lock:
            # Create a snapshot (list copy) to avoid "Set changed size during iteration"
            # when WeakSet references are garbage collected or modified concurrently
            connections = list(self._connections.get(session_id, WeakSet()))

        # Add timestamp if not present
        if "timestamp" not in message:
            message["timestamp"] = datetime.now(timezone.utc).isoformat()

        message_json = json.dumps(message, default=str)
        disconnected: list[WebSocket] = []

        for websocket in connections:
            try:
                await websocket.send_text(message_json)
            except Exception:
                disconnected.append(websocket)

        # Clean up disconnected
        if disconnected:
            async with self._lock:
                for ws in disconnected:
                    self._connections.get(session_id, WeakSet()).discard(ws)

    async def send_personal(self, websocket: WebSocket, message: dict[str, Any]) -> None:
        """Send message to a specific connection."""
        if "timestamp" not in message:
            message["timestamp"] = datetime.now(timezone.utc).isoformat()
        await websocket.send_text(json.dumps(message, default=str))

    def get_connection_count(self, session_id: str) -> int:
        """Get number of active connections for a session."""
        return len(self._connections.get(session_id, WeakSet()))


# Global connection manager instance
manager = ConnectionManager()


async def websocket_handler(websocket: WebSocket, session_id: str, user_id: str | None = None) -> None:
    """Handle WebSocket connection for a session.

    Args:
        websocket: The WebSocket connection
        session_id: The session to connect to
        user_id: The user_id to verify ownership (optional for backwards compatibility)
    """
    # Verify user owns the session if user_id is provided
    if user_id:
        sessions_file = user_service.get_user_sessions_file(user_id)
        if sessions_file.exists():
            import json as json_module
            with open(sessions_file, "r", encoding="utf-8") as f:
                sessions = json_module.load(f)
            session_ids = [s.get("id") for s in sessions]
            if session_id not in session_ids:
                # Allow connection for backwards compatibility with orphaned sessions
                # Just log a warning instead of rejecting
                import logging
                logging.getLogger(__name__).warning(
                    f"Session {session_id} not in user {user_id}'s session list (orphaned session)"
                )

    await manager.connect(websocket, session_id)

    try:
        # Send initial connection confirmation
        await manager.send_personal(
            websocket,
            {
                "type": "connected",
                "payload": {
                    "session_id": session_id,
                    "message": "Connected to session",
                    "active_users": manager.get_connection_count(session_id),
                },
            },
        )

        # Broadcast user joined
        await manager.broadcast_to_session(
            session_id,
            {
                "type": "user_joined",
                "payload": {"active_users": manager.get_connection_count(session_id)},
            },
        )

        # Listen for messages
        while True:
            data = await websocket.receive_text()
            try:
                message = json.loads(data)
                await handle_client_message(websocket, session_id, message)
            except json.JSONDecodeError:
                await manager.send_personal(
                    websocket,
                    {"type": "error", "payload": {"message": "Invalid JSON"}},
                )

    except WebSocketDisconnect:
        pass
    finally:
        await manager.disconnect(websocket, session_id)
        # Broadcast user left
        await manager.broadcast_to_session(
            session_id,
            {
                "type": "user_left",
                "payload": {"active_users": manager.get_connection_count(session_id)},
            },
        )


async def handle_client_message(
    websocket: WebSocket, session_id: str, message: dict[str, Any]
) -> None:
    """Handle incoming client messages."""
    msg_type = message.get("type", "")
    payload = message.get("payload", {})

    if msg_type == "subscribe":
        # Client wants to subscribe to updates (already connected)
        await manager.send_personal(
            websocket,
            {"type": "subscribed", "payload": {"session_id": session_id}},
        )

    elif msg_type == "command":
        command = payload.get("command")
        if command in ("pause", "resume", "stop"):
            # Broadcast command to all users
            await manager.broadcast_to_session(
                session_id,
                {"type": "command", "payload": {"command": command, "from": "user"}},
            )
        else:
            await manager.send_personal(
                websocket,
                {"type": "error", "payload": {"message": f"Unknown command: {command}"}},
            )

    elif msg_type == "user_input":
        # User is providing input to the agent
        user_message = payload.get("message", "")

        # Import here to avoid circular imports
        from prompt_auto_lab.api.routes.sessions import submit_user_input

        # Submit the input to the session's queue for the orchestrator
        await submit_user_input(session_id, user_message)

        # Broadcast as agent event so it shows in MessageStream
        await manager.broadcast_to_session(
            session_id,
            {
                "type": "agent_event",
                "payload": {
                    "type": "agent_message",
                    "agent": "UserProxy",
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                    "payload": {"message": f"User: {user_message}"},
                },
            },
        )

    elif msg_type == "state_patch":
        # User wants to update state - broadcast to all
        await manager.broadcast_to_session(
            session_id,
            {"type": "state_patch", "payload": payload},
        )

    else:
        await manager.send_personal(
            websocket,
            {"type": "error", "payload": {"message": f"Unknown message type: {msg_type}"}},
        )
