"""Authenticated MCP orchestration around the unchanged speech wire protocol."""

import asyncio
import json
from dataclasses import dataclass, field
from time import monotonic
from typing import Any

import websockets
from fastapi import WebSocket, WebSocketDisconnect
from starlette.websockets import WebSocketState
from websockets.asyncio.client import ClientConnection
from websockets.typing import Subprotocol

from vocero_api.core import config as settings
from vocero_api.core.logging import logger
from vocero_api.core.urls import _proxy_upstream_url
from vocero_api.models.mcp import IntegrationError
from vocero_api.services import content_context, conversations, mcp_tools

_TEXT_DELTA = {
    "response.text.delta",
    "response.output_text.delta",
    "response.audio_transcript.delta",
    "response.output_audio_transcript.delta",
}
_TEXT_DONE = {
    "response.text.done",
    "response.output_text.done",
    "response.audio_transcript.done",
    "response.output_audio_transcript.done",
}
_AUDIO = {"response.audio.delta", "response.output_audio.delta", "response.audio.done", "response.output_audio.done"}


@dataclass
class ResponseBuffer:
    generation: int
    text: str = ""
    segment: str = ""
    frames: list[str | bytes] = field(default_factory=list)
    size: int = 0
    plain: bool = False

    def collect(self, raw: str | bytes, event: dict[str, Any]) -> None:
        kind = event.get("type")
        if kind in _TEXT_DELTA:
            self.segment += event.get("delta", "")
        elif kind in _TEXT_DONE:
            text = event.get("transcript") or event.get("text") or self.segment
            if isinstance(text, str) and text:
                self.text = (self.text + " " + text).strip()
            self.segment = ""
        candidate = (self.text + self.segment).lstrip()
        self.plain = bool(
            candidate
            and "<" not in candidate
            and not candidate.startswith(">")
            and "VOCERO_MCP" not in candidate.upper()
            and not "VOCERO_MCP".startswith(candidate[:10].upper())
        )
        self.frames.append(raw)
        self.size += len(raw)
        if self.size > 8_000_000 or len(candidate) > 64000:
            raise IntegrationError(502, "La respuesta de voz excedió el tamaño permitido.")


class McpRelay:
    def __init__(
        self, browser: WebSocket, upstream: ClientConnection, conversation: conversations.Conversation
    ) -> None:
        self.browser = browser
        self.upstream = upstream
        self.conversation = conversation
        self.turn = mcp_tools.ToolTurn(conversation, conversation.generation)
        self.responses: dict[str, ResponseBuffer] = {}
        self.finished: set[str] = set()
        self.active_id = ""
        self.accept_response = True
        self.upstream_lock = asyncio.Lock()
        self.browser_lock = asyncio.Lock()
        self.tool_tasks: set[asyncio.Task[None]] = set()
        self.base_instructions = settings.S2S_PRIMARY_PROMPT or ""

    def instructions_event(self) -> str:
        return json.dumps(
            {
                "type": "session.update",
                "session": {"instructions": mcp_tools.instructions(self.base_instructions, self.conversation)},
            },
            ensure_ascii=False,
        )

    async def send_browser(self, raw: str | bytes) -> None:
        async with self.browser_lock:
            try:
                if isinstance(raw, bytes):
                    await self.browser.send_bytes(raw)
                else:
                    await self.browser.send_text(raw)
            except (RuntimeError, WebSocketDisconnect):
                raise asyncio.CancelledError from None

    async def browser_event(self, event: dict[str, Any]) -> None:
        await self.send_browser(json.dumps(event, ensure_ascii=False))

    async def send_upstream(self, raw: str | bytes) -> None:
        async with self.upstream_lock:
            await self.upstream.send(raw)

    def new_turn(self, *, speaking: bool = False) -> None:
        self.turn = mcp_tools.ToolTurn(self.conversation, self.conversation.interrupt())
        self.accept_response = not speaking

    async def flush(self, buffer: ResponseBuffer) -> None:
        for frame in buffer.frames:
            await self.send_browser(frame)
        buffer.frames.clear()
        buffer.size = 0

    async def tool_call(self, name: str, arguments: dict[str, Any], turn: mcp_tools.ToolTurn) -> None:
        try:
            await self.browser_event({"type": "vocero.tool.started", "name": name})
            previous_content = self.conversation.current_content
            result = await turn.call(name, arguments)
            async with self.upstream_lock:
                turn.require_current()
                if previous_content != self.conversation.current_content:
                    await self.upstream.send(self.instructions_event())
                await self.upstream.send(
                    json.dumps(
                        {
                            "type": "conversation.item.create",
                            "item": {
                                "type": "message",
                                "role": "user",
                                "content": [{"type": "input_text", "text": mcp_tools.result_text(name, result)}],
                            },
                        },
                        ensure_ascii=False,
                    )
                )
                turn.require_current()
                await self.upstream.send(json.dumps({"type": "response.create"}))
        except IntegrationError as exc:
            if turn.generation != self.conversation.generation:
                raise asyncio.CancelledError from None
            await self.browser_event({"type": "error", "error": {"code": "mcp_error", "message": exc.message}})
            await self.browser_event({"type": "vocero.tool.failed"})
            if exc.status == 401:
                conversations.store.remove(self.conversation.id)
        except websockets.ConnectionClosed:
            pass
        except Exception as exc:
            logger.warning("mcp_relay_tool_failed error_type=%s", type(exc).__name__)
            await self.browser_event({"type": "error", "error": {"message": "No se pudo completar la consulta."}})
            await self.browser_event({"type": "vocero.tool.failed"})
        finally:
            if self.conversation.pending is asyncio.current_task():
                self.conversation.pending = None

    async def client_to_speech(self) -> None:
        while True:
            self.conversation.require_active()
            message = await self.browser.receive()
            if message["type"] == "websocket.disconnect":
                return
            if message.get("bytes") is not None:
                await self.send_upstream(message["bytes"])
                continue
            raw = message.get("text")
            if raw is None:
                continue
            try:
                event = json.loads(raw)
            except ValueError:
                await self.send_upstream(raw)
                continue
            if isinstance(event, dict):
                kind = event.get("type")
                if settings.S2S_REASONING_TYPE == "PROXY":
                    if kind == "response.cancel":
                        self.new_turn()
                    item = event.get("item")
                    if kind == "conversation.item.create" and isinstance(item, dict) and item.get("role") == "user":
                        self.new_turn()
                    session = event.get("session")
                    if kind == "session.update" and isinstance(session, dict):
                        # Browser instructions also carry the current session's RAG notes.
                        self.base_instructions = str(session.get("instructions") or settings.S2S_PRIMARY_PROMPT or "")
                        session["instructions"] = mcp_tools.instructions(self.base_instructions, self.conversation)
                        raw = json.dumps(event, ensure_ascii=False)
            await self.send_upstream(raw)

    async def speech_to_client(self) -> None:
        async for raw in self.upstream:
            self.conversation.require_active()
            if settings.S2S_REASONING_TYPE != "PROXY":
                await self.send_browser(raw)
                continue
            if isinstance(raw, bytes):
                buffer = self.responses.get(self.active_id)
                if buffer is None:
                    await self.send_browser(raw)
                elif buffer.generation == self.conversation.generation:
                    buffer.collect(raw, {})
                    if buffer.plain:
                        await self.flush(buffer)
                continue
            try:
                event = json.loads(raw)
            except ValueError:
                await self.send_browser(raw)
                continue
            if not isinstance(event, dict):
                await self.send_browser(raw)
                continue
            kind = event.get("type", "")
            if kind == "input_audio_buffer.speech_started":
                self.new_turn(speaking=True)
            if kind == "conversation.item.input_audio_transcription.completed":
                self.accept_response = True
            response = event.get("response")
            response = response if isinstance(response, dict) else {}
            response_id = event.get("response_id") or response.get("id") or self.active_id
            if kind == "response.created":
                if len(self.responses) >= 8:
                    raise IntegrationError(502, "El servicio de voz envió demasiadas respuestas simultáneas.")
                self.active_id = str(response_id)
                self.responses.setdefault(
                    self.active_id, ResponseBuffer(self.conversation.generation if self.accept_response else -1)
                )
                if response.get("output") is None:
                    response["output"] = []
                await self.browser_event(event)
                continue
            if response_id in self.finished and kind.startswith("response."):
                continue
            protected = kind in _TEXT_DELTA | _TEXT_DONE | _AUDIO or kind in {
                "response.content_part.added",
                "response.content_part.done",
                "response.output_item.added",
                "response.output_item.done",
                "response.done",
            }
            if not protected:
                await self.send_browser(raw)
                continue
            # Some compatible speech servers emit text before response.created.
            response_id = str(response_id or "unlabelled")
            buffer = self.responses.setdefault(
                response_id, ResponseBuffer(self.conversation.generation if self.accept_response else -1)
            )
            if kind == "response.done":
                self.responses.pop(response_id, None)
                self.finished.add(response_id)
                if len(self.finished) > 256:
                    self.finished = {response_id}
                if self.active_id == response_id:
                    self.active_id = ""
                full_text = (
                    " ".join(
                        str(part.get("transcript") or part.get("text") or "")
                        for item in response.get("output") or []
                        if isinstance(item, dict)
                        for part in item.get("content") or []
                        if isinstance(part, dict)
                    ).strip()
                    or (buffer.text + buffer.segment).strip()
                )
                chosen = None
                error = None
                if (
                    buffer.generation == self.conversation.generation
                    and response.get("status", "completed") == "completed"
                ):
                    try:
                        chosen = mcp_tools.selection(full_text)
                    except IntegrationError as exc:
                        error = exc
                stale = buffer.generation != self.conversation.generation
                interrupted_marker = response.get("status", "completed") != "completed" and (
                    not buffer.plain or "VOCERO_MCP" in full_text.upper()
                )
                if chosen or error or stale or interrupted_marker:
                    response["output"] = []
                    await self.browser_event({"type": "response.done", "response": response})
                    if chosen:
                        if self.conversation.pending is not None:
                            raise IntegrationError(502, "El servicio de voz solicitó consultas simultáneas.")
                        task = asyncio.create_task(self.tool_call(*chosen, self.turn))
                        self.conversation.pending = task
                        self.tool_tasks.add(task)
                        task.add_done_callback(self.tool_tasks.discard)
                    if error:
                        await self.browser_event({"type": "error", "error": {"message": error.message}})
                        await self.browser_event({"type": "vocero.tool.failed"})
                else:
                    await self.flush(buffer)
                    if response.get("output") is None:
                        response["output"] = []
                    await self.browser_event(event)
                    if self.conversation.greeting_pending and response.get("status", "completed") == "completed":
                        self.conversation.greeting_pending = False
                        await self.send_upstream(self.instructions_event())
                continue
            if buffer.generation == self.conversation.generation:
                buffer.collect(raw, event)
                if buffer.plain:
                    await self.flush(buffer)


async def realtime_mcp_proxy(ws: WebSocket) -> None:
    conversation: conversations.Conversation | None = None
    claimed = False
    tasks: list[asyncio.Task[None]] = []
    relay: McpRelay | None = None
    try:
        conversation = conversations.resolve(ws)
        if not settings.S2S_PROXY_TARGET:
            raise IntegrationError(503, "El proxy de voz no está configurado.")
        if conversation.socket_attached:
            raise IntegrationError(409, "Esta conversación ya tiene una conexión de voz.")
        conversation.socket_attached = True
        claimed = True
        await ws.accept()
        if settings.S2S_REASONING_TYPE == "PROXY":
            conversation.catalog = await content_context.catalog(conversation.backend_token)
            conversation.require_active()
        # Only the existing speech protocols cross the upstream boundary.
        protocols = [
            Subprotocol(part.strip())
            for part in ws.headers.get("sec-websocket-protocol", "").split(",")
            if part.strip()
        ]
        async with websockets.connect(
            _proxy_upstream_url(ws), subprotocols=protocols or None, open_timeout=15, max_size=8_000_000
        ) as upstream:
            relay = McpRelay(ws, upstream, conversation)

            async def expired_or_closed() -> None:
                try:
                    await asyncio.wait_for(
                        conversation.closed.wait(), timeout=max(0, conversation.expires_at - monotonic())
                    )
                except TimeoutError:
                    await relay.browser_event(
                        {"type": "error", "error": {"message": "La conversación expiró. Inicia una nueva sesión."}}
                    )

            tasks = [
                asyncio.create_task(relay.client_to_speech()),
                asyncio.create_task(relay.speech_to_client()),
                asyncio.create_task(expired_or_closed()),
            ]
            done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
            for task in done:
                task.result()
    except IntegrationError as exc:
        if exc.status == 401 and claimed and conversation is not None:
            conversations.store.remove(conversation.id)
        if ws.application_state == WebSocketState.CONNECTING:
            await ws.accept()
        await ws.send_json({"type": "error", "error": {"code": "mcp_error", "message": exc.message}})
    except (WebSocketDisconnect, websockets.ConnectionClosed):
        pass
    except (OSError, TimeoutError):
        await ws.send_json({"type": "error", "error": {"message": "No se pudo conectar con el servicio de voz."}})
    finally:
        for task in tasks:
            task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)
        if relay is not None:
            pending_tools = list(relay.tool_tasks)
            for task in pending_tools:
                task.cancel()
            await asyncio.gather(*pending_tools, return_exceptions=True)
        if claimed and conversation is not None:
            pending = conversation.pending
            conversations.store.remove(conversation.id)
            if pending is not None:
                await asyncio.gather(pending, return_exceptions=True)
        try:
            await ws.close()
        except (RuntimeError, WebSocketDisconnect):
            pass