"""Regenerate LLM descriptions from the Ember model registry.

Reads the existing ``configs/LLM_Descriptions.json`` file, looks up each
configured model in the Ember catalog/pricing APIs, and rewrites the JSON with
fresh feature text and normalized per-token pricing. This keeps the
GraphRouter fork aligned with the authoritative model metadata without
hand-editing values.
"""

from __future__ import annotations

import json
from pathlib import Path

from ember.models.catalog import get_model_info
from ember.models.pricing import get_model_cost


BASE_DIR = Path(__file__).resolve().parents[1]
CONFIG_PATH = BASE_DIR / "configs" / "LLM_Descriptions.json"
CACHE_PATH = BASE_DIR / "cache" / "llm_description_runs.jsonl"
PROVIDER_NAMES = {
    "openai": "OpenAI",
    "anthropic": "Anthropic",
    "google": "Google",
}


def _build_feature(
    model_name: str,
    *,
    description: str,
    provider: str,
    context: int,
    input_price: float,
    output_price: float,
) -> str:
    """Compose the feature string surfaced to GraphRouter."""

    provider_name = PROVIDER_NAMES.get(provider, provider.title())
    return (
        f"{description} Provided by {provider_name} with a {context:,}-token context window. "
        f"Estimated pricing: ${input_price:.2f} per 1M input tokens and "
        f"${output_price:.2f} per 1M output tokens."
    )


def _load_overrides() -> dict[str, str]:
    """Return the most recent cached LLM-authored descriptions, if any."""

    if not CACHE_PATH.exists():
        return {}
    try:
        with CACHE_PATH.open("r", encoding="utf-8") as f:
            lines = [line.strip() for line in f if line.strip()]
        if not lines:
            return {}
        latest = json.loads(lines[-1])
        descriptions = latest.get("descriptions", {})
        if isinstance(descriptions, dict):
            return {k: str(v) for k, v in descriptions.items()}
    except Exception:
        return {}
    return {}


def update_llm_descriptions() -> None:
    """Rewrite ``LLM_Descriptions.json`` using catalog data and cached prose."""

    raw = json.loads(CONFIG_PATH.read_text())
    overrides = _load_overrides()
    updated: dict[str, dict[str, object]] = {}

    for display_name, entry in raw.items():
        model_id = entry.get("model")
        if not isinstance(model_id, str):
            raise ValueError(f"Model id missing for {display_name}")

        info = get_model_info(model_id)
        cost = get_model_cost(model_id)

        input_per_million = cost["input"]
        output_per_million = cost["output"]
        if input_per_million == 0 and output_per_million == 0:
            raise ValueError(f"Pricing missing for {model_id}; ensure pricing.yaml has an entry.")

        feature = _build_feature(
            display_name,
            description=info.description,
            provider=info.provider,
            context=info.context_window,
            input_price=input_per_million,
            output_price=output_per_million,
        )

        generated = overrides.get(display_name)
        if generated:
            feature = (
                f"{generated.strip()} Pricing: ${input_per_million:.2f} per 1M input tokens "
                f"and ${output_per_million:.2f} per 1M output tokens (context {info.context_window:,} tokens)."
            )

        updated[display_name] = {
            "feature": feature,
            "input_price": round(input_per_million, 6),
            "output_price": round(output_per_million, 6),
            "model": model_id,
        }

    CONFIG_PATH.write_text(json.dumps(updated, indent=2) + "\n")


if __name__ == "__main__":
    update_llm_descriptions()
