# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Provider-neutral interface: one async call that returns a validated Pydantic object."""
from __future__ import annotations

import asyncio
import hashlib
import json
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, TypeVar

from pydantic import BaseModel, ValidationError

from ..config import load_config, model_info, provider_mode
from .ledger import CostLedger

T = TypeVar("T", bound=BaseModel)


@dataclass
class ProviderResponse:
    parsed: BaseModel | None
    raw_text: str
    model: str
    provider: str
    input_tokens: int = 0
    output_tokens: int = 0
    cached_tokens: int = 0
    cost_usd: float = 0.0
    latency_s: float = 0.0
    request_id: str | None = None
    prompt_hash: str = ""
    validation_retries: int = 0
    error: str | None = None
    extra: dict[str, Any] = field(default_factory=dict)


def prompt_hash(system: str, user: str, schema_name: str) -> str:
    return hashlib.sha256(f"{system}\n\x00\n{user}\n\x00\n{schema_name}".encode()).hexdigest()[:16]


def json_schema_for(schema: type[BaseModel]) -> dict[str, Any]:
    """Strict JSON schema (additionalProperties=false everywhere, all properties required)."""
    js = schema.model_json_schema()

    def tighten(node: Any) -> None:
        if isinstance(node, dict):
            if node.get("type") == "object" and "properties" in node:
                node["additionalProperties"] = False
                node["required"] = list(node["properties"].keys())
            for k in ("default",):
                node.pop(k, None)
            for v in node.values():
                tighten(v)
        elif isinstance(node, list):
            for v in node:
                tighten(v)

    tighten(js)
    return js


class LLMProvider(ABC):
    name: str = "base"

    def __init__(self, ledger: CostLedger | None = None, concurrency: int | None = None):
        cfg = load_config("models")["request_defaults"]
        self.ledger = ledger or CostLedger()
        self.sem = asyncio.Semaphore(concurrency or cfg.get("concurrency", 6))
        self.max_retries = cfg.get("max_retries", 4)
        self.timeout_s = cfg.get("timeout_s", 180)

    @abstractmethod
    async def _raw_complete(self, *, system: str, user: str, schema: type[BaseModel], model: str, effort: str, max_tokens: int) -> tuple[str, dict[str, Any]]:
        """Return (json_text, meta) where meta has input_tokens/output_tokens/cached_tokens/request_id."""

    async def complete(
        self,
        *,
        system: str,
        user: str,
        schema: type[T],
        model: str,
        effort: str = "medium",
        max_tokens: int = 4096,
        tag: str = "",
        validation_retries: int = 2,
    ) -> ProviderResponse:
        ph = prompt_hash(system, user, schema.__name__)
        t0 = time.perf_counter()
        retries = 0
        last_err: str | None = None
        user_msg = user
        tot_in = tot_out = tot_cached = 0
        text = ""
        meta: dict[str, Any] = {}
        parsed: BaseModel | None = None
        async with self.sem:
            for attempt in range(validation_retries + 1):
                try:
                    text, meta = await asyncio.wait_for(
                        self._raw_complete(system=system, user=user_msg, schema=schema, model=model, effort=effort, max_tokens=max_tokens),
                        timeout=self.timeout_s * 2,
                    )
                except Exception as e:  # provider/network errors after SDK retries
                    last_err = f"{type(e).__name__}: {e}"[:500]
                    retries = attempt
                    break
                tot_in += int(meta.get("input_tokens", 0))
                tot_out += int(meta.get("output_tokens", 0))
                tot_cached += int(meta.get("cached_tokens", 0))
                try:
                    parsed = schema.model_validate_json(_extract_json(text))
                    last_err = None
                    break
                except (ValidationError, ValueError) as e:
                    last_err = f"validation: {str(e)[:400]}"
                    retries = attempt + 1
                    user_msg = user + f"\n\nYour previous answer failed validation ({str(e)[:300]}). Return a corrected JSON object only."
        latency = time.perf_counter() - t0
        cost = CostLedger.estimate_cost(model, tot_in, tot_out, tot_cached)
        self.ledger.record(
            provider=self.name, model=model, tag=tag, prompt_hash=ph, input_tokens=tot_in, output_tokens=tot_out,
            cached_tokens=tot_cached, cost_usd=cost, latency_s=round(latency, 3), ok=parsed is not None,
            validation_retries=retries, error=last_err, request_id=meta.get("request_id"),
        )
        return ProviderResponse(
            parsed=parsed, raw_text=text, model=model, provider=self.name, input_tokens=tot_in, output_tokens=tot_out,
            cached_tokens=tot_cached, cost_usd=cost, latency_s=latency, request_id=meta.get("request_id"), prompt_hash=ph,
            validation_retries=retries, error=last_err,
        )


def _extract_json(text: str) -> str:
    text = text.strip()
    if text.startswith("```"):
        text = text.strip("`")
        if text.startswith("json"):
            text = text[4:]
    i, j = text.find("{"), text.rfind("}")
    return text[i : j + 1] if i >= 0 and j > i else text


def get_provider(model: str, ledger: CostLedger | None = None, force_mock: bool | None = None) -> LLMProvider:
    """Instantiate the adapter for a registered model id (mock when keys are absent or AV_PROVIDER_MODE=mock)."""
    from .mock import MockProvider

    mock = force_mock if force_mock is not None else (provider_mode() == "mock")
    if mock or model.startswith("mock"):
        return MockProvider(ledger=ledger)
    vendor = model_info(model)["provider"]
    if vendor == "anthropic":
        from .anthropic_provider import AnthropicProvider

        return AnthropicProvider(ledger=ledger)
    if vendor == "openai":
        from .openai_provider import OpenAIProvider

        return OpenAIProvider(ledger=ledger)
    raise ValueError(vendor)