"""WebSocket connection manager."""

import asyncio
import json
import logging
import time
import uuid
from typing import TYPE_CHECKING, Any

from ieum.server.websocket.protocol import (
    MessageType,
    WSMessage,
    connected,
    error_message,
    interview_questions,
    plan_mode_changed,
    session_policy_updated,
    system_message,
)

if TYPE_CHECKING:
    from fastapi import WebSocket

    from ieum.server.services.session_service import Session, SessionService

logger = logging.getLogger(__name__)
REPLAY_HISTORY_MAX_EVENTS = 2000
# Maximum allowed WebSocket message size (1 MB)
WS_MAX_MESSAGE_SIZE = 1 * 1024 * 1024


class WebSocketConnection:
    """Represents a single WebSocket connection."""

    def __init__(self, websocket: Any, session_id: str) -> None:  # WebSocket
        """Initialize the connection.

        Args:
            websocket: FastAPI WebSocket instance.
            session_id: Associated session ID.
        """
        self.websocket = websocket
        self.session_id = session_id
        self.id = str(uuid.uuid4())[:8]
        self._closed = False

    async def send(self, message: WSMessage) -> None:
        """Send a message to the client.

        Args:
            message: Message to send.
        """
        if self._closed:
            return
        try:
            await self.websocket.send_json(message.to_dict())
        except Exception:
            self._closed = True

    async def receive(self) -> WSMessage | None:
        """Receive a message from the client.

        Returns:
            Received message or None on error.
        """
        if self._closed:
            return None
        try:
            raw = await self.websocket.receive_text()
            if len(raw) > WS_MAX_MESSAGE_SIZE:
                logger.warning(
                    "WebSocket message too large (%d bytes, max %d), dropping",
                    len(raw),
                    WS_MAX_MESSAGE_SIZE,
                )
                return None
            data = json.loads(raw)
            return WSMessage.from_dict(data)
        except Exception:
            self._closed = True
            return None

    async def close(self) -> None:
        """Close the connection."""
        if not self._closed:
            self._closed = True
            try:
                await self.websocket.close()
            except Exception:
                pass

    @property
    def is_closed(self) -> bool:
        """Check if connection is closed."""
        return self._closed


class WebSocketManager:
    """Manages WebSocket connections for all sessions."""

    def __init__(self) -> None:
        """Initialize the manager."""
        self._connections: dict[str, set[WebSocketConnection]] = {}
        self._lock = asyncio.Lock()
        # Pending approval futures per session
        self._pending_approvals: dict[str, asyncio.Future] = {}
        # Pending interview futures per interview_id
        self._pending_interviews: dict[str, asyncio.Future] = {}
        # Track background tasks to prevent garbage collection and enable cleanup
        self._background_tasks: set[asyncio.Task] = set()
        # Reconnect replay history (per-session sequence + event payloads)
        self._session_sequences: dict[str, int] = {}
        self._session_history: dict[str, list[dict[str, Any]]] = {}
        # Per-session chat.send rate-limit buckets
        self._session_rate_windows: dict[str, list[float]] = {}
        # Shared AgentService instance (created lazily, reused across messages)
        self._agent_service: Any = None

    async def connect(
        self,
        websocket: Any,  # WebSocket
        session_id: str,
        already_accepted: bool = False,
    ) -> WebSocketConnection:
        """Register a new WebSocket connection.

        Args:
            websocket: FastAPI WebSocket instance.
            session_id: Session ID to connect to.
            already_accepted: Whether the connection was already accepted.

        Returns:
            WebSocket connection wrapper.
        """
        if not already_accepted:
            await websocket.accept()
        conn = WebSocketConnection(websocket, session_id)

        async with self._lock:
            if session_id not in self._connections:
                self._connections[session_id] = set()
            self._connections[session_id].add(conn)

        return conn

    async def disconnect(self, connection: WebSocketConnection) -> None:
        """Unregister a WebSocket connection.

        Args:
            connection: Connection to unregister.
        """
        async with self._lock:
            session_conns = self._connections.get(connection.session_id)
            if session_conns:
                session_conns.discard(connection)
                if not session_conns:
                    del self._connections[connection.session_id]

        await connection.close()

    async def remove_session(self, session_id: str) -> None:
        """Remove all session-specific data (rate windows, history, sequences).

        Should be called when a session is deleted to prevent memory leaks.

        Args:
            session_id: Session ID to clean up.
        """
        async with self._lock:
            self._session_rate_windows.pop(session_id, None)
            self._session_history.pop(session_id, None)
            self._session_sequences.pop(session_id, None)
            # Also close and remove any remaining connections for this session
            connections = self._connections.pop(session_id, set())

        for conn in connections:
            await conn.close()

    async def broadcast(self, session_id: str, message: WSMessage) -> None:
        """Broadcast a message to all connections for a session.

        Args:
            session_id: Session ID to broadcast to.
            message: Message to send.
        """
        async with self._lock:
            next_seq = self._session_sequences.get(session_id, 0) + 1
            self._session_sequences[session_id] = next_seq
            if isinstance(message.data, dict):
                message.data["_seq"] = next_seq

            payload = message.to_dict()
            history = self._session_history.setdefault(session_id, [])
            history.append(payload)
            if len(history) > REPLAY_HISTORY_MAX_EVENTS:
                del history[: len(history) - REPLAY_HISTORY_MAX_EVENTS]

            connections = self._connections.get(session_id, set()).copy()

        for conn in connections:
            await conn.send(WSMessage.from_dict(payload))

    async def _replay_history(
        self,
        conn: WebSocketConnection,
        *,
        session_id: str,
        last_seq: int | None,
    ) -> None:
        """Replay buffered messages newer than ``last_seq`` to the connection."""
        if last_seq is None or last_seq < 0:
            return

        async with self._lock:
            history = list(self._session_history.get(session_id, []))

        replayed = 0
        for payload in history:
            seq = (payload.get("data") or {}).get("_seq")
            if isinstance(seq, int) and seq > last_seq:
                await conn.send(WSMessage.from_dict(payload))
                replayed += 1

        if replayed:
            logger.info(
                "Replayed %d websocket events for session=%s after seq=%s",
                replayed,
                session_id,
                last_seq,
            )

    async def _allow_session_message(self, session: Any) -> bool:
        """Return whether an incoming chat message is allowed by session rate limit."""
        max_requests = getattr(session, "rate_limit_max_requests", None)
        window_seconds = getattr(session, "rate_limit_window_seconds", None)
        if not max_requests or not window_seconds:
            return True

        now = time.monotonic()
        async with self._lock:
            timestamps = self._session_rate_windows.setdefault(session.id, [])
            timestamps[:] = [t for t in timestamps if now - t < window_seconds]
            if len(timestamps) >= max_requests:
                return False
            timestamps.append(now)
            return True

    async def _record_audit_event(
        self,
        session_service: Any,
        session: Any,
        event_type: str,
        payload: dict[str, Any],
    ) -> None:
        """Best-effort audit logging hook."""
        try:
            await session_service.append_audit_event(
                session_id=session.id,
                user_id=session.user_id,
                workspace_id=session.workspace_id,
                event_type=event_type,
                payload=payload,
            )
        except Exception:
            logger.debug(
                "Failed to append audit event for session=%s event=%s",
                session.id,
                event_type,
                exc_info=True,
            )

    async def handle_connection(
        self,
        websocket: Any,  # WebSocket
        session_id: str,
        session_service: Any,  # SessionService
        user_id: str | None = None,
        last_seq: int | None = None,
    ) -> None:
        """Handle a WebSocket connection lifecycle.

        Args:
            websocket: FastAPI WebSocket instance.
            session_id: Session ID to connect to.
            session_service: Session service instance.
        """
        # Accept connection first (required before any communication)
        await websocket.accept()

        # Verify session exists
        session = session_service.get_session(session_id, user_id=user_id)
        if session is None:
            await websocket.close(code=4004, reason="Session not found")
            return

        # Register connection (already accepted)
        conn = await self.connect(websocket, session_id, already_accepted=True)

        # Send connected message
        await conn.send(connected(session_id))
        await self._replay_history(conn, session_id=session_id, last_seq=last_seq)

        try:
            # Handle messages
            await self._handle_messages(conn, session, session_service)
        except Exception as e:
            await conn.send(error_message(str(e)))
        finally:
            await self.disconnect(conn)

    async def _handle_messages(
        self,
        conn: WebSocketConnection,
        session: Any,  # Session
        session_service: Any,  # SessionService
    ) -> None:
        """Handle incoming messages for a connection.

        Args:
            conn: WebSocket connection.
            session: Associated session.
            session_service: Session service.
        """
        from ieum.server.services.agent_service import AgentService
        from ieum.server.websocket.adapter import WebSocketUIAdapter

        # Reuse a single AgentService instance across all messages/connections
        if self._agent_service is None:
            self._agent_service = AgentService()
        agent_service = self._agent_service

        while not conn.is_closed:
            message = await conn.receive()
            if message is None:
                break

            # Keep in-memory and persisted activity timestamps in sync.
            try:
                await session_service.touch_session(session.id)
            except Exception:
                logger.debug(
                    "Failed to persist session activity for %s",
                    session.id,
                    exc_info=True,
                )

            if message.type == MessageType.CHAT_SEND:
                # Handle chat message
                user_input = message.data.get("content", "")
                if not user_input:
                    continue
                if not await self._allow_session_message(session):
                    await conn.send(
                        error_message(
                            "Session rate limit exceeded. Please wait before sending more messages.",
                            code="RATE_LIMIT",
                        )
                    )
                    continue
                await self._record_audit_event(
                    session_service,
                    session,
                    "chat.send",
                    {"content": str(user_input)[:4000], "length": len(str(user_input))},
                )

                # Inject plan mode preamble if active
                if (
                    getattr(session, "plan_mode", False)
                    and "## PLAN MODE ACTIVE" not in user_input
                ):
                    from ieum.plan_mode import PLAN_MODE_PREAMBLE
                    user_input = f"{PLAN_MODE_PREAMBLE}\n\n## Task\n{user_input}"

                # Save user message to conversation history
                from ieum.sessions import save_conversation_message
                await save_conversation_message(session.thread_id, "user", str(user_input))

                # Create async wrapper for broadcast
                async def send_message(msg: WSMessage, sid: str = session.id) -> None:
                    await self.broadcast(sid, msg)

                async def request_approval(
                    action_requests: list[dict[str, Any]],
                    assistant_id: str | None,
                    sid: str = session.id,
                ) -> asyncio.Future:
                    return await self._request_approval(sid, action_requests, assistant_id)

                async def request_interview(
                    interview_id: str,
                    questions: list[dict[str, Any]],
                    sid: str = session.id,
                ) -> asyncio.Future:
                    return await self._request_interview(sid, interview_id, questions)

                async def audit_event(
                    event_type: str,
                    payload: dict[str, Any],
                    sess: Any = session,
                ) -> None:
                    await self._record_audit_event(session_service, sess, event_type, payload)

                # Create adapter for this execution
                adapter = WebSocketUIAdapter(
                    send_message=send_message,
                    request_approval_fn=request_approval,
                    request_interview_fn=request_interview,
                    audit_event_fn=audit_event,
                )
                # Set model name for cost tracking
                adapter._model_name = getattr(session, "model_name", None)

                # Execute in background with proper task tracking
                task = asyncio.create_task(
                    self._execute_and_capture_plan(
                        agent_service, session, user_input, adapter
                    )
                )
                self._background_tasks.add(task)
                task.add_done_callback(self._task_done_callback)

            elif message.type == MessageType.CHAT_CANCEL:
                # Cancel current execution
                await agent_service.cancel(session.id)
                await self._record_audit_event(
                    session_service,
                    session,
                    "chat.cancel",
                    {"source": "client"},
                )

            elif message.type == MessageType.APPROVAL_RESPOND:
                # Handle approval response
                interrupt_id = message.data.get("interrupt_id")
                decision = message.data.get("decision")
                await self._handle_approval_response(session.id, interrupt_id, decision)
                await self._record_audit_event(
                    session_service,
                    session,
                    "approval.respond",
                    {
                        "interrupt_id": interrupt_id,
                        "decision": decision,
                    },
                )

            elif message.type == MessageType.SESSION_SET_APPROVAL_POLICY:
                requested_policy = message.data.get("approval_policy")
                from ieum.approval_policy import is_valid_approval_policy

                if not is_valid_approval_policy(str(requested_policy or "")):
                    await conn.send(
                        error_message(
                            "Invalid approval policy",
                            code="INVALID_APPROVAL_POLICY",
                        )
                    )
                    continue
                updated = await session_service.update_session_approval_policy(
                    session.id,
                    str(requested_policy or "manual"),
                    user_id=session.user_id,
                )
                if updated is None:
                    await conn.send(error_message("Session not found", code="SESSION_NOT_FOUND"))
                    continue
                session = updated
                await self.broadcast(
                    session.id,
                    session_policy_updated(session.approval_policy),
                )
                await self._record_audit_event(
                    session_service,
                    session,
                    "session.approval_policy",
                    {"approval_policy": session.approval_policy},
                )

            elif message.type == MessageType.AUDIT_REPLAY:
                since_id = message.data.get("since_id")
                limit = message.data.get("limit", 200)
                try:
                    normalized_limit = max(1, min(int(limit), 1000))
                except (TypeError, ValueError):
                    normalized_limit = 200
                try:
                    since_value = int(since_id) if since_id is not None else None
                except (TypeError, ValueError):
                    since_value = None

                events = await session_service.list_audit_events(
                    session_id=session.id,
                    user_id=session.user_id,
                    since_id=since_value,
                    limit=normalized_limit,
                )
                await conn.send(
                    WSMessage(
                        type=MessageType.AUDIT_REPLAY_RESULT,
                        data={
                            "session_id": session.id,
                            "events": events,
                        },
                    )
                )

            elif message.type == MessageType.INTERVIEW_ANSWER:
                # Handle interview answer
                iview_id = message.data.get("interview_id")
                answers = message.data.get("answers")
                await self._handle_interview_response(iview_id, answers)

            elif message.type == MessageType.INTERVIEW_CANCEL:
                # Handle interview cancellation
                iview_id = message.data.get("interview_id")
                await self._handle_interview_response(iview_id, {"_cancelled": True})

            elif message.type == MessageType.PLAN_SET_MODE:
                # Toggle plan mode or set explicitly
                enable = message.data.get("enable")
                if enable is None:
                    # Toggle
                    if getattr(session, "plan_mode", False):
                        session.exit_plan_mode()
                    else:
                        session.enter_plan_mode()
                elif enable:
                    if not getattr(session, "plan_mode", False):
                        session.enter_plan_mode()
                else:
                    if getattr(session, "plan_mode", False):
                        session.exit_plan_mode()
                await self.broadcast(
                    session.id,
                    plan_mode_changed(
                        session.plan_mode,
                        session.approval_policy,
                        session.current_plan,
                    ),
                )
                await self._record_audit_event(
                    session_service,
                    session,
                    "plan.set_mode",
                    {"plan_mode": session.plan_mode},
                )

            elif message.type == MessageType.PLAN_APPROVE:
                # Approve and execute the captured plan
                plan = getattr(session, "current_plan", None)
                task_text = message.data.get("task", "")
                if not plan:
                    await conn.send(
                        error_message("No plan to approve", code="NO_PLAN")
                    )
                    continue
                session.exit_plan_mode()
                await self.broadcast(
                    session.id,
                    plan_mode_changed(
                        session.plan_mode,
                        session.approval_policy,
                    ),
                )
                # Build execution message from plan
                from ieum.plan_mode import PLAN_EXECUTE_TEMPLATE
                execute_msg = PLAN_EXECUTE_TEMPLATE.format(
                    plan=plan,
                    task=task_text or "(plan execution)",
                )
                # Save and execute
                await save_conversation_message(session.thread_id, "user", execute_msg)

                async def send_message_plan(msg: WSMessage, sid: str = session.id) -> None:
                    await self.broadcast(sid, msg)

                async def request_approval_plan(
                    action_requests: list[dict[str, Any]],
                    assistant_id: str | None,
                    sid: str = session.id,
                ) -> asyncio.Future:
                    return await self._request_approval(sid, action_requests, assistant_id)

                async def request_interview_plan(
                    interview_id: str,
                    questions: list[dict[str, Any]],
                    sid: str = session.id,
                ) -> asyncio.Future:
                    return await self._request_interview(sid, interview_id, questions)

                async def audit_event_plan(
                    event_type: str,
                    payload: dict[str, Any],
                    sess: Any = session,
                ) -> None:
                    await self._record_audit_event(session_service, sess, event_type, payload)

                plan_adapter = WebSocketUIAdapter(
                    send_message=send_message_plan,
                    request_approval_fn=request_approval_plan,
                    request_interview_fn=request_interview_plan,
                    audit_event_fn=audit_event_plan,
                )
                plan_adapter._model_name = getattr(session, "model_name", None)
                plan_task = asyncio.create_task(
                    self._execute_and_capture_plan(
                        agent_service, session, execute_msg, plan_adapter
                    )
                )
                self._background_tasks.add(plan_task)
                plan_task.add_done_callback(self._task_done_callback)
                await self._record_audit_event(
                    session_service,
                    session,
                    "plan.approve",
                    {"task": task_text[:4000]},
                )

            elif message.type == MessageType.TASK_SUBMIT:
                # Create a background task
                description = message.data.get("description", "")
                if not description:
                    await conn.send(
                        error_message("Task description is required", code="MISSING_DESCRIPTION")
                    )
                    continue
                task_svc = getattr(self, "_task_service", None)
                if task_svc is None:
                    await conn.send(
                        error_message("Background tasks not available", code="TASK_SERVICE_UNAVAILABLE")
                    )
                    continue
                user_id = getattr(session, "user_id", "anonymous")
                await task_svc.create_task(
                    session_id=session.id,
                    user_id=user_id,
                    description=description,
                )
                await self._record_audit_event(
                    session_service,
                    session,
                    "task.submit",
                    {"description": description[:4000]},
                )

            elif message.type == MessageType.TASK_CANCEL:
                # Cancel a background task
                task_id = message.data.get("task_id")
                if not task_id:
                    await conn.send(
                        error_message("Task ID is required", code="MISSING_TASK_ID")
                    )
                    continue
                task_svc = getattr(self, "_task_service", None)
                if task_svc is None:
                    await conn.send(
                        error_message("Background tasks not available", code="TASK_SERVICE_UNAVAILABLE")
                    )
                    continue
                cancelled = await task_svc.cancel_task(task_id)
                if not cancelled:
                    await conn.send(
                        error_message(f"Cannot cancel task {task_id}", code="TASK_CANCEL_FAILED")
                    )
                await self._record_audit_event(
                    session_service,
                    session,
                    "task.cancel",
                    {"task_id": task_id},
                )

            elif message.type == MessageType.PING:
                # Respond to ping
                await conn.send(WSMessage(type=MessageType.PONG))

    async def _execute_and_capture_plan(
        self,
        agent_service: Any,
        session: Any,
        user_input: str,
        adapter: Any,
    ) -> None:
        """Execute agent and capture plan if in plan mode."""
        await agent_service.execute(session, user_input, adapter)
        # Sync adapter cost to session object for REST API access
        adapter_cost = getattr(adapter, "_session_cost", 0.0)
        if adapter_cost > 0:
            session.session_cost = adapter_cost
        # After execution, capture plan if in plan mode
        if getattr(session, "plan_mode", False):
            plan_text = getattr(adapter, "_last_completed_content", "")
            if plan_text:
                session.current_plan = plan_text
                await self.broadcast(
                    session.id,
                    system_message(
                        "Plan captured. Use '/plan show' to review, '/plan approve' to execute."
                    ),
                )
                await self.broadcast(
                    session.id,
                    plan_mode_changed(
                        session.plan_mode,
                        session.approval_policy,
                        session.current_plan,
                    ),
                )

    async def _request_approval(
        self,
        session_id: str,
        action_requests: list[dict[str, Any]],
        assistant_id: str | None,
    ) -> asyncio.Future:
        """Request approval from the client.

        Args:
            session_id: Session ID.
            action_requests: Tool calls requiring approval.
            assistant_id: Agent identifier.

        Returns:
            Future that resolves with the approval decision.
        """
        from ieum.file_ops import build_approval_preview
        from ieum.server.websocket.protocol import approval_request

        # Generate interrupt ID
        interrupt_id = str(uuid.uuid4())[:8]

        # Build preview with diff for file operations
        enriched_requests = []
        for action in action_requests:
            tool_name = action.get("tool") or action.get("action", {}).get("tool")
            tool_args = action.get("tool_input") or action.get("action", {}).get("tool_input") or {}

            enriched_action = dict(action)

            # Add diff preview for file operations
            if tool_name in ("write_file", "edit_file"):
                preview = build_approval_preview(tool_name, tool_args, assistant_id)
                if preview:
                    enriched_action["preview"] = {
                        "title": preview.title,
                        "details": preview.details,
                        "diff": preview.diff,
                        "error": preview.error,
                    }

            enriched_requests.append(enriched_action)

        # Create future for response
        future: asyncio.Future = asyncio.get_running_loop().create_future()

        async with self._lock:
            self._pending_approvals[interrupt_id] = future

        # Send approval request
        await self.broadcast(
            session_id,
            approval_request(interrupt_id, enriched_requests, assistant_id),
        )

        return future

    async def _handle_approval_response(
        self,
        session_id: str,
        interrupt_id: str | None,
        decision: dict[str, Any] | None,
    ) -> None:
        """Handle an approval response from the client.

        Args:
            session_id: Session ID.
            interrupt_id: Interrupt ID being responded to.
            decision: Approval decision.
        """
        if interrupt_id is None or decision is None:
            return

        async with self._lock:
            future = self._pending_approvals.pop(interrupt_id, None)

        if future and not future.done():
            future.set_result(decision)

    async def _request_interview(
        self,
        session_id: str,
        interview_id: str,
        questions: list[dict[str, Any]],
    ) -> asyncio.Future:
        """Request interview answers from the client.

        Args:
            session_id: Session ID.
            interview_id: Unique interview interrupt ID.
            questions: List of question dicts to present.

        Returns:
            Future that resolves with the interview answers.
        """
        future: asyncio.Future = asyncio.get_running_loop().create_future()

        async with self._lock:
            self._pending_interviews[interview_id] = future

        # Send interview questions to all connected clients
        await self.broadcast(
            session_id,
            interview_questions(interview_id, questions),
        )

        return future

    async def _handle_interview_response(
        self,
        interview_id: str | None,
        answers: dict[str, Any] | None,
    ) -> None:
        """Handle an interview response from the client.

        Args:
            interview_id: Interview ID being responded to.
            answers: User's answers or cancellation.
        """
        if interview_id is None or answers is None:
            return

        async with self._lock:
            future = self._pending_interviews.pop(interview_id, None)

        if future and not future.done():
            future.set_result(answers)

    def _task_done_callback(self, task: asyncio.Task) -> None:
        """Callback for completed background tasks.

        Removes the task from the tracking set and logs unhandled exceptions.

        Args:
            task: The completed task.
        """
        self._background_tasks.discard(task)
        if not task.cancelled():
            exc = task.exception()
            if exc is not None:
                logger.error(
                    "Unhandled exception in background task: %s", exc, exc_info=exc,
                )

    async def shutdown(self) -> None:
        """Cancel and await all pending background tasks.

        Should be called during application shutdown to ensure clean cleanup.
        """
        tasks = list(self._background_tasks)
        if not tasks:
            return

        logger.info("Shutting down %d background tasks", len(tasks))
        for task in tasks:
            task.cancel()

        results = await asyncio.gather(*tasks, return_exceptions=True)
        for result in results:
            if isinstance(result, Exception) and not isinstance(result, asyncio.CancelledError):
                logger.error("Error during task shutdown: %s", result)

        self._background_tasks.clear()
