"""
Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>

AISBF - AI Service Broker Framework || AI Should Be Free

Process-wide runtime for the "shodan" last-resort provider.

This module owns a single shared discovery/failover engine (vendored from
Pasquale Minervini's open-router project, see open_router_engine.py) and adapts
it to AISBF's provider contract. It builds the engine's DiscoveryConfig /
ProbeConfig from an AISBF provider's ``shodan_config`` dict, runs the background
discovery + health loops once per process, and exposes a small OpenAI-compatible
failover routing helper that returns AISBF-friendly types (a plain dict for
non-streaming, an async generator of OpenAI SSE bytes for streaming).

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""
import asyncio
import json
import logging
import time
from http import HTTPStatus
from pathlib import Path
from typing import Any, AsyncIterator, Dict, List, Optional, Sequence

import aiohttp

from . import open_router_engine as engine

logger = logging.getLogger(__name__)

# Model names that mean "route to whatever the discovery pool currently offers"
# rather than a specific upstream model id (which for discovered Ollama/vLLM
# hosts is unpredictable and changes over time).
AUTO_MODEL_NAMES = frozenset({"", "auto", "any", "*", "shodan/auto", "shodan"})


class ShodanRuntime:
    """Owns one engine.RouterRuntime + shared upstream session for the process."""

    def __init__(self, router_config: "engine.RouterConfig", upstream_timeout: float) -> None:
        self._router_config = router_config
        self._upstream_timeout = upstream_timeout
        self._session: Optional[aiohttp.ClientSession] = None
        self._runtime: Optional["engine.RouterRuntime"] = None
        self._started = False

    async def _ensure_started(self) -> None:
        if self._started:
            return
        # The credential-free upstream client. Discovered hosts are untrusted, so
        # we never attach auth headers (engine._upstream_headers already omits
        # them); this session only carries content negotiation headers per call.
        self._session = aiohttp.ClientSession()
        self._runtime = engine.RouterRuntime(self._router_config, logger, self._session)
        self._runtime.start_background()
        self._started = True
        logger.info(
            "ShodanRuntime started: discovery mode=%s, masscan lists=%d, "
            "discovery interval=%.0fs, health interval=%.0fs",
            self._router_config.discovery.shodan_mode,
            len(self._router_config.discovery.masscan_list_paths),
            self._router_config.discovery.interval,
            self._router_config.probe.health_interval,
        )

    @property
    def registry(self) -> "engine.EndpointRegistry":
        assert self._runtime is not None
        return self._runtime.registry

    async def available_models(self) -> List[str]:
        await self._ensure_started()
        return self.registry.models()

    async def endpoint_count(self) -> int:
        await self._ensure_started()
        return self.registry.endpoint_count()

    def _resolve_model(self, requested: str) -> Optional[str]:
        """Map a requested model to a concrete model id the pool can serve.

        For an explicit model name we honour it as-is (candidates() will match).
        For an ``auto`` alias we pick the fastest-servable model currently in the
        registry: prefer a non-"cloud" model with at least one live candidate.
        Returns None when the pool is empty.
        """
        requested_norm = (requested or "").strip()
        if requested_norm.lower() not in AUTO_MODEL_NAMES:
            return requested_norm
        models = self.registry.models()
        if not models:
            return None
        # Prefer a model that actually has a live candidate right now, and prefer
        # local (non-"cloud") models so we don't lean on remote paid services.
        best: Optional[str] = None
        for model in models:
            if not self.registry.candidates(model):
                continue
            if "cloud" not in model.lower():
                return model
            best = best or model
        return best or models[0]

    async def route(
        self,
        resource: str,
        body: Dict[str, Any],
        stream: bool,
    ) -> Any:
        """Try the discovered upstreams in priority order until one answers.

        Returns a plain OpenAI-compatible dict for non-streaming requests, or an
        async generator yielding OpenAI-compatible SSE byte chunks for streaming.
        Raises on total failure (no endpoint served the request) so the caller
        can record a provider failure and, in a rotation, fail over.
        """
        await self._ensure_started()

        requested_model = body.get("model", "")
        model = self._resolve_model(requested_model)
        if model is None:
            raise RuntimeError(
                "shodan: no discovered endpoints are available yet "
                "(discovery may still be warming up)"
            )
        # Send the resolved concrete model id upstream.
        body = {**body, "model": model}
        payload = json.dumps(body).encode("utf-8")

        candidates = self.registry.candidates(model)
        if not candidates:
            raise RuntimeError(
                f"shodan: no live upstream advertises model {model!r}"
            )

        t = self._upstream_timeout
        timeout = aiohttp.ClientTimeout(total=None, sock_connect=t, sock_read=t)
        headers = {
            "content-type": "application/json",
            "accept": "text/event-stream" if stream else "application/json",
        }

        last_error: Optional[str] = None
        for state in candidates:
            endpoint = state.config
            url = endpoint.api_url(resource)
            started = time.monotonic()
            try:
                response = await self._session.post(
                    url,
                    data=payload,
                    headers=headers,
                    timeout=timeout,
                    allow_redirects=False,
                )
            except (aiohttp.ClientError, asyncio.TimeoutError, ValueError) as error:
                last_error = f"{endpoint.base_url}: request failed: {error}"
                self.registry.discard(endpoint.base_url)
                logger.warning("shodan upstream %s discarded: %s", endpoint.base_url, error)
                continue

            if self._is_failure_status(response.status):
                error_body = await self._safe_read(response)
                response.close()
                last_error = f"{endpoint.base_url}: HTTP {response.status}"
                self._quarantine_on_status(endpoint, model, response.status)
                logger.warning(
                    "shodan upstream %s returned HTTP %s for model %s; trying next",
                    endpoint.base_url, response.status, model,
                )
                continue

            # Success — this endpoint served the request.
            if stream:
                logger.info("shodan routing stream to %s (model %s)", endpoint.base_url, model)
                return self._relay_stream(response, endpoint.base_url, model, started)

            data = await self._safe_read(response)
            response.close()
            self.registry.record_generation_latency(
                endpoint.base_url, model, time.monotonic() - started
            )
            logger.info("shodan served non-stream from %s (model %s)", endpoint.base_url, model)
            try:
                return json.loads(data)
            except json.JSONDecodeError as error:
                last_error = f"{endpoint.base_url}: invalid JSON ({error})"
                self.registry.discard(endpoint.base_url)
                continue

        raise RuntimeError(
            f"shodan: all {len(candidates)} candidate upstream(s) failed for "
            f"model {model!r} (last error: {last_error})"
        )

    async def _relay_stream(
        self,
        response: aiohttp.ClientResponse,
        base_url: str,
        model: str,
        started: float,
    ) -> AsyncIterator[bytes]:
        """Relay an upstream SSE body chunk-by-chunk as OpenAI-compatible bytes.

        A failure after the first bytes cannot be retried on another endpoint
        (the client stream has already begun), matching the upstream engine's
        behaviour: we quarantine the endpoint and let the error surface.
        """
        try:
            async for chunk in response.content.iter_any():
                if chunk:
                    yield chunk
            self.registry.record_generation_latency(
                base_url, model, time.monotonic() - started
            )
        except (aiohttp.ClientError, asyncio.TimeoutError) as error:
            self.registry.discard(base_url)
            logger.warning("shodan stream from %s failed mid-flight: %s", base_url, error)
            raise
        finally:
            response.close()

    @staticmethod
    def _is_failure_status(status: int) -> bool:
        return (
            engine._is_upstream_failure_status(status)
            or status in engine.ENDPOINT_FAILURE_STATUSES
            or status in engine.MODEL_FAILURE_STATUSES
            or status >= 500
        )

    def _quarantine_on_status(self, endpoint: "engine.EndpointConfig", model: str, status: int) -> None:
        # 401 means the endpoint demands credentials we deliberately never send:
        # drop the whole endpoint. 403/410 blacklist just this endpoint/model
        # pairing. Everything else drops the endpoint for this pass.
        if status in engine.MODEL_FAILURE_STATUSES:
            self.registry.discard_model(endpoint.base_url, model)
        else:
            self.registry.discard(endpoint.base_url)

    @staticmethod
    async def _safe_read(response: aiohttp.ClientResponse) -> bytes:
        try:
            return await response.content.read(engine.MAX_MODELS_RESPONSE_BYTES + 1)
        except (aiohttp.ClientError, asyncio.TimeoutError):
            return b""

    async def stop(self) -> None:
        if self._runtime is not None:
            await self._runtime.stop()
        if self._session is not None:
            await self._session.close()
        self._started = False


# --- Process-wide singleton -------------------------------------------------

_runtime_singleton: Optional[ShodanRuntime] = None
_runtime_lock = asyncio.Lock()


def _build_router_config(shodan_config: Optional[Dict[str, Any]]) -> "engine.RouterConfig":
    cfg = dict(shodan_config or {})

    masscan_paths = tuple(
        Path(p).expanduser() for p in (cfg.get("masscan_lists") or []) if str(p).strip()
    )
    discovery = engine.DiscoveryConfig(
        shodan_api_key=cfg.get("shodan_api_key") or None,
        shodan_mode=str(cfg.get("shodan_mode", "auto")),
        retrieval_timeout=float(cfg.get("retrieval_timeout", engine.DEFAULT_RETRIEVAL_TIMEOUT)),
        retrieval_retries=int(cfg.get("retrieval_retries", engine.DEFAULT_RETRIEVAL_RETRIES)),
        shodan_max_results=cfg.get("shodan_max_results"),
        masscan_list_paths=masscan_paths,
        interval=float(cfg.get("discovery_interval", engine.DEFAULT_DISCOVERY_INTERVAL)),
        blacklist_duration=float(cfg.get("blacklist_duration", engine.DEFAULT_BLACKLIST_DURATION)),
    )
    probe = engine.ProbeConfig(
        workers=int(cfg.get("workers", engine.DEFAULT_WORKERS)),
        timeout=float(cfg.get("probe_timeout", engine.DEFAULT_PROBE_TIMEOUT)),
        readiness_timeout=float(cfg.get("readiness_timeout", engine.DEFAULT_READINESS_TIMEOUT)),
        readiness_model_attempts=int(
            cfg.get("readiness_model_attempts", engine.DEFAULT_READINESS_MODEL_ATTEMPTS)
        ),
        health_interval=float(cfg.get("health_interval", engine.DEFAULT_HEALTH_INTERVAL)),
        progress_every=int(cfg.get("progress_every", engine.DEFAULT_PROGRESS_EVERY)),
    )
    server = engine.ServerOptions(
        upstream_timeout=float(cfg.get("upstream_timeout", engine.DEFAULT_UPSTREAM_TIMEOUT)),
    )
    return engine.RouterConfig(discovery=discovery, probe=probe, server=server)


async def get_runtime(shodan_config: Optional[Dict[str, Any]]) -> ShodanRuntime:
    """Return the process-wide ShodanRuntime, creating it on first use.

    The first caller's ``shodan_config`` configures discovery for the whole
    process; later callers reuse the same runtime (a single discovery pool is
    the intent — one broker, one shared set of discovered endpoints).
    """
    global _runtime_singleton
    if _runtime_singleton is not None:
        return _runtime_singleton
    async with _runtime_lock:
        if _runtime_singleton is None:
            router_config = _build_router_config(shodan_config)
            _runtime_singleton = ShodanRuntime(
                router_config, router_config.server.upstream_timeout
            )
        return _runtime_singleton
