from __future__ import annotations

import json
import logging
import re
from typing import Any

from .budget import BudgetLedger
from .pricing import TokenRates, cost_usd
from .secrets import get_cloud_credentials


logger = logging.getLogger(__name__)


_JSON_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)```", re.S | re.I)


def _extract_json_object(text: str) -> str:
    text = (text or "").strip()
    if not text:
        return text
    fence = _JSON_FENCE_RE.search(text)
    if fence:
        text = fence.group(1).strip()
    start = text.find("{")
    end = text.rfind("}")
    if start >= 0 and end > start:
        return text[start : end + 1]
    return text


class CloudLLMClient:
    def __init__(
        self,
        *,
        budget: BudgetLedger,
        rates: TokenRates,
        text_model: str = "qwen-plus",
        vlm_model: str = "qwen2-omni-flash",
        fast_model: str | None = None,
        timeout: float = 60.0,
    ) -> None:
        self.budget = budget
        self.rates = rates
        self.text_model = text_model
        # cheap/low-latency model for scope checks and query rewriting
        self.fast_model = fast_model or text_model
        self.vlm_model = vlm_model
        self.embed_model = "text-embedding-v3"
        # Large diagnosis prompts (retrieved context + 3 full causes) routinely take
        # >20s to generate, so default to a roomier timeout.
        self.timeout = float(timeout)

        api_key, base_url = get_cloud_credentials()
        try:
            from openai import OpenAI
        except Exception as e:
            raise RuntimeError("pip install openai") from e

        self._client = OpenAI(api_key=api_key, base_url=base_url, timeout=self.timeout, max_retries=0)

    def _usage_cost(
        self,
        usage: Any,
        *,
        assumed_vision_tokens: int = 0,
        multimodal_output: bool = True,
    ) -> float:
        if usage is None:
            return 0.0
        prompt = int(getattr(usage, "prompt_tokens", 0) or 0)
        completion = int(getattr(usage, "completion_tokens", 0) or 0)
        vision = min(prompt, assumed_vision_tokens) if assumed_vision_tokens else 0
        text_in = max(0, prompt - vision)
        return cost_usd(
            rates=self.rates,
            vision_tokens=vision,
            text_in_tokens=text_in,
            text_out_tokens=completion,
            multimodal_output=multimodal_output,
        )

    def chat_text_json(
        self,
        *,
        messages: list[dict],
        category: str = "diagnosis",
        estimated_usd: float = 0.02,
        temperature: float = 0.15,
        model: str | None = None,
    ) -> str:
        self.budget.require(estimated_usd, category=category)
        use_model = model or self.text_model
        try:
            resp = self._client.chat.completions.create(
                model=use_model,
                messages=messages,
                temperature=temperature,
                response_format={"type": "json_object"},
                timeout=self.timeout,
            )
        except Exception:
            logger.warning("JSON-mode text chat failed; retrying without response_format", exc_info=True)
            resp = self._client.chat.completions.create(
                model=use_model,
                messages=messages,
                temperature=temperature,
                timeout=self.timeout,
            )
        content = resp.choices[0].message.content or ""
        actual = self._usage_cost(resp.usage, multimodal_output=False)
        if actual <= 0:
            actual = estimated_usd
        self.budget.charge(actual, category, note=f"text chat ({use_model})")
        self.last_usage = {
            "model": use_model,
            "prompt_tokens": int(getattr(resp.usage, "prompt_tokens", 0) or 0) if resp.usage else 0,
            "completion_tokens": int(getattr(resp.usage, "completion_tokens", 0) or 0) if resp.usage else 0,
            "usd": round(actual, 6),
        }
        return _extract_json_object(content)

    def chat_vision_json(
        self,
        *,
        prompt: str,
        image_b64_url: str,
        category: str = "video_vlm",
        estimated_usd: float = 0.01,
        vision_tokens: int = 0,
        temperature: float = 0.2,
    ) -> str:
        self.budget.require(estimated_usd, category=category)
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": image_b64_url}},
                ],
            }
        ]
        resp = self._client.chat.completions.create(
            model=self.vlm_model,
            messages=messages,
            temperature=temperature,
            response_format={"type": "json_object"},
        )
        content = resp.choices[0].message.content or ""
        actual = self._usage_cost(
            resp.usage,
            assumed_vision_tokens=vision_tokens,
            multimodal_output=True,
        )
        if actual <= 0:
            actual = estimated_usd
        self.budget.charge(actual, category, note="vision chat")
        return content

    def embed_texts(
        self,
        texts: list[str],
        *,
        model: str | None = None,
        category: str = "embedding",
        batch_size: int = 96,
    ) -> list[list[float]]:
        """Batched embeddings (one HTTP call per ``batch_size`` texts)."""
        out: list[list[float]] = []
        for i in range(0, len(texts), batch_size):
            chunk = [t[:8000] for t in texts[i : i + batch_size]]
            self.budget.require(0.0001 * len(chunk), category=category)
            resp = self._client.embeddings.create(model=model or self.embed_model, input=chunk)
            out.extend(d.embedding for d in resp.data)
            usage = getattr(resp, "usage", None)
            tokens = int(getattr(usage, "total_tokens", 0) or 0) if usage else sum(max(50, len(t) // 4) for t in chunk)
            actual = cost_usd(rates=self.rates, text_in_tokens=tokens)
            if actual > 0:
                self.budget.charge(actual, category, note=f"embedding x{len(chunk)}")
        return out

    def embed_text(
        self,
        text: str,
        *,
        model: str | None = None,
        category: str = "embedding",
        estimated_usd: float = 0.0001,
    ) -> list[float]:
        self.budget.require(estimated_usd, category=category)
        resp = self._client.embeddings.create(
            model=model or self.embed_model,
            input=text[:8000],
        )
        vec = resp.data[0].embedding
        usage = getattr(resp, "usage", None)
        tokens = int(getattr(usage, "total_tokens", 0) or 0) if usage else max(50, len(text) // 4)
        actual = cost_usd(rates=self.rates, text_in_tokens=tokens)
        if actual <= 0:
            actual = estimated_usd
        self.budget.charge(actual, category, note="embedding")
        return vec