"""LLMHosts Relay tunnel provider.

Primary tunnel provider: connects to the LLMHosts relay server via WebSocket
using the Rust ``_llmhosts.relay`` extension (pyo3 + maturin). Falls back gracefully
if the Rust extension is not built — the provider simply won't appear in the
detector cascade.

Protocol stack:
  WSS → Noise NK encryption (M2) → yamux multiplexing (M3) → HTTP forwarding

M2 behaviour: set ``RELAY_SERVER_PUBKEY`` environment variable (or config key
``relay.server_pubkey``) to the 64-char hex Curve25519 public key of your relay
server. When set, all relay traffic is end-to-end encrypted — the VPS cannot
read request/response payloads (zero-knowledge relay).

When ``RELAY_SERVER_PUBKEY`` is not set, falls back to M1 (unencrypted WSS).
This allows development without a keypair and backward compatibility during
relay server rollout.

Reconnect behaviour: after a successful connection, a background asyncio task
monitors the connection every 5 seconds. On disconnect, it retries with
exponential backoff (1s → 2s → 4s → 8s → 16s → 30s cap) plus ±20% random
jitter to prevent thundering herd on relay server restarts.
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import os
import random
from datetime import datetime, timezone

from llmhosts.tunnel.models import TunnelProvider, TunnelStatus

logger = logging.getLogger(__name__)

# --- Rust extension import (graceful degradation) --------------------------

RELAY_AVAILABLE: bool = False
_RelayClient: type | None = None
RELAY_VERSION: str = "0.0.0"

try:
    from _llmhosts import relay as _relay_core_mod  # type: ignore[import]

    _RelayClient = _relay_core_mod.RelayClient
    RELAY_VERSION = _relay_core_mod.RELAY_VERSION
    RELAY_AVAILABLE = True
    logger.debug("Rust relay core loaded (_llmhosts.relay %s)", RELAY_VERSION)
except ImportError:
    # Fallback: try standalone module (development mode with individual maturin develop)
    try:
        import _relay_core  # type: ignore[import]

        _RelayClient = _relay_core.RelayClient
        RELAY_VERSION = _relay_core.RELAY_VERSION
        RELAY_AVAILABLE = True
        logger.debug("Rust relay core loaded (standalone _relay_core %s)", RELAY_VERSION)
    except ImportError:
        logger.debug("Rust relay core not available — relay provider disabled")

# Default relay server URL
DEFAULT_RELAY_URL = "wss://llmhosts.com/api/relay/connect"

# Reconnect backoff schedule in seconds (exponential, capped at 30s)
_RECONNECT_BACKOFF: list[float] = [1.0, 2.0, 4.0, 8.0, 16.0, 30.0]
# Random jitter fraction applied to each backoff delay (±20%)
_RECONNECT_JITTER: float = 0.2
# How often to poll the connection health (seconds)
_HEALTH_POLL_INTERVAL: float = 5.0


class RelayTunnel:
    """Manage a connection to the LLMHosts relay server.

    Interface: ``start()``, ``stop()``, ``status()``, ``get_url()``.

    After a successful connection, a background task monitors the relay
    and reconnects automatically with exponential backoff + jitter.
    """

    def __init__(self) -> None:
        self._client: object | None = None  # _relay_core.RelayClient instance
        self._active: bool = False
        self._url: str | None = None
        self._started_at: datetime | None = None
        self._port: int = 4000
        self._relay_url: str = DEFAULT_RELAY_URL
        self._device_token: str | None = None
        self._reconnect_task: asyncio.Task | None = None  # type: ignore[type-arg]
        self._reconnect_count: int = 0
        self._stopping: bool = False

    async def start(
        self,
        port: int = 4000,
        relay_url: str | None = None,
        device_token: str | None = None,
    ) -> TunnelStatus:
        """Connect to the relay server.

        Args:
            port: Local port that the FastAPI proxy listens on.
            relay_url: WebSocket URL of the relay server.
            device_token: Authentication token for this device.
        """
        self._port = port
        self._relay_url = relay_url or DEFAULT_RELAY_URL
        self._device_token = device_token
        self._stopping = False

        if not RELAY_AVAILABLE or _RelayClient is None:
            return TunnelStatus(
                provider=TunnelProvider.RELAY,
                active=False,
                local_port=port,
                error="Rust relay core not available (install with maturin develop)",
            )

        if not device_token:
            return TunnelStatus(
                provider=TunnelProvider.RELAY,
                active=False,
                local_port=port,
                error="No relay device token configured",
            )

        status = await self._connect_once()

        if status.active:
            # Launch background reconnect monitor
            if self._reconnect_task is not None and not self._reconnect_task.done():
                self._reconnect_task.cancel()
            self._reconnect_task = asyncio.create_task(self._reconnect_loop())

        return status

    async def _connect_once(self) -> TunnelStatus:
        """Attempt a single connection to the relay server."""
        assert _RelayClient is not None  # guarded by caller

        try:
            # M2: pass server_pubkey for Noise NK encryption if configured.
            server_pubkey = os.environ.get("RELAY_SERVER_PUBKEY") or None
            self._client = _RelayClient(self._relay_url, self._device_token, self._port, server_pubkey=server_pubkey)
            await self._client.connect()  # type: ignore[union-attr]

            if self._client.is_connected:  # type: ignore[union-attr]
                self._active = True
                self._started_at = datetime.now(tz=timezone.utc)
                self._url = self._relay_url

                encrypted = getattr(self._client, "is_noise_encrypted", False)
                logger.info(
                    "Relay tunnel connected: %s (port %d, version %s, encrypted=%s)",
                    self._relay_url,
                    self._port,
                    RELAY_VERSION,
                    encrypted,
                )

                return TunnelStatus(
                    provider=TunnelProvider.RELAY,
                    active=True,
                    url=self._url,
                    local_port=self._port,
                    started_at=self._started_at,
                )

            return TunnelStatus(
                provider=TunnelProvider.RELAY,
                active=False,
                local_port=self._port,
                error="Relay connection established but client reports disconnected",
            )

        except Exception as exc:
            logger.error("Relay connection failed: %s", exc)
            self._active = False
            self._client = None
            return TunnelStatus(
                provider=TunnelProvider.RELAY,
                active=False,
                local_port=self._port,
                error=f"Relay connection failed: {exc}",
            )

    async def _reconnect_loop(self) -> None:
        """Background task: monitor relay health and reconnect on disconnect.

        Uses exponential backoff (1s → 30s cap) with ±20% random jitter to
        prevent thundering herd when the relay server restarts.
        """
        attempt = 0

        while not self._stopping:
            # Wait for the next health poll tick
            try:
                await asyncio.sleep(_HEALTH_POLL_INTERVAL)
            except asyncio.CancelledError:
                return

            if self._stopping:
                return

            # Check if still connected
            try:
                connected = self._client is not None and bool(self._client.is_connected)  # type: ignore[attr-defined]
            except Exception:
                connected = False

            if connected:
                attempt = 0  # Stable — reset backoff counter
                continue

            # Connection lost
            logger.warning(
                "Relay connection lost — scheduling reconnect (attempt %d)",
                attempt + 1,
            )
            self._active = False

            # Compute backoff with jitter
            backoff = _RECONNECT_BACKOFF[min(attempt, len(_RECONNECT_BACKOFF) - 1)]
            jitter = backoff * _RECONNECT_JITTER * (random.random() * 2.0 - 1.0)
            delay = max(0.1, backoff + jitter)

            logger.info("Relay reconnect in %.1fs (attempt %d)...", delay, attempt + 1)

            try:
                await asyncio.sleep(delay)
            except asyncio.CancelledError:
                return

            if self._stopping:
                return

            if _RelayClient is None:
                attempt += 1
                continue

            try:
                server_pubkey = os.environ.get("RELAY_SERVER_PUBKEY") or None
                self._client = _RelayClient(
                    self._relay_url,
                    self._device_token,
                    self._port,
                    server_pubkey=server_pubkey,
                )
                await self._client.connect()  # type: ignore[union-attr]

                if self._client.is_connected:  # type: ignore[union-attr]
                    self._active = True
                    self._reconnect_count += 1
                    encrypted = getattr(self._client, "is_noise_encrypted", False)
                    logger.info(
                        "Relay reconnected after %d attempt(s) (encrypted=%s)",
                        attempt + 1,
                        encrypted,
                    )
                    attempt = 0  # Reset backoff on successful reconnect
                else:
                    logger.warning("Relay reconnect attempt %d: not connected after connect()", attempt + 1)
                    self._client = None
                    attempt += 1

            except Exception as exc:
                logger.error("Relay reconnect attempt %d failed: %s", attempt + 1, exc)
                self._client = None
                attempt += 1

    async def stop(self) -> None:
        """Disconnect from the relay server and cancel the reconnect monitor."""
        self._stopping = True

        # Cancel the background reconnect loop
        if self._reconnect_task is not None and not self._reconnect_task.done():
            self._reconnect_task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._reconnect_task
        self._reconnect_task = None

        if self._client is not None:
            try:
                await self._client.disconnect()  # type: ignore[attr-defined]
            except Exception as exc:
                logger.warning("Error disconnecting relay: %s", exc)

        self._active = False
        self._client = None
        self._url = None
        self._started_at = None
        logger.info("Relay tunnel stopped")

    async def status(self) -> TunnelStatus:
        """Check relay connection status."""
        if self._client is not None and self._active:
            try:
                connected = self._client.is_connected  # type: ignore[attr-defined]
            except Exception:
                connected = False

            if connected:
                return TunnelStatus(
                    provider=TunnelProvider.RELAY,
                    active=True,
                    url=self._url,
                    local_port=self._port,
                    started_at=self._started_at,
                )

            # Client reports disconnected — reconnect loop will handle it
            self._active = False
            return TunnelStatus(
                provider=TunnelProvider.RELAY,
                active=False,
                local_port=self._port,
                error="Relay connection lost",
            )

        return TunnelStatus(
            provider=TunnelProvider.RELAY,
            active=False,
            local_port=self._port,
        )

    async def get_url(self) -> str | None:
        """Get the relay URL."""
        return self._url

    @property
    def bytes_relayed(self) -> int:
        """Total bytes relayed through the connection."""
        if self._client is not None:
            try:
                return int(self._client.bytes_relayed)  # type: ignore[attr-defined]
            except Exception:
                return 0
        return 0

    @property
    def uptime_seconds(self) -> float:
        """Connection uptime in seconds."""
        if self._client is not None:
            try:
                return float(self._client.uptime_seconds)  # type: ignore[attr-defined]
            except Exception:
                return 0.0
        return 0.0

    @property
    def reconnect_count(self) -> int:
        """Number of successful reconnects since initial connection."""
        return self._reconnect_count
