"""
KivotClient - Thin async client for the Go backend HTTP/WS API.

Provides typed access to run lifecycle operations:
  - client.runs.create()         POST /v1/runs/start
  - client.runs.cancel()         POST /v1/runs/cancel
  - client.runs.status()         GET  /v1/runs/{run_id}/status
  - client.runs.events()         GET  /v1/runs/{run_id}/events
  - client.runs.stream()         WS   /v1/runs/stream
  - client.runs.create_and_stream()  create + stream in one call
  - client.runs.approve()        POST /v1/runs/{run_id}/approval
  - client.healthz()             GET  /healthz
  - client.readyz()              GET  /readyz
  - client.wait_for_healthy()    Poll /healthz until 200

Uses aiohttp for both HTTP and WebSocket transport.
"""

import asyncio
import json
import logging
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from urllib.parse import quote

logger = logging.getLogger(__name__)

import aiohttp

from .errors import (
    KivotAPIError,
    KivotConnectionError,
    RunNotFoundError,
    StreamError,
)
from .event_types import EventType, StreamEvent, TERMINAL_EVENT_TYPES


class KivotClient:
    """
    Thin async client for the Kivot Go backend API.

    Usage::

        async with KivotClient("http://localhost:8080") as client:
            handle = await client.runs.create(agent_id="hello")
            stream = client.runs.stream(handle.run_id)
            async for event in stream:
                print(event.type_name, event.content)
    """

    def __init__(
        self,
        base_url: str = "http://localhost:8080",
        *,
        headers: Optional[Dict[str, str]] = None,
        timeout: float = 30.0,
        ws_timeout: float = 300.0,
        max_reconnect_attempts: int = 3,
        reconnect_delay: float = 1.0,
    ):
        """
        Initialize KivotClient.

        Args:
            base_url: Base URL of the Go backend (e.g. "http://localhost:8080").
            headers: Extra headers to send with every request (for auth tokens, etc.).
            timeout: HTTP request timeout in seconds.
            ws_timeout: WebSocket receive timeout in seconds.
            max_reconnect_attempts: Max reconnect attempts for WebSocket streams.
            reconnect_delay: Base delay in seconds for exponential backoff on reconnect.
        """
        self.base_url = base_url.rstrip("/")
        self.headers = dict(headers) if headers else {}
        self.timeout = aiohttp.ClientTimeout(total=timeout)
        self.ws_timeout = ws_timeout
        self.max_reconnect_attempts = max_reconnect_attempts
        self.reconnect_delay = reconnect_delay
        self._session: Optional[aiohttp.ClientSession] = None
        self.runs = _RunsNamespace(self)

    # --- Lifecycle ---

    async def __aenter__(self) -> "KivotClient":
        await self._ensure_session()
        return self

    async def __aexit__(self, *exc: Any) -> None:
        await self.close()

    async def close(self) -> None:
        """Close the underlying HTTP session."""
        if self._session and not self._session.closed:
            await self._session.close()
            self._session = None

    async def _ensure_session(self) -> aiohttp.ClientSession:
        """Create or return the aiohttp session."""
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession(
                timeout=self.timeout,
                headers=self.headers,
            )
        return self._session

    # --- Health endpoints ---

    async def healthz(self) -> bool:
        """Check GET /healthz. Returns True on 200, False otherwise."""
        try:
            session = await self._ensure_session()
            async with session.get(f"{self.base_url}/healthz") as resp:
                return resp.status == 200
        except (aiohttp.ClientError, asyncio.TimeoutError):
            return False

    async def readyz(self) -> bool:
        """Check GET /readyz. Returns True on 200, False otherwise."""
        try:
            session = await self._ensure_session()
            async with session.get(f"{self.base_url}/readyz") as resp:
                return resp.status == 200
        except (aiohttp.ClientError, asyncio.TimeoutError):
            return False

    async def wait_for_healthy(
        self,
        timeout: float = 30.0,
        interval: float = 0.5,
    ) -> None:
        """
        Poll /healthz until it returns 200 or timeout is reached.

        Args:
            timeout: Max seconds to wait.
            interval: Seconds between polls.

        Raises:
            KivotConnectionError: If the service is not healthy within timeout.
        """
        deadline = asyncio.get_event_loop().time() + timeout
        while asyncio.get_event_loop().time() < deadline:
            if await self.healthz():
                return
            await asyncio.sleep(interval)
        raise KivotConnectionError(
            f"Service not healthy after {timeout}s",
        )

    # --- Base HTTP request ---

    async def _request(
        self,
        method: str,
        path: str,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        """
        Execute an HTTP request with error mapping.

        Maps HTTP status codes to exception types:
          - 200 -> parsed JSON dict
          - 404 -> RunNotFoundError
          - 400 -> KivotAPIError
          - 413 -> KivotAPIError
          - 503 -> KivotAPIError
          - other -> KivotAPIError

        Raises:
            KivotAPIError: On 4xx/5xx responses.
            RunNotFoundError: On 404.
            KivotConnectionError: On connection failures.
        """
        session = await self._ensure_session()
        url = f"{self.base_url}{path}"

        try:
            async with session.request(method, url, **kwargs) as resp:
                body = await resp.text()

                if 200 <= resp.status < 300:
                    return json.loads(body) if body else {}

                if resp.status == 404:
                    raise RunNotFoundError(
                        f"Not found: {body}",
                        status_code=404,
                        response_body=body,
                    )

                if resp.status == 400:
                    raise KivotAPIError(
                        f"Bad request: {body}",
                        status_code=400,
                        response_body=body,
                    )

                if resp.status == 413:
                    raise KivotAPIError(
                        f"Request too large: {body}",
                        status_code=413,
                        response_body=body,
                    )

                if resp.status == 503:
                    raise KivotAPIError(
                        f"Service unavailable: {body}",
                        status_code=503,
                        response_body=body,
                    )

                raise KivotAPIError(
                    f"HTTP {resp.status}: {body}",
                    status_code=resp.status,
                    response_body=body,
                )
        except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
            raise KivotConnectionError(
                f"Connection failed: {exc}",
            ) from exc


@dataclass
class RunHandle:
    """
    Result of client.runs.create().

    Contains the run_id and ws_url returned by the server.
    The stream() convenience method creates an EventStream for this run.
    """

    run_id: str
    ws_url: str
    thread_id: Optional[str] = None
    _client: "KivotClient" = field(default=None, repr=False)

    def stream(self, *, cursor: int = 0) -> "EventStream":
        """Create an EventStream for this run."""
        return EventStream(self._client, self.run_id, cursor=cursor)


class _RunsNamespace:
    """
    Namespace for run operations: client.runs.create(), client.runs.stream(), etc.
    """

    def __init__(self, client: "KivotClient"):
        self._client = client

    async def create(
        self,
        *,
        agent_id: Optional[str] = None,
        input: Optional[Any] = None,
        model_config: Optional[Dict[str, Any]] = None,
        thread_id: Optional[str] = None,
    ) -> RunHandle:
        """
        Start a new run via POST /v1/runs/start.

        Args:
            agent_id: Optional agent identifier.
            input: Optional input data (any JSON-serializable value).
            model_config: Optional model configuration dict.
            thread_id: Optional thread identifier for multi-turn conversations.

        Returns:
            RunHandle with run_id, ws_url, thread_id, and stream() convenience method.
        """
        body: Dict[str, Any] = {}
        if agent_id is not None:
            body["agent_id"] = agent_id
        if input is not None:
            body["input"] = input
        if model_config is not None:
            body["model_config"] = model_config
        if thread_id is not None:
            body["thread_id"] = thread_id

        data = await self._client._request(
            "POST",
            "/v1/runs/start",
            json=body if body else None,
        )
        return RunHandle(
            run_id=data["run_id"],
            ws_url=data.get("ws_url", ""),
            thread_id=data.get("thread_id"),
            _client=self._client,
        )

    async def cancel(self, run_id: str) -> Dict[str, Any]:
        """
        Cancel a run via POST /v1/runs/cancel.

        Args:
            run_id: Run identifier to cancel.

        Returns:
            Dict with "ok" key on success.

        Raises:
            RunNotFoundError: If the run_id is not found.
        """
        return await self._client._request(
            "POST",
            "/v1/runs/cancel",
            json={"run_id": run_id},
        )

    async def status(self, run_id: str) -> Dict[str, Any]:
        """
        Get run status via GET /v1/runs/{run_id}/status.

        Returns:
            Dict with run_id, state, event_count, latest_cursor, etc.
        """
        return await self._client._request("GET", f"/v1/runs/{run_id}/status")

    async def events(
        self,
        run_id: str,
        *,
        cursor: int = 0,
        limit: int = 100,
    ) -> Dict[str, Any]:
        """
        Get events from the Postgres event store via GET /v1/runs/{run_id}/events.

        Args:
            run_id: Run identifier.
            cursor: Return events with seq > cursor.
            limit: Max number of events to return (max 1000).

        Returns:
            Dict with run_id, events list, count, and next_cursor.
        """
        params: Dict[str, Any] = {}
        if cursor > 0:
            params["cursor"] = cursor
        if limit != 100:
            params["limit"] = limit
        return await self._client._request(
            "GET",
            f"/v1/runs/{run_id}/events",
            params=params,
        )

    def stream(self, run_id: str, *, cursor: int = 0) -> "EventStream":
        """
        Create an EventStream for WebSocket streaming.

        Returns an async iterable that yields StreamEvent objects.
        Supports auto-reconnect with cursor-based resume.

        Args:
            run_id: Run identifier.
            cursor: Start streaming from events after this cursor.
        """
        return EventStream(self._client, run_id, cursor=cursor)

    async def create_and_stream(
        self,
        *,
        agent_id: Optional[str] = None,
        input: Optional[Any] = None,
        model_config: Optional[Dict[str, Any]] = None,
        thread_id: Optional[str] = None,
    ) -> "EventStream":
        """
        Create a run and start streaming events in one call.

        This is the most common pattern: start a run and immediately
        iterate over its events.

        Returns:
            EventStream async iterable yielding StreamEvent objects.
        """
        handle = await self.create(
            agent_id=agent_id,
            input=input,
            model_config=model_config,
            thread_id=thread_id,
        )
        return self.stream(handle.run_id)

    async def approve(
        self,
        run_id: str,
        request_id: str,
        *,
        approved: bool = True,
        decided_by: str = "",
        reason: str = "",
    ) -> Dict[str, Any]:
        """
        Submit an approval decision via POST /v1/runs/{run_id}/approval.

        Args:
            run_id: Run identifier.
            request_id: Approval request identifier.
            approved: Whether to approve or reject.
            decided_by: Who made the decision (e.g. email).
            reason: Optional reason for the decision.

        Returns:
            Dict with ok, run_id, request_id, approved keys.
        """
        body: Dict[str, Any] = {
            "request_id": request_id,
            "approved": approved,
        }
        if decided_by:
            body["decided_by"] = decided_by
        if reason:
            body["reason"] = reason

        return await self._client._request(
            "POST",
            f"/v1/runs/{run_id}/approval",
            json=body,
        )


class EventStream:
    """
    Async iterator over WebSocket stream events.

    Supports:
    - ``async for event in stream: ...``
    - Auto-reconnect with exponential backoff and cursor-based resume
    - Automatic stop on terminal event (RUN_COMPLETED, RUN_FAILED, RUN_CANCELLED)
    - Manual close via ``stream.close()``

    Usage::

        stream = client.runs.stream(run_id)
        async for event in stream:
            print(event.type_name, event.content)
            if event.is_terminal:
                break
    """

    def __init__(
        self,
        client: "KivotClient",
        run_id: str,
        cursor: int = 0,
        max_events: int = 10000,
    ):
        self._client = client
        self._run_id = run_id
        self._cursor = cursor
        self._max_events = max_events
        self._ws: Optional[aiohttp.ClientWebSocketResponse] = None
        self._closed = False
        self._events: List[StreamEvent] = []
        self._events_dropped: int = 0

    def __aiter__(self) -> "EventStream":
        return self

    async def __anext__(self) -> StreamEvent:
        """
        Get the next event from the stream.

        Reconnects automatically on connection loss using the last cursor
        for seamless resume. Raises StopAsyncIteration after a terminal
        event or when the stream is closed.
        """
        if self._closed:
            raise StopAsyncIteration

        for attempt in range(self._client.max_reconnect_attempts + 1):
            try:
                if self._ws is None or self._ws.closed:
                    await self._connect()

                msg = await self._ws.receive(timeout=self._client.ws_timeout)

                if msg.type == aiohttp.WSMsgType.TEXT:
                    event = self._parse_message(msg.data)
                    self._cursor = event.cursor
                    self._events.append(event)
                    if len(self._events) > self._max_events:
                        dropped = len(self._events) - self._max_events
                        self._events_dropped += dropped
                        self._events = self._events[-self._max_events:]
                        if self._events_dropped == dropped:
                            logger.warning(
                                "EventStream buffer overflow: dropped %d oldest events "
                                "(max_events=%d, run_id=%s)",
                                dropped, self._max_events, self._run_id,
                            )

                    if event.is_terminal:
                        self._closed = True
                        await self._close_ws()

                    return event

                if msg.type in (
                    aiohttp.WSMsgType.CLOSED,
                    aiohttp.WSMsgType.CLOSING,
                    aiohttp.WSMsgType.ERROR,
                    aiohttp.WSMsgType.CLOSE,
                ):
                    # Server closed the connection -- attempt reconnect.
                    self._ws = None
                    if attempt < self._client.max_reconnect_attempts:
                        delay = self._client.reconnect_delay * (2**attempt)
                        await asyncio.sleep(delay)
                        continue
                    # Exhausted all attempts -- stream ended normally if server
                    # closed cleanly (no more events), otherwise raise.
                    self._closed = True
                    raise StopAsyncIteration

            except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
                self._ws = None
                if attempt < self._client.max_reconnect_attempts:
                    delay = self._client.reconnect_delay * (2**attempt)
                    await asyncio.sleep(delay)
                    continue
                self._closed = True
                raise KivotConnectionError(
                    f"WebSocket failed after {attempt + 1} attempts: {exc}",
                ) from exc

        # Should not reach here, but be safe.
        self._closed = True
        raise StopAsyncIteration

    async def _connect(self) -> None:
        """Open a WebSocket connection with the current cursor for resume."""
        session = await self._client._ensure_session()
        ws_base = self._client.base_url.replace("http://", "ws://").replace(
            "https://", "wss://"
        )
        url = f"{ws_base}/v1/runs/stream?run_id={quote(self._run_id, safe='')}"
        if self._cursor > 0:
            url += f"&cursor={self._cursor}"

        self._ws = await session.ws_connect(url, heartbeat=30.0)

    def _parse_message(self, data: str) -> StreamEvent:
        """Parse a JSON WebSocket message into a StreamEvent."""
        raw = json.loads(data)
        event_data = raw.get("event", {})

        event_type = event_data.get("type", 0)
        try:
            type_name = EventType(event_type).name
        except ValueError:
            type_name = f"UNKNOWN_{event_type}"

        return StreamEvent(
            cursor=raw.get("cursor", 0),
            run_id=event_data.get("run_id", ""),
            event_id=event_data.get("event_id", ""),
            seq=event_data.get("seq", 0),
            type=event_type,
            type_name=type_name,
            payload=event_data.get("payload", {}),
            timestamp=str(event_data.get("ts", "")),
            step_id=event_data.get("step_id", ""),
            author=event_data.get("author", ""),
            trace_id=event_data.get("trace_id", ""),
            agent_id=event_data.get("agent_id", ""),
            raw=raw,
        )

    async def close(self) -> None:
        """Close the stream. Subsequent iteration will raise StopAsyncIteration."""
        self._closed = True
        await self._close_ws()

    async def _close_ws(self) -> None:
        """Close the underlying WebSocket connection."""
        if self._ws and not self._ws.closed:
            await self._ws.close()
            self._ws = None

    # --- Convenience properties ---

    @property
    def events(self) -> List[StreamEvent]:
        """All events received so far (copy)."""
        return list(self._events)

    @property
    def last_cursor(self) -> int:
        """Current cursor position (for manual resume)."""
        return self._cursor

    @property
    def events_dropped(self) -> int:
        """Total events dropped due to buffer overflow."""
        return self._events_dropped
