"""
WebSocket client — connects to the SigmanticAI backend.

Handles:
  - Connection with authentication
  - Auto-reconnect with event replay
  - Message sending (text, file uploads)
  - Async event iterator for the display layer
"""

from __future__ import annotations

import asyncio
import json
import logging
import ssl
import time
from typing import Any, AsyncIterator, Dict, Optional

import sys

import certifi
import websockets
from websockets.exceptions import (
    ConnectionClosed,
    ConnectionClosedError,
    ConnectionClosedOK,
    InvalidStatusCode,
)

from sigmanticai.config import API_URL

logger = logging.getLogger("sigmanticai.client")


def _status(msg: str) -> None:
    """Print a reconnect status line visible to the user."""
    print(f"\033[33m{msg}\033[0m", file=sys.stderr, flush=True)

# Reconnect settings
MAX_RECONNECT_ATTEMPTS = 10
RECONNECT_BASE_DELAY = 1.0  # seconds
RECONNECT_MAX_DELAY = 30.0
_SESSION_LOCK_RETRY_DELAY = 5.0  # wait for old gateway's session lock to expire


class SessionMigrating(Exception):
    """Raised when the server returns 4009 (session lock held by old gateway)."""
    pass


class SigmanticClient:
    """
    WebSocket client for the SigmanticAI backend.

    Usage:
        client = SigmanticClient(access_token="...")
        await client.connect()

        # Send a message
        await client.send_message("Build verification for a FIFO")

        # Iterate over events
        async for event in client.events():
            if event["type"] == "text":
                print(event["content"])
            elif event["type"] == "done":
                break

        await client.close()
    """

    def __init__(
        self,
        access_token: str,
        api_url: str = API_URL,
        session_id: Optional[str] = None,
    ):
        self._access_token = access_token
        self._api_url = api_url
        self._session_id = session_id
        self._ws: Optional[websockets.WebSocketClientProtocol] = None
        self._last_event_id: Optional[str] = None
        self._connected = False
        self._closing = False

    @property
    def session_id(self) -> Optional[str]:
        return self._session_id

    async def connect(self) -> Dict[str, Any]:
        """
        Connect and authenticate.

        Returns the session info dict from the server.
        Raises ConnectionError on failure.
        """
        try:
            # Build SSL context with certifi's CA bundle so connections
            # work on macOS even when system certs aren't configured
            # (e.g. python.org Python without "Install Certificates").
            ssl_ctx = ssl.create_default_context(cafile=certifi.where())

            self._ws = await websockets.connect(
                self._api_url,
                ssl=ssl_ctx,
                ping_interval=30,
                ping_timeout=1800,
                close_timeout=10,
                max_size=10 * 1024 * 1024,  # 10 MB max message
            )
        except Exception as e:
            raise ConnectionError(f"Failed to connect to {self._api_url}: {e}")

        # Send auth message
        auth_msg = {"api_key": self._access_token}
        if self._session_id:
            auth_msg["session_id"] = self._session_id
        if self._last_event_id:
            auth_msg["after"] = self._last_event_id

        await self._ws.send(json.dumps(auth_msg))

        # Wait for session response
        try:
            raw = await asyncio.wait_for(self._ws.recv(), timeout=30)
        except ConnectionClosedError as e:
            if e.rcvd and e.rcvd.code == 4009:
                raise SessionMigrating(
                    "Session is being migrated to a new server — retrying..."
                )
            reason = e.reason or "Invalid or expired credentials"
            raise ConnectionError(reason)
        session_info = json.loads(raw)

        if session_info.get("type") == "error":
            msg = session_info.get("message", "Unknown error")
            raise ConnectionError(msg)

        self._session_id = session_info.get("session_id", self._session_id)
        self._connected = True
        return session_info

    async def send_message(self, text: str) -> None:
        """Send a user message to the agent."""
        if not self._ws or not self._connected:
            raise RuntimeError("Not connected")
        await self._ws.send(json.dumps({"type": "message", "text": text}))

    async def send_cancel(self) -> None:
        """Send a cancel signal to abort the current response without disconnecting."""
        if not self._ws or not self._connected:
            return
        await self._ws.send(json.dumps({"type": "cancel"}))

    async def send_file(self, filename: str, content: str) -> None:
        """Upload a file (e.g., RTL for --rtl flag)."""
        if not self._ws or not self._connected:
            raise RuntimeError("Not connected")
        await self._ws.send(
            json.dumps({"type": "upload", "filename": filename, "content": content})
        )

    async def send_delete(self, filename: str) -> None:
        """Notify the server that a local file was deleted."""
        if not self._ws or not self._connected:
            raise RuntimeError("Not connected")
        await self._ws.send(
            json.dumps({"type": "delete", "filename": filename})
        )

    async def send_tool_response(self, request_id: str, result: str) -> None:
        """Send the result of a locally-executed tool call back to the server."""
        if not self._ws or not self._connected:
            raise RuntimeError("Not connected")
        await self._ws.send(json.dumps({
            "type": "tool_response",
            "request_id": request_id,
            "result": result[:500000],
        }))

    async def send_project_context(
        self, tree: str, tool_config: str = "", cache: str = "",
        os_info: str = "", checkpoints: str = "",
    ) -> None:
        """Send project directory tree, tool config, cache, and OS info to the server."""
        if not self._ws or not self._connected:
            raise RuntimeError("Not connected")
        msg: dict = {
            "type": "project_context",
            "tree": tree[:100000],
        }
        if tool_config:
            msg["tool_config"] = tool_config[:10000]
        if cache:
            msg["cache"] = cache[:500000]
        if os_info:
            msg["os_info"] = os_info[:200]
        if checkpoints:
            msg["checkpoints"] = checkpoints[:500000]
        await self._ws.send(json.dumps(msg))

    async def send_command_response(
        self, request_id: str, stdout: str, stderr: str, exit_code: int
    ) -> None:
        """Send the result of a locally-executed command back to the server."""
        if not self._ws or not self._connected:
            raise RuntimeError("Not connected")
        await self._ws.send(json.dumps({
            "type": "command_response",
            "request_id": request_id,
            "stdout": stdout[:500000],
            "stderr": stderr[:100000],
            "exit_code": exit_code,
        }))

    async def events(self) -> AsyncIterator[Dict[str, Any]]:
        """
        Async iterator over events from the server.

        Handles auto-reconnect transparently. Yields event dicts.
        """
        reconnect_attempts = 0

        while not self._closing:
            try:
                async for raw_msg in self._ws:
                    event = json.loads(raw_msg)
                    event_type = event.get("type", "")

                    # Track last event ID for replay on reconnect
                    if "id" in event:
                        self._last_event_id = event["id"]

                    # Handle pings
                    if event_type == "ping":
                        await self._ws.send(json.dumps({"type": "pong"}))
                        continue

                    # Reset reconnect counter on data
                    reconnect_attempts = 0
                    yield event

                    # If server says done or failed, stop
                    if event_type in ("done", "failed"):
                        return

            except ConnectionClosed as cc:
                if self._closing:
                    return

                reason = str(cc)
                is_timeout = "keepalive" in reason.lower() or "ping timeout" in reason.lower()

                if is_timeout:
                    logger.warning(
                        "Connection timed out (idle too long). "
                        "Run `sigmanticai` again to resume your session."
                    )
                    raise ConnectionError(
                        "Connection timed out — the session was idle too long. "
                        "Run `sigmanticai` to reconnect and keep building "
                        "(your progress is saved)."
                    )

                _status("Connection lost — reconnecting (your progress is saved)...")
                self._connected = False

                if reconnect_attempts >= MAX_RECONNECT_ATTEMPTS:
                    raise ConnectionError(
                        f"Lost connection after {MAX_RECONNECT_ATTEMPTS} reconnect attempts. "
                        f"Run `sigmanticai` to reconnect and resume your session."
                    )

                # Exponential backoff
                delay = min(
                    RECONNECT_BASE_DELAY * (2**reconnect_attempts),
                    RECONNECT_MAX_DELAY,
                )
                reconnect_attempts += 1
                _status(f"Reconnect attempt {reconnect_attempts}/{MAX_RECONNECT_ATTEMPTS} in {delay:.0f}s...")
                await asyncio.sleep(delay)

                try:
                    await self.connect()
                    _status("Reconnected! Resuming session...")
                except SessionMigrating:
                    _status("Server update in progress — waiting for session migration...")
                    reconnect_attempts -= 1  # don't count migration retries
                    await asyncio.sleep(_SESSION_LOCK_RETRY_DELAY)
                    continue
                except Exception as e:
                    _status(f"Reconnect failed: {e}")
                    continue

            except Exception as e:
                if self._closing:
                    return
                logger.error("Event stream error: %s", e)
                raise

    async def close(self) -> None:
        """Close the WebSocket connection."""
        self._closing = True
        self._connected = False
        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass


