"""LLM client — single `chat()` entry point that dispatches to the active provider.

Provider selection is via `settings.llm_provider`. Supported:
- "groq"      — hits Groq via the openai SDK (groq_base_url + groq_api_key).
- "anthropic" — native Anthropic SDK.
- "openai"    — openai SDK with the standard base_url (also handles vLLM /
                any OpenAI-compatible endpoint).

Every call opens a Langfuse `generation` span automatically, tagged with the
model name and populated with token usage from the provider's response.
Spans nest under whatever observation the caller has open, so the trace
hierarchy shows `text_to_sql > llm.chat` with usage on the inner node.
"""

from __future__ import annotations

from dataclasses import dataclass
from functools import lru_cache
from typing import Callable

from anthropic import Anthropic
from openai import OpenAI

from api import langfuse_client as lf
from api.config import get_settings


# ── Provider clients (cached) ──

@lru_cache(maxsize=1)
def _anthropic() -> Anthropic:
    return Anthropic(api_key=get_settings().anthropic_api_key)


@lru_cache(maxsize=1)
def _groq() -> OpenAI:
    s = get_settings()
    return OpenAI(api_key=s.groq_api_key, base_url=s.groq_base_url)


@lru_cache(maxsize=1)
def _openai() -> OpenAI:
    s = get_settings()
    return OpenAI(api_key=s.openai_api_key, base_url=s.openai_base_url)


# ── Per-provider chat impls ──

@dataclass
class _Call:
    text: str
    usage: dict[str, int]   # {"input": N, "output": M, "total": N+M}


def _chat_anthropic(system: str, user: str, max_tokens: int) -> _Call:
    s = get_settings()
    msg = _anthropic().messages.create(
        model=s.anthropic_model,
        max_tokens=max_tokens,
        system=system,
        messages=[{"role": "user", "content": user}],
    )
    text = "".join(b.text for b in msg.content if getattr(b, "type", None) == "text").strip()
    return _Call(
        text=text,
        usage={
            "input": int(msg.usage.input_tokens),
            "output": int(msg.usage.output_tokens),
            "total": int(msg.usage.input_tokens + msg.usage.output_tokens),
        },
    )


def _chat_openai_compat(client: OpenAI, model: str, system: str, user: str, max_tokens: int) -> _Call:
    resp = client.chat.completions.create(
        model=model,
        max_tokens=max_tokens,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
    )
    text = (resp.choices[0].message.content or "").strip()
    u = resp.usage
    usage = {
        "input": int(u.prompt_tokens) if u else 0,
        "output": int(u.completion_tokens) if u else 0,
        "total": int(u.total_tokens) if u else 0,
    }
    return _Call(text=text, usage=usage)


def _chat_groq(system: str, user: str, max_tokens: int) -> _Call:
    return _chat_openai_compat(_groq(), get_settings().groq_model, system, user, max_tokens)


def _chat_openai(system: str, user: str, max_tokens: int) -> _Call:
    return _chat_openai_compat(_openai(), get_settings().openai_model, system, user, max_tokens)


_PROVIDERS: dict[str, Callable[[str, str, int], _Call]] = {
    "groq": _chat_groq,
    "anthropic": _chat_anthropic,
    "openai": _chat_openai,
}


def _active_model() -> str:
    s = get_settings()
    return {
        "groq": s.groq_model,
        "anthropic": s.anthropic_model,
        "openai": s.openai_model,
    }.get(s.llm_provider.lower(), s.llm_provider)


# ── Public surface ──

def chat(*, system: str, user: str, max_tokens: int = 1024, span_name: str = "llm.chat") -> str:
    """Run a single chat completion. Opens a Langfuse generation span around
    the call with the model name + token usage. `span_name` is the label
    shown in Langfuse — pass something descriptive (e.g. "text_to_sql.gen")
    so traces are scannable."""
    provider = get_settings().llm_provider.lower()
    impl = _PROVIDERS.get(provider)
    if impl is None:
        raise ValueError(
            f"unknown llm_provider {provider!r}; supported: {sorted(_PROVIDERS)}"
        )

    with lf.span(
        span_name,
        as_type="generation",
        model=_active_model(),
        input={"system": system, "user": user},
    ) as gen:
        call = impl(system, user, max_tokens)
        gen.update(output=call.text, usage_details=call.usage)
        return call.text