from __future__ import annotations

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


@dataclass
class ModelResponse:
    ok: bool
    content: str = ""
    parsed: dict[str, Any] | None = None
    duration_seconds: float = 0.0
    error: str = ""


def _chat_completions_url(base_url: str) -> str:
    base = base_url.rstrip("/")
    if base.endswith("/chat/completions"):
        return base
    if base.endswith("/v1"):
        return f"{base}/chat/completions"
    return f"{base}/v1/chat/completions"


def parse_json_object(text: str) -> dict[str, Any] | None:
    raw = text.strip()
    fence = re.search(r"```(?:json)?\s*(.*?)```", raw, re.IGNORECASE | re.DOTALL)
    if fence:
        raw = fence.group(1).strip()
    start = raw.find("{")
    end = raw.rfind("}")
    if start >= 0 and end > start:
        raw = raw[start : end + 1]
    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError:
        return None
    return parsed if isinstance(parsed, dict) else None


def _network_error_message(label: str, exc: BaseException, timeout: float) -> str:
    reason = getattr(exc, "reason", exc)
    if isinstance(exc, TimeoutError) or isinstance(reason, TimeoutError):
        detail = f"timed out after {timeout:g}s"
    else:
        detail = str(reason) or type(reason).__name__
    return f"{label} -> {detail}"


class ModelClient:
    def __init__(self, model_config: dict[str, Any], timeout_seconds: float = 30):
        self.model_config = model_config
        self.timeout_seconds = timeout_seconds

    @property
    def available(self) -> bool:
        return bool(
            self.model_config.get("controller_model_proxy_url")
            or (self.model_config.get("vllm_base_url") and self.model_config.get("vllm_model"))
        )

    def chat_json(self, messages: list[dict[str, str]], temperature: float | None = None) -> ModelResponse:
        if not self.available:
            return ModelResponse(ok=False, error="Model config is incomplete.")

        proxy_url = str(self.model_config.get("controller_model_proxy_url") or "").strip()
        if proxy_url:
            return self._chat_json_via_controller(proxy_url, messages, temperature)

        url = _chat_completions_url(str(self.model_config.get("vllm_base_url") or ""))
        payload = {
            "model": self.model_config.get("vllm_model"),
            "messages": messages,
            "temperature": temperature if temperature is not None else self.model_config.get("temperature", 0),
        }
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json; charset=utf-8",
        }
        api_key = str(self.model_config.get("api_key") or "")
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"

        started = time.monotonic()
        request = Request(url, data=body, headers=headers, method="POST")
        try:
            with urlopen(request, timeout=self.timeout_seconds) as response:
                raw = response.read()
        except HTTPError as exc:
            error_body = exc.read().decode("utf-8", errors="replace")
            return ModelResponse(
                ok=False,
                duration_seconds=round(time.monotonic() - started, 3),
                error=f"{url} -> HTTP {exc.code}: {error_body or exc.reason}",
            )
        except (URLError, TimeoutError, OSError) as exc:
            return ModelResponse(
                ok=False,
                duration_seconds=round(time.monotonic() - started, 3),
                error=_network_error_message(url, exc, self.timeout_seconds),
            )

        duration = round(time.monotonic() - started, 3)
        try:
            data = json.loads(raw.decode("utf-8"))
            choice = data["choices"][0]
            message = choice.get("message") or {}
            content = (
                message.get("content")
                or message.get("reasoning_content")
                or message.get("reasoning")
                or choice.get("text")
                or ""
            )
        except (KeyError, IndexError, TypeError, json.JSONDecodeError) as exc:
            return ModelResponse(ok=False, duration_seconds=duration, error=f"Invalid model response: {exc}")
        if not isinstance(content, str):
            content = json.dumps(content, ensure_ascii=False)

        return ModelResponse(
            ok=True,
            content=str(content),
            parsed=parse_json_object(str(content)),
            duration_seconds=duration,
        )

    def _chat_json_via_controller(
        self,
        url: str,
        messages: list[dict[str, str]],
        temperature: float | None = None,
    ) -> ModelResponse:
        payload = {
            "messages": messages,
            "temperature": temperature if temperature is not None else self.model_config.get("temperature", 0),
        }
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json; charset=utf-8",
        }
        token = str(self.model_config.get("controller_model_proxy_token") or "")
        if token:
            headers["Authorization"] = f"Bearer {token}"

        started = time.monotonic()
        request = Request(url, data=body, headers=headers, method="POST")
        try:
            with urlopen(request, timeout=self.timeout_seconds) as response:
                raw = response.read()
        except HTTPError as exc:
            error_body = exc.read().decode("utf-8", errors="replace")
            return ModelResponse(
                ok=False,
                duration_seconds=round(time.monotonic() - started, 3),
                error=f"controller model proxy -> HTTP {exc.code}: {error_body or exc.reason}",
            )
        except (URLError, TimeoutError, OSError) as exc:
            return ModelResponse(
                ok=False,
                duration_seconds=round(time.monotonic() - started, 3),
                error=_network_error_message("controller model proxy", exc, self.timeout_seconds),
            )

        duration = round(time.monotonic() - started, 3)
        try:
            data = json.loads(raw.decode("utf-8")) if raw else {}
        except json.JSONDecodeError as exc:
            return ModelResponse(ok=False, duration_seconds=duration, error=f"Invalid controller proxy response: {exc}")

        content = str(data.get("content") or "")
        parsed = data.get("parsed") if isinstance(data.get("parsed"), dict) else parse_json_object(content)
        proxy_duration = data.get("duration_seconds")
        try:
            duration = float(proxy_duration) if proxy_duration is not None else duration
        except (TypeError, ValueError):
            pass
        return ModelResponse(
            ok=bool(data.get("ok")),
            content=content,
            parsed=parsed,
            duration_seconds=round(duration, 3),
            error=str(data.get("error") or ""),
        )