"""WebSocket adapter for enterprise AI platform workflow execution.

Uses a simple open-send-receive-close pattern:
  1. Open WebSocket with JWT auth
  2. Send SendMessageV2 payload
  3. Collect all streaming responses until agent_execution_stats arrives
  4. Close WebSocket

A lightweight heartbeat runs ONLY during the active execution window
(enterprise platforms expect pings every ~27s; workflows can run 60-90s+).
No persistent connections, no reconnection logic.

SECURITY: JWT token is held in memory only, sent via WSS (TLS),
and never logged or persisted.
"""

from __future__ import annotations

import asyncio
import json
import logging
import time
from collections.abc import AsyncGenerator
from typing import Any

from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

class WebSocketConfig(BaseModel):
    """Configuration for a single workflow execution over WebSocket."""

    url_template: str = Field(
        description="WebSocket URL with {token} placeholder. "
        "Example: wss://gateway.example.com/prod?Authorization={token}"
    )
    heartbeat_interval: int = Field(
        default=25,
        description="Heartbeat interval in seconds during execution "
        "(slightly under the server's 27s window to be safe)",
    )
    execution_timeout: float = Field(
        default=600.0,
        description="Max seconds to wait for a workflow to complete before giving up",
    )


# ---------------------------------------------------------------------------
# Stream Buffer -- accumulates all messages from one execution
# ---------------------------------------------------------------------------

class StreamBuffer:
    """Accumulates streaming answer chunks and tool calls into a complete result."""

    def __init__(self) -> None:
        self._chunks: list[str] = []
        self._tool_requests: list[dict[str, Any]] = []
        self._tool_responses: list[dict[str, Any]] = []
        self._execution_stats: dict[str, Any] | None = None
        self._reasoning: dict[str, Any] | None = None
        self._is_complete: bool = False
        self._start_time: float = time.time()

    def add_chunk(self, text: str) -> None:
        """Add a streaming answer chunk."""
        self._chunks.append(text)

    def add_tool_request(self, request: dict[str, Any]) -> None:
        """Record a tool invocation request."""
        self._tool_requests.append(request)

    def add_tool_response(self, response: dict[str, Any]) -> None:
        """Record a tool invocation response."""
        self._tool_responses.append(response)

    def set_execution_stats(self, stats: dict[str, Any]) -> None:
        """Set the final execution statistics -- marks execution as complete."""
        self._execution_stats = stats
        self._is_complete = True

    def set_reasoning(self, reasoning: dict[str, Any]) -> None:
        """Set the agent reasoning traces."""
        self._reasoning = reasoning

    @property
    def full_text(self) -> str:
        """Get the complete accumulated answer text."""
        return "".join(self._chunks)

    @property
    def is_complete(self) -> bool:
        return self._is_complete

    @property
    def duration(self) -> float:
        return time.time() - self._start_time

    @property
    def tool_call_count(self) -> int:
        return len(self._tool_requests)

    @property
    def error_count(self) -> int:
        return sum(
            1
            for r in self._tool_responses
            if r.get("tool_status") == "error"
        )

    def to_dict(self) -> dict[str, Any]:
        """Export the complete execution result."""
        return {
            "answer": self.full_text,
            "tool_calls": {
                "requests": self._tool_requests,
                "responses": self._tool_responses,
                "total": len(self._tool_requests),
                "errors": self.error_count,
            },
            "execution_stats": self._execution_stats,
            "reasoning": self._reasoning,
            "duration_seconds": round(self.duration, 2),
            "is_complete": self._is_complete,
        }


# ---------------------------------------------------------------------------
# Parsed message envelope
# ---------------------------------------------------------------------------

class WebSocketMessage(BaseModel):
    """A single parsed WebSocket message."""

    direction: str = Field(description="'send' or 'receive'")
    message_type: str = Field(default="unknown")
    agent_id: str | None = Field(default=None)
    data: dict[str, Any] = Field(default_factory=dict)
    timestamp: float = Field(default_factory=time.time)
    raw: str = Field(default="", exclude=True)  # excluded from serialisation by default


# ---------------------------------------------------------------------------
# Message parser (stateless, reusable)
# ---------------------------------------------------------------------------

def parse_ws_message(raw: str, direction: str = "receive") -> WebSocketMessage:
    """Parse a raw WebSocket JSON frame into a typed WebSocketMessage.

    Recognises the following server→client message shapes:
      {"message": "pong"}                       → heartbeat_response
      {"agent_<ID>": {"answer": "..."}}          → answer_chunk
      {"agent_<ID>": {"tool_use_request": {}}}   → tool_use_request
      {"agent_<ID>": {"tool_use_response": {}}}  → tool_use_response
      {"agent_<ID>": {"agent_execution_stats":…}} → execution_stats
      {"agent_<ID>": {"agent_reasoning": …}}     → agent_reasoning
    """
    try:
        data = json.loads(raw)
    except json.JSONDecodeError:
        return WebSocketMessage(direction=direction, raw=raw)

    if not isinstance(data, dict):
        return WebSocketMessage(direction=direction, raw=raw, data={"value": data})

    # Heartbeat pong
    if data.get("message") == "pong":
        return WebSocketMessage(direction=direction, message_type="heartbeat_response", raw=raw)

    # Agent-keyed messages
    for key, value in data.items():
        if not key.startswith("agent_") or not isinstance(value, dict):
            continue

        if "answer" in value:
            return WebSocketMessage(
                direction=direction,
                message_type="answer_chunk",
                agent_id=key,
                data={"text": value["answer"]},
                raw=raw,
            )
        if "tool_use_request" in value:
            return WebSocketMessage(
                direction=direction,
                message_type="tool_use_request",
                agent_id=key,
                data=value["tool_use_request"],
                raw=raw,
            )
        if "tool_use_response" in value:
            return WebSocketMessage(
                direction=direction,
                message_type="tool_use_response",
                agent_id=key,
                data=value["tool_use_response"],
                raw=raw,
            )
        if "agent_execution_stats" in value:
            return WebSocketMessage(
                direction=direction,
                message_type="execution_stats",
                agent_id=key,
                data=value["agent_execution_stats"],
                raw=raw,
            )
        if "agent_reasoning" in value:
            return WebSocketMessage(
                direction=direction,
                message_type="agent_reasoning",
                agent_id=key,
                data=value["agent_reasoning"],
                raw=raw,
            )
        break  # only process the first agent_ key

    return WebSocketMessage(direction=direction, data=data, raw=raw)


# ---------------------------------------------------------------------------
# WebSocket Adapter -- open / execute / close
# ---------------------------------------------------------------------------

class WebSocketAdapter:
    """Execute a single workflow over a transient WebSocket connection.

    Lifecycle:
        adapter = WebSocketAdapter(config)
        result  = await adapter.run(token, workflow_id, query, ...)
        # connection is now closed; result contains the full answer

    For streaming access, use ``run_streaming`` which yields messages
    as they arrive and closes the socket when done.

    SECURITY: The JWT token is only held in memory and transmitted over
    WSS (TLS). It is never logged or written to disk.
    """

    def __init__(self, config: WebSocketConfig) -> None:
        self._config = config

    # ------------------------------------------------------------------
    # Public API -- blocking (returns final result dict)
    # ------------------------------------------------------------------

    async def run(
        self,
        token: str,
        workflow_id: str,
        query: str,
        modelparams: dict[str, dict[str, str]] | None = None,
        input_variables: dict[str, str] | None = None,
        conversation_id: str | None = None,
        is_persistence_allowed: bool = False,
    ) -> dict[str, Any]:
        """Execute a workflow and return the complete result.

        Opens a WebSocket, sends ``SendMessageV2``, collects all streamed
        messages until ``agent_execution_stats`` arrives, then closes the
        connection and returns the accumulated result.

        Args:
            token: Platform JWT. NEVER logged.
            workflow_id: UUID of the workflow chain.
            query: User prompt / instruction.
            modelparams: Per-component auth tokens and settings.
            input_variables: Template variable values from the UI.
            conversation_id: For continuing a previous conversation.
            is_persistence_allowed: Whether the platform should persist history.

        Returns:
            Dict with keys: answer, tool_calls, execution_stats,
            reasoning, duration_seconds, is_complete.
        """
        async for _msg in self.run_streaming(
            token=token,
            workflow_id=workflow_id,
            query=query,
            modelparams=modelparams,
            input_variables=input_variables,
            conversation_id=conversation_id,
            is_persistence_allowed=is_persistence_allowed,
        ):
            pass  # run_streaming feeds _last_result internally

        return self._last_result or {"answer": "", "is_complete": False}

    # ------------------------------------------------------------------
    # Public API -- streaming (yields messages as they arrive)
    # ------------------------------------------------------------------

    async def run_streaming(
        self,
        token: str,
        workflow_id: str,
        query: str,
        modelparams: dict[str, dict[str, str]] | None = None,
        input_variables: dict[str, str] | None = None,
        conversation_id: str | None = None,
        is_persistence_allowed: bool = False,
    ) -> AsyncGenerator[WebSocketMessage, None]:
        """Execute a workflow, yielding parsed messages as they stream in.

        The WebSocket is opened before the first yield and automatically
        closed after the execution completes (or times out).

        Yields:
            WebSocketMessage for each received frame.
        """
        try:
            import websockets  # type: ignore[import-untyped]
        except ImportError:
            logger.error("websockets package required. Run: pip install websockets")
            raise RuntimeError("Missing dependency: websockets")

        url = self._config.url_template.replace("{token}", token)
        safe_url = self._config.url_template.replace("{token}", "<REDACTED>")
        logger.info("Opening WebSocket: %s", safe_url)

        buf = StreamBuffer()
        heartbeat_task: asyncio.Task[None] | None = None

        async with websockets.connect(
            url,
            ping_interval=None,  # we send our own heartbeat
            close_timeout=10,
            max_size=10 * 1024 * 1024,  # 10 MB
        ) as ws:
            logger.info("WebSocket connected — sending workflow execution")

            # --- lightweight heartbeat (only during this execution) --------
            async def _heartbeat() -> None:
                try:
                    while True:
                        await asyncio.sleep(self._config.heartbeat_interval)
                        await ws.send(json.dumps({"action": "heartbeat"}))
                        logger.debug("Heartbeat ping sent")
                except asyncio.CancelledError:
                    pass
                except Exception as exc:
                    logger.warning("Heartbeat stopped: %s", exc)

            heartbeat_task = asyncio.create_task(_heartbeat())

            # --- send the execution payload --------------------------------
            payload: dict[str, Any] = {
                "action": "SendMessageV2",
                "workflow_id": workflow_id,
                "query": query,
                "is_persistence_allowed": is_persistence_allowed,
                "modelparams": modelparams or {},
                "input_variables": input_variables or {},
                "conversation_id": conversation_id,
            }
            await ws.send(json.dumps(payload))
            logger.info("SendMessageV2 dispatched (workflow_id=%s)", workflow_id)

            # --- receive loop ----------------------------------------------
            try:
                async for raw_frame in ws:
                    msg = parse_ws_message(str(raw_frame), direction="receive")

                    # Skip heartbeat pongs silently
                    if msg.message_type == "heartbeat_response":
                        continue

                    # Route into buffer
                    if msg.message_type == "answer_chunk":
                        buf.add_chunk(msg.data.get("text", ""))
                    elif msg.message_type == "tool_use_request":
                        buf.add_tool_request(msg.data)
                    elif msg.message_type == "tool_use_response":
                        buf.add_tool_response(msg.data)
                    elif msg.message_type == "execution_stats":
                        buf.set_execution_stats(msg.data)
                    elif msg.message_type == "agent_reasoning":
                        buf.set_reasoning(msg.data)

                    yield msg

                    if buf.is_complete:
                        logger.info(
                            "Workflow complete: duration=%.1fs, tool_calls=%d, errors=%d",
                            buf.duration,
                            buf.tool_call_count,
                            buf.error_count,
                        )
                        break

            except asyncio.TimeoutError:
                logger.warning(
                    "Execution timed out after %.0fs", self._config.execution_timeout
                )

            finally:
                # Cancel heartbeat — connection is about to close
                if heartbeat_task is not None:
                    heartbeat_task.cancel()
                    try:
                        await heartbeat_task
                    except asyncio.CancelledError:
                        pass

        logger.info("WebSocket closed")

        # Store the buffer on self so callers of `run()` can access it
        self._last_result = buf.to_dict()

    # ------------------------------------------------------------------
    # Convenience
    # ------------------------------------------------------------------

    @property
    def last_result(self) -> dict[str, Any] | None:
        """Result from the most recent ``run`` or ``run_streaming`` call."""
        return getattr(self, "_last_result", None)
