"""Strict client for the local LM Studio server.

The client uses LM Studio's OpenAI-compatible chat-completions endpoint because
that endpoint supports custom tools. Model discovery is performed before
inference, and model identity mismatches are fatal.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass
import json
import os
import re
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

from .specs import ModelSpec


class LMStudioError(RuntimeError):
    """Raised when discovery, identity validation, or inference fails."""


class LMStudioTransportError(LMStudioError):
    """Raised when no valid server response was observed and one retry is safe."""


def normalize_identity(value: str) -> str:
    return re.sub(r"[^a-z0-9]+", "", value.lower())


def _record_identity(record: dict[str, Any]) -> str:
    values = [
        str(record.get("id", "")),
        str(record.get("key", "")),
        str(record.get("display_name", "")),
        str(record.get("name", "")),
    ]
    return " ".join(item for item in values if item)


@dataclass(frozen=True, slots=True)
class DiscoveryResult:
    openai_models: tuple[dict[str, Any], ...]
    native_models: tuple[dict[str, Any], ...]
    endpoint_errors: tuple[str, ...]

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True, slots=True)
class ResolvedModel:
    inference_key: str
    openai_record: dict[str, Any]
    native_record: dict[str, Any] | None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


def select_expected_model(
    spec: ModelSpec,
    openai_models: tuple[dict[str, Any], ...] | list[dict[str, Any]],
    native_models: tuple[dict[str, Any], ...] | list[dict[str, Any]] = (),
) -> ResolvedModel:
    """Resolve one and only one inference-visible model matching the fixed identity."""

    expected = normalize_identity(spec.expected_identity)
    matches = [
        record
        for record in openai_models
        if expected in normalize_identity(_record_identity(record))
    ]
    if not matches:
        visible = sorted(filter(None, (_record_identity(item) for item in openai_models)))
        raise LMStudioError(
            f"Expected {spec.canonical_name}, but no matching model is visible through "
            f"{spec.discovery_endpoint}. Visible models: {visible or ['<none>']}"
        )
    if len(matches) > 1:
        exact = [
            record
            for record in matches
            if normalize_identity(str(record.get("id", ""))) == expected
            or normalize_identity(str(record.get("key", ""))) == expected
        ]
        if len(exact) == 1:
            matches = exact
        else:
            raise LMStudioError(
                "Model identity is ambiguous; refusing to select a quantization or variant silently: "
                + ", ".join(_record_identity(item) for item in matches)
            )

    record = matches[0]
    inference_key = str(record.get("id") or record.get("key") or "")
    if not inference_key:
        raise LMStudioError("Matching LM Studio model record has no inference identifier")
    if inference_key != spec.expected_inference_key:
        raise LMStudioError(
            f"Expected inference key {spec.expected_inference_key}, but LM Studio exposed {inference_key}"
        )

    native_match: dict[str, Any] | None = None
    native_candidates = [
        item for item in native_models if expected in normalize_identity(_record_identity(item))
    ]
    if len(native_candidates) == 1:
        native_match = native_candidates[0]
    elif native_candidates:
        selected = [item for item in native_candidates if item.get("selected_variant")]
        if len(selected) == 1:
            native_match = selected[0]

    if native_match is None:
        raise LMStudioError(
            "The matching model has no unique native /api/v1/models record; "
            "variant and runtime metadata cannot be verified"
        )

    quantization = native_match.get("quantization", {})
    quantization_name = quantization.get("name") if isinstance(quantization, dict) else quantization
    expected_runtime = {
        "selected_variant": spec.expected_variant,
        "format": spec.expected_format,
        "quantization": spec.expected_quantization,
    }
    actual_runtime = {
        "selected_variant": native_match.get("selected_variant"),
        "format": native_match.get("format"),
        "quantization": quantization_name,
    }
    mismatches = [
        f"{field}: expected {expected_runtime[field]!r}, observed {actual_runtime[field]!r}"
        for field in expected_runtime
        if expected_runtime[field] != actual_runtime[field]
    ]

    loaded_instances = native_match.get("loaded_instances", [])
    loaded_contexts = {
        item.get("config", {}).get("context_length")
        for item in loaded_instances
        if isinstance(item, dict) and isinstance(item.get("config"), dict)
    }
    if loaded_contexts != {spec.context_length}:
        observed_contexts = sorted(loaded_contexts, key=lambda value: str(value))
        mismatches.append(
            f"loaded context length: expected only {spec.context_length}, observed {observed_contexts}"
        )

    capabilities = native_match.get("capabilities", {})
    reasoning = capabilities.get("reasoning", {}) if isinstance(capabilities, dict) else {}
    observed_reasoning = reasoning.get("default") if isinstance(reasoning, dict) else None
    if observed_reasoning is None:
        observed_reasoning = "none"
    if observed_reasoning != spec.reasoning_mode:
        mismatches.append(
            f"reasoning mode: expected {spec.reasoning_mode!r}, observed {observed_reasoning!r}"
        )
    if mismatches:
        raise LMStudioError(
            f"LM Studio runtime does not match {spec.model_id}: " + "; ".join(mismatches)
        )

    return ResolvedModel(
        inference_key=inference_key,
        openai_record=dict(record),
        native_record=dict(native_match),
    )


class LMStudioClient:
    def __init__(self, spec: ModelSpec, timeout_seconds: float = 10.0):
        self.spec = spec
        self.timeout_seconds = timeout_seconds

    def _headers(self) -> dict[str, str]:
        headers = {"Content-Type": "application/json"}
        token = os.environ.get(self.spec.api_token_env, "").strip()
        if token:
            headers["Authorization"] = f"Bearer {token}"
        return headers

    def _request(
        self,
        method: str,
        endpoint: str,
        payload: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        data = None if payload is None else json.dumps(payload).encode("utf-8")
        request = Request(
            self.spec.base_url + endpoint,
            data=data,
            method=method,
            headers=self._headers(),
        )
        try:
            with urlopen(request, timeout=self.timeout_seconds) as response:
                body = response.read().decode("utf-8")
        except HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")
            error_type = (
                LMStudioTransportError
                if exc.code in {408, 429, 500, 502, 503, 504}
                else LMStudioError
            )
            raise error_type(
                f"LM Studio returned HTTP {exc.code} for {endpoint}: {detail}"
            ) from exc
        except URLError as exc:
            raise LMStudioTransportError(
                f"Cannot connect to LM Studio at {self.spec.base_url}. "
                "Start the server on port 1234 and load the frozen model. "
                f"Underlying error: {exc.reason}"
            ) from exc
        try:
            decoded = json.loads(body)
        except json.JSONDecodeError as exc:
            raise LMStudioError(f"LM Studio returned non-JSON data for {endpoint}") from exc
        if not isinstance(decoded, dict):
            raise LMStudioError(f"LM Studio returned an unexpected response for {endpoint}")
        return decoded

    def discover(self) -> DiscoveryResult:
        errors: list[str] = []
        openai_models: tuple[dict[str, Any], ...] = ()
        native_models: tuple[dict[str, Any], ...] = ()
        try:
            response = self._request("GET", self.spec.discovery_endpoint)
            data = response.get("data", [])
            if isinstance(data, list):
                openai_models = tuple(item for item in data if isinstance(item, dict))
        except LMStudioError as exc:
            errors.append(str(exc))
        try:
            response = self._request("GET", self.spec.native_discovery_endpoint)
            data = response.get("models", [])
            if isinstance(data, list):
                native_models = tuple(item for item in data if isinstance(item, dict))
        except LMStudioError as exc:
            errors.append(str(exc))
        if not openai_models and not native_models:
            raise LMStudioError("; ".join(errors) or "LM Studio returned no model records")
        return DiscoveryResult(openai_models, native_models, tuple(errors))

    def resolve(self, discovery: DiscoveryResult | None = None) -> tuple[DiscoveryResult, ResolvedModel]:
        result = discovery or self.discover()
        resolved = select_expected_model(self.spec, result.openai_models, result.native_models)
        return result, resolved

    def chat_completions(
        self,
        model_key: str,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        max_tokens: int | None = None,
        seed: int | None = None,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "model": model_key,
            "messages": messages,
            "temperature": self.spec.temperature,
            "top_p": self.spec.top_p,
            "max_tokens": max_tokens or self.spec.max_tokens,
            "seed": self.spec.seed if seed is None else seed,
            "stream": False,
        }
        if tools is not None:
            payload["tools"] = tools
        return self._request("POST", self.spec.inference_endpoint, payload)

    def inference_probe(self, model_key: str) -> dict[str, Any]:
        response = self.chat_completions(
            model_key,
            messages=[
                {
                    "role": "user",
                    "content": "Reply with exactly MODEL_OK and no other text.",
                }
            ],
            max_tokens=256,
        )
        try:
            message = response["choices"][0]["message"]
            content = message["content"]
        except (KeyError, IndexError, TypeError) as exc:
            raise LMStudioError("Inference probe returned an invalid chat-completions response") from exc
        if not isinstance(content, str) or content.strip() != "MODEL_OK":
            finish_reason = response.get("choices", [{}])[0].get("finish_reason")
            reasoning_content = message.get("reasoning_content", "")
            raise LMStudioError(
                "Inference probe did not produce the required MODEL_OK marker "
                f"(finish_reason={finish_reason!r}, visible={content!r}, "
                f"reasoning_chars={len(reasoning_content)})"
            )
        return response