"""WebSocket UI adapter for agent execution."""

import asyncio
import json
import uuid
from collections.abc import Awaitable
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, Callable

from ieum.file_ops import FileOpTracker
from ieum.server.sanitize import sanitize_agent_output
from ieum.server.websocket.protocol import (
    WSMessage,
    chat_message_end,
    chat_message_start,
    chat_token,
    error_message,
    file_changed,
    interview_complete,
    stream_debug,
    stream_state,
    system_message,
    thinking_end,
    thinking_start,
    thinking_token,
    tool_error,
    tool_rejected,
    tool_start,
    tool_success,
)
from ieum.stream_handler import BaseStreamHandler

if TYPE_CHECKING:
    from ieum.server.services.session_service import Session


class WebSocketUIAdapter(BaseStreamHandler):
    """Adapter for rendering agent output to WebSocket clients.

    Inherits from BaseStreamHandler to reuse shared event processing logic.
    Implements the on_* callbacks to send WebSocket protocol messages.
    """

    def __init__(
        self,
        send_message: Callable[[WSMessage], Awaitable[None]],
        request_approval_fn: Callable[[list[dict], str | None], Awaitable[asyncio.Future]],
        request_interview_fn: Callable[[str, list[dict]], Awaitable[asyncio.Future]] | None = None,
        audit_event_fn: Callable[[str, dict[str, Any]], Awaitable[None]] | None = None,
    ) -> None:
        """Initialize the adapter.

        Args:
            send_message: Async callable to send a WebSocket message.
            request_approval_fn: Callable that returns a Future for HITL approval.
            request_interview_fn: Callable that returns a Future for interview answers.
            audit_event_fn: Optional callback for audit event capture.
        """
        self._send_message = send_message
        self._request_approval_fn = request_approval_fn
        self._request_interview_fn = request_interview_fn
        self._audit_event_fn = audit_event_fn

        # State tracking
        self._current_message_id: str | None = None
        self._current_tool_ids: set[str] = set()
        self._pending_text = ""
        self._is_thinking = False
        self._thinking_message_id: str | None = None
        self._accumulated_text = ""
        self._thread_id: str | None = None
        self._last_completed_content: str = ""
        # Cost tracking
        self._model_name: str | None = None
        self._session_cost: float = 0.0

    # --- Legacy compatibility aliases ---
    # These are used by manager.py's closure-based wrappers
    @property
    def _request_approval(self):
        return self._request_approval_fn

    @property
    def _request_interview(self):
        return self._request_interview_fn

    # --- Direct send methods (used by manager and other code) ---

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

    async def send_token(self, content: str) -> None:
        """Send a streaming token (sanitized for XSS safety)."""
        if self._current_message_id is None:
            self._current_message_id = str(uuid.uuid4())[:8]
            await self.send(chat_message_start(self._current_message_id))
        # Accumulate raw text for conversation persistence
        self._accumulated_text += content
        # Security: Sanitize agent output before sending to browser
        safe_content = sanitize_agent_output(content)
        await self.send(chat_token(safe_content, self._current_message_id))

    async def end_message(self, tokens: int = 0) -> None:
        """End the current message."""
        if self._current_message_id:
            # Calculate cost for this message
            msg_cost = 0.0
            breakdown = getattr(self, "_token_breakdown", {})
            if self._model_name and breakdown:
                from ieum.pricing import calculate_cost
                msg_cost = calculate_cost(
                    self._model_name,
                    breakdown.get("input_tokens", 0),
                    breakdown.get("output_tokens", 0),
                )
                if msg_cost > 0:
                    self._session_cost += msg_cost

            await self.send(chat_message_end(
                self._current_message_id, tokens,
                cost=msg_cost, session_cost=self._session_cost,
            ))
            # Save assistant message to conversation history
            if self._accumulated_text and self._thread_id:
                from ieum.sessions import save_conversation_message
                await save_conversation_message(
                    self._thread_id, "assistant", self._accumulated_text, tokens
                )
            self._accumulated_text = ""
            self._current_message_id = None

    async def send_tool_start(
        self, tool_id: str, tool_name: str, args: dict[str, Any],
    ) -> None:
        """Send a tool start message."""
        self._current_tool_ids.add(tool_id)
        await self.send(tool_start(tool_id, tool_name, args))
        if self._audit_event_fn:
            await self._audit_event_fn(
                "tool.start",
                {"tool_id": tool_id, "tool_name": tool_name, "args": args},
            )

    async def send_tool_success(self, tool_id: str, output: str) -> None:
        """Send a tool success message (sanitized for XSS safety)."""
        self._current_tool_ids.discard(tool_id)
        safe_output = sanitize_agent_output(output)
        await self.send(tool_success(tool_id, safe_output))
        if self._audit_event_fn:
            await self._audit_event_fn(
                "tool.success",
                {"tool_id": tool_id, "output": safe_output[:20000]},
            )

    async def send_tool_error(self, tool_id: str, err: str) -> None:
        """Send a tool error message (sanitized for XSS safety)."""
        self._current_tool_ids.discard(tool_id)
        safe_err = sanitize_agent_output(err)
        await self.send(tool_error(tool_id, safe_err))
        if self._audit_event_fn:
            await self._audit_event_fn(
                "tool.error",
                {"tool_id": tool_id, "error": safe_err[:20000]},
            )

    async def send_tool_rejected(self, tool_id: str) -> None:
        """Send a tool rejected message."""
        self._current_tool_ids.discard(tool_id)
        await self.send(tool_rejected(tool_id))

    async def send_file_changed(
        self, path: str, action: str, diff: str | None = None,
    ) -> None:
        """Send a file changed message."""
        await self.send(file_changed(path, action, diff))
        if self._audit_event_fn:
            await self._audit_event_fn(
                "file.changed",
                {"path": path, "action": action, "diff": (diff or "")[:20000]},
            )

    async def send_system_message(self, content: str) -> None:
        """Send a system message."""
        await self.send(system_message(content))
        if self._audit_event_fn:
            await self._audit_event_fn(
                "system.message",
                {"content": content[:4000]},
            )

    async def send_error(self, err: str, code: str | None = None) -> None:
        """Send an error message."""
        await self.send(error_message(err, code))
        if self._audit_event_fn:
            await self._audit_event_fn(
                "error",
                {"error": err[:4000], "code": code},
            )

    async def send_state(self, node: str, state: dict[str, Any]) -> None:
        """Send a state snapshot (values mode)."""
        safe_state = {}
        for key, value in state.items():
            try:
                json.dumps(value)
                safe_state[key] = value
            except (TypeError, ValueError):
                safe_state[key] = str(value)[:500]
        await self.send(stream_state(node, safe_state))

    async def send_debug(
        self, event_type: str, node: str | None = None, data: dict[str, Any] | None = None,
    ) -> None:
        """Send a debug event."""
        await self.send(stream_debug(event_type, node, data))

    async def send_thinking_start(self) -> None:
        """Send thinking start message."""
        if not self._is_thinking:
            self._is_thinking = True
            self._thinking_message_id = str(uuid.uuid4())[:8]
            await self.send(thinking_start(self._thinking_message_id))

    async def send_thinking_token(self, content: str) -> None:
        """Send a thinking token."""
        if not self._is_thinking:
            await self.send_thinking_start()
        await self.send(thinking_token(content, self._thinking_message_id))

    async def send_thinking_end(self) -> None:
        """Send thinking end message."""
        if self._is_thinking:
            await self.send(thinking_end(self._thinking_message_id))
            self._is_thinking = False
            self._thinking_message_id = None

    def mark_all_tools_rejected(self) -> None:
        """Mark all current tools as rejected (for cleanup)."""
        self._current_tool_ids.clear()

    # --- BaseStreamHandler callback implementations ---

    def should_process_values_stream(self) -> bool:
        return True

    def should_process_debug_stream(self) -> bool:
        return True

    async def on_chat_token(self, token: str, ns_key: tuple) -> None:
        # End thinking before regular text if active
        if self._is_thinking:
            await self.send_thinking_end()
        await self.send_token(token)

    async def on_tool_start(
        self, tool_id: str, tool_name: str, tool_input: dict, ns_key: tuple,
    ) -> None:
        await self.send_tool_start(tool_id, tool_name, tool_input)

    async def on_tool_end(
        self, tool_id: str, tool_name: str, output: str, status: str,
    ) -> None:
        if status == "success":
            await self.send_tool_success(tool_id, output)
        else:
            await self.send_tool_error(tool_id, output or "Error")

    async def on_file_changed(self, record: Any) -> None:
        await self.send_file_changed(record.display_path, record.tool_name, record.diff)

    async def on_thinking_token(self, token: str) -> None:
        await self.send_thinking_token(token)

    async def on_thinking_end(self) -> None:
        await self.send_thinking_end()

    async def on_message_complete(self, ns_key: tuple, tokens: int) -> None:
        if self._accumulated_text:
            self._last_completed_content = self._accumulated_text
        await self.end_message(tokens)

    async def on_human_message(self, ns_key: tuple) -> None:
        await self.end_message()

    async def on_stream_error(self, err: str) -> None:
        await self.send_error(err)

    async def on_interrupted(self) -> None:
        await self.send_system_message("Interrupted by user")
        for tid in list(self._current_tool_ids):
            await self.send_tool_rejected(tid)
        self.mark_all_tools_rejected()

    async def on_tool_rejected_cleanup(self) -> None:
        for tid in list(self._current_tool_ids):
            await self.send_tool_rejected(tid)
        self.mark_all_tools_rejected()

    async def on_rejection_message(self) -> None:
        await self.send_system_message(
            "Command rejected. Tell the agent what you'd like instead."
        )

    async def request_approval(
        self, action_requests: list[dict], assistant_id: str | None,
    ) -> asyncio.Future:
        return await self._request_approval_fn(action_requests, assistant_id)

    async def request_interview(
        self, interview_id: str, questions: list[dict],
    ) -> asyncio.Future:
        if self._request_interview_fn:
            return await self._request_interview_fn(interview_id, questions)
        # No interview handler - return a resolved future with empty answers
        loop = asyncio.get_running_loop()
        fut: asyncio.Future = loop.create_future()
        fut.set_result({})
        return fut

    async def on_approval_auto_approved(self) -> None:
        pass  # WS adapter does not need special handling here

    async def on_pre_stream(self) -> None:
        pass  # No spinner for WebSocket

    async def on_post_stream(self, tokens: int) -> None:
        pass  # No token tracker for WebSocket

    async def on_interview_complete(self, interview_id: str) -> None:
        await self.send(interview_complete(interview_id))

    async def on_tool_call_flush_text(self, ns_key: tuple) -> None:
        await self.end_message()

    async def on_values_update(self, data: dict) -> None:
        node_name = "state_update"
        messages = data.get("messages", [])
        state_summary = {
            "message_count": len(messages) if isinstance(messages, list) else 0,
            "has_tool_calls": any(
                hasattr(m, "tool_calls") and m.tool_calls
                for m in (messages if isinstance(messages, list) else [])
            ),
        }
        await self.send_state(node_name, state_summary)

    async def on_debug_event(self, data: dict) -> None:
        event_type = data.get("type", "unknown")
        node = None
        debug_data = {}

        if event_type == "task":
            node = data.get("name")
            debug_data = {"id": data.get("id")}
            if data.get("path"):
                debug_data["path"] = data.get("path")
        elif event_type == "checkpoint":
            debug_data = {"checkpoint_ns": data.get("checkpoint_ns")}

        await self.send_debug(event_type, node, debug_data)

    async def on_tool_running(self) -> None:
        pass  # No running animation for WebSocket


async def execute_task_websocket(
    user_input: str,
    agent: Any,
    assistant_id: str | None,
    session: Any,  # Session
    adapter: WebSocketUIAdapter,
    backend: Any = None,
) -> None:
    """Execute a task with output directed to WebSocket.

    Uses BaseStreamHandler.process_stream for the core event loop,
    delegating to the WebSocketUIAdapter callbacks for WS-specific behavior.

    Args:
        user_input: The user's input message.
        agent: The LangGraph agent to execute.
        assistant_id: The agent identifier.
        session: Session state.
        adapter: The WebSocketUIAdapter for streaming output.
        backend: Optional backend for file operations.
    """
    thread_id = session.thread_id
    adapter._thread_id = thread_id
    config = {
        "configurable": {"thread_id": thread_id},
        "metadata": {
            "assistant_id": assistant_id,
            "agent_name": assistant_id,
            "updated_at": datetime.now(UTC).isoformat(),
        }
        if assistant_id
        else {},
    }

    file_op_tracker = FileOpTracker(assistant_id=assistant_id, backend=backend)
    stream_input: dict = {"messages": [{"role": "user", "content": user_input}]}

    def on_auto_approve():
        session.auto_approve = True

    try:
        should_continue = True
        current_input: Any = stream_input
        while should_continue:
            should_continue, next_input, _tokens = await adapter.process_stream(
                agent=agent,
                stream_input=current_input,
                config=config,
                file_op_tracker=file_op_tracker,
                session_auto_approve=session.auto_approve,
                assistant_id=assistant_id,
                approval_policy=getattr(session, "approval_policy", "manual"),
                on_auto_approve_enabled=on_auto_approve,
            )
            if next_input is not None:
                current_input = next_input

    except asyncio.CancelledError:
        await adapter.on_interrupted()
        raise

    except Exception as e:
        await adapter.send_error(str(e))
        raise
