"""上流クライアントとモデル一覧キャッシュ。

httpx は既定で ``Accept`` ``Accept-Encoding`` ``Connection`` ``User-Agent`` を補う。
無改変中継ではこれらの自動付与が差分になるため、クライアントが送っていない
ヘッダは組み立て後に取り除く。応答は ``aiter_raw()`` で受け、自動解凍しない。
"""

from __future__ import annotations

import asyncio
import json
import logging
from dataclasses import dataclass
from typing import Sequence

import httpx

from .config import Settings
from .db import iso_jst, now_jst
from .logging_utils import scrub

logger = logging.getLogger(__name__)

#: httpx が自動で足すヘッダ。クライアント由来でなければ落とす。
_HTTPX_DEFAULT_HEADERS: frozenset[str] = frozenset({"accept", "accept-encoding", "connection", "user-agent"})

UPSTREAM_OK = "ok"
UPSTREAM_UNREACHABLE = "unreachable"
UPSTREAM_UNKNOWN = "unknown"


@dataclass(slots=True)
class ModelsSnapshot:
    """``/v1/models`` の取得結果。"""

    models: tuple[str, ...]
    fetched_at: str | None
    stale: bool


class UpstreamClient:
    """上流 llama-server への接続を束ねる。"""

    def __init__(self, settings: Settings, *, client: httpx.AsyncClient | None = None) -> None:
        self._settings = settings
        self._base = httpx.URL(settings.upstream_base_url)
        self._base_raw_path = self._base.raw_path.rstrip(b"/")
        self._owns_client = client is None
        self._client = client or httpx.AsyncClient(
            timeout=httpx.Timeout(
                connect=settings.connect_timeout,
                read=settings.idle_timeout,
                write=settings.idle_timeout,
                pool=settings.connect_timeout,
            ),
            limits=httpx.Limits(max_connections=None, max_keepalive_connections=64),
            follow_redirects=False,
        )
        self._models_lock = asyncio.Lock()
        self._models: ModelsSnapshot = ModelsSnapshot(models=(), fetched_at=None, stale=False)
        self._models_fetched_monotonic: float | None = None
        self._upstream_status = UPSTREAM_UNKNOWN

    # ------------------------------------------------------------------ 基本

    @property
    def client(self) -> httpx.AsyncClient:
        return self._client

    @property
    def upstream_status(self) -> str:
        """直近の上流到達性。``ok`` / ``unreachable`` / ``unknown``。"""

        return self._upstream_status

    async def aclose(self) -> None:
        if self._owns_client:
            await self._client.aclose()

    def target_url(self, raw_path: bytes, query_string: bytes) -> httpx.URL:
        """受信した生のパスをそのまま上流 URL へ組み立てる。"""

        path = raw_path if raw_path.startswith(b"/") else b"/" + raw_path
        full = self._base_raw_path + path
        if query_string:
            full = full + b"?" + query_string
        return self._base.copy_with(raw_path=full)

    def build_request(
        self,
        *,
        method: str,
        raw_path: bytes,
        query_string: bytes,
        headers: Sequence[tuple[bytes, bytes]],
        content: bytes,
    ) -> httpx.Request:
        """上流向けリクエストを組み立てる。

        ``content`` は受信したバイト列そのもの。再シリアライズはしない。
        """

        request = self._client.build_request(
            method,
            self.target_url(raw_path, query_string),
            headers=httpx.Headers(list(headers)),
            content=content,
        )
        sent = {name.decode("latin-1").lower() for name, _ in headers}
        for name in _HTTPX_DEFAULT_HEADERS:
            if name not in sent and name in request.headers:
                del request.headers[name]
        return request

    async def send(self, request: httpx.Request) -> httpx.Response:
        """ストリーミングで送出する。応答本体は ``aiter_raw()`` で読む。"""

        response = await self._client.send(request, stream=True)
        self._upstream_status = UPSTREAM_OK
        return response

    def mark_unreachable(self) -> None:
        self._upstream_status = UPSTREAM_UNREACHABLE

    # ------------------------------------------------------------ モデル一覧

    def cached_models(self) -> ModelsSnapshot:
        return self._models

    async def get_models(self, *, force: bool = False) -> ModelsSnapshot:
        """上流のモデル一覧をキャッシュ付きで返す。

        到達不能時は直近のキャッシュを ``stale: true`` で返す。
        """

        loop_now = asyncio.get_running_loop().time()
        if (
            not force
            and self._models_fetched_monotonic is not None
            and (loop_now - self._models_fetched_monotonic) < self._settings.models_cache_ttl
        ):
            return self._models

        async with self._models_lock:
            loop_now = asyncio.get_running_loop().time()
            if (
                not force
                and self._models_fetched_monotonic is not None
                and (loop_now - self._models_fetched_monotonic) < self._settings.models_cache_ttl
            ):
                return self._models
            try:
                response = await self._client.get(
                    self.target_url(b"/v1/models", b""),
                    headers=self._internal_headers(),
                    timeout=httpx.Timeout(self._settings.connect_timeout),
                )
                response.raise_for_status()
                models = _parse_models(response.content)
            except Exception as exc:
                self._upstream_status = UPSTREAM_UNREACHABLE
                logger.warning("上流のモデル一覧を取得できませんでした: %s", scrub(str(exc)))
                self._models = ModelsSnapshot(
                    models=self._models.models,
                    fetched_at=self._models.fetched_at,
                    stale=True,
                )
                return self._models
            self._upstream_status = UPSTREAM_OK
            self._models_fetched_monotonic = asyncio.get_running_loop().time()
            self._models = ModelsSnapshot(models=models, fetched_at=iso_jst(now_jst()), stale=False)
            return self._models

    def _internal_headers(self) -> dict[str, str]:
        headers = {"accept": "application/json"}
        if self._settings.upstream_api_key:
            headers["authorization"] = f"Bearer {self._settings.upstream_api_key}"
        return headers

    async def probe(self) -> str:
        """ヘルスチェック用に上流到達性を確かめる。"""

        snapshot = await self.get_models()
        if snapshot.fetched_at is None and snapshot.stale:
            return UPSTREAM_UNREACHABLE
        return UPSTREAM_UNREACHABLE if snapshot.stale else UPSTREAM_OK


def _parse_models(raw: bytes) -> tuple[str, ...]:
    try:
        payload = json.loads(raw)
    except (ValueError, UnicodeDecodeError):
        return ()
    if not isinstance(payload, dict):
        return ()
    data = payload.get("data")
    if not isinstance(data, list):
        return ()
    models: list[str] = []
    for item in data:
        if isinstance(item, dict):
            identifier = item.get("id")
            if isinstance(identifier, str) and identifier:
                models.append(identifier)
    return tuple(dict.fromkeys(models))