"""OpenAI-compatible LLM client for multi-server support."""

import os
from pathlib import Path
from collections.abc import Awaitable, Callable

from openai import AsyncOpenAI

from backend.models import LLMConfig


# ── Server Config Loading ─────────────────────────────

def _load_llm_configs() -> list[LLMConfig]:
    config_path = os.getenv(
        "LLMS_CONFIG_PATH", str(Path(__file__).parent.parent / "config" / "llms.json")
    )
    import json
    with open(config_path, encoding="utf-8") as f:
        data = json.load(f)
    return [LLMConfig(**item) for item in data]


_llm_configs: list[LLMConfig] = []
_clients: dict[str, AsyncOpenAI] = {}  # alias -> client instance
LLM_TIMEOUT_SECONDS = float(os.getenv("LLM_TIMEOUT_SECONDS", "180"))


def get_llm_configs() -> list[LLMConfig]:
    """Return all LLM configurations."""
    return _llm_configs


def get_client_for_alias(alias: str) -> tuple[AsyncOpenAI, LLMConfig]:
    """Get the AsyncOpenAI client and config for a given alias.
    
    Creates a lazily cached client per alias.
    """
    if alias not in _clients:
        cfg = None
        for c in _llm_configs:
            if c.alias == alias:
                cfg = c
                break
        if cfg is None:
            raise ValueError(f"Unknown LLM alias: {alias}")

        client = AsyncOpenAI(
            api_key=cfg.api_key,
            base_url=cfg.base_url,
            timeout=LLM_TIMEOUT_SECONDS,
            max_retries=0,
        )
        _clients[alias] = client

    return _clients[alias], next(c for c in _llm_configs if c.alias == alias)


async def chat_complete(
    alias: str,
    system_prompt: str,
    user_content: str,
    temperature: float = 0.3,
    on_delta: Callable[[str], Awaitable[None]] | None = None,
) -> str:
    """Send a single-turn chat completion and return the assistant message."""
    client, cfg = get_client_for_alias(alias)
    request = {
        "model": cfg.model,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content},
        ],
        "temperature": temperature,
    }
    if on_delta is None:
        response = await client.chat.completions.create(**request)
        content = response.choices[0].message.content or ""
        return content.strip()

    stream = await client.chat.completions.create(**request, stream=True)
    parts: list[str] = []
    async for event in stream:
        if not event.choices:
            continue
        delta = event.choices[0].delta.content or ""
        if delta:
            parts.append(delta)
            await on_delta(delta)
    return "".join(parts).strip()


def get_context_size(alias: str) -> int | None:
    """Get the context window size for a given alias."""
    for c in _llm_configs:
        if c.alias == alias:
            return c.context_size
    return None


# ── Initialize on import ──────────────────────────────

_llm_configs = _load_llm_configs()