"""Automated pricing updater — fetches litellm community pricing data.

Syncs cost_per_1k_input, cost_per_1k_output, and max_tokens from the
litellm model_prices_and_context_window.json database.  Capability
scores are NEVER touched — they remain manually curated.
"""

from __future__ import annotations

import asyncio
import datetime
import json
import os
import tempfile
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import httpx
import structlog

from robot_resources.routing.selector import _DB_PATH, MODELS_DB, validate_models_db

logger = structlog.get_logger()

LITELLM_URL = (
    "https://raw.githubusercontent.com/BerriAI/litellm/main/"
    "model_prices_and_context_window.json"
)

# Explicit litellm key overrides for models whose litellm key differs
# from our model name.  Only Gemini models need the "gemini/" prefix.
_LITELLM_PREFIX_OVERRIDES: dict[str, list[str]] = {
    "gemini-2.5-pro": ["gemini/gemini-2.5-pro", "gemini-2.5-pro"],
    "gemini-2.5-flash": ["gemini/gemini-2.5-flash", "gemini-2.5-flash"],
    "gemini-2.5-flash-lite": ["gemini/gemini-2.5-flash-lite", "gemini-2.5-flash-lite"],
}

# Auto-derived from MODELS_DB at module load.  Adding a model to
# models_db.json requires zero changes here — only Gemini overrides
# need explicit entries above.
MODEL_NAME_MAP: dict[str, list[str]] = {
    m["name"]: _LITELLM_PREFIX_OVERRIDES.get(m["name"], [m["name"]])
    for m in MODELS_DB
}

_PRECISION = 6  # decimal places for float comparison

# ------------------------------------------------------------------
# Cooldown — prevents excessive litellm fetches
# ------------------------------------------------------------------

_DEFAULT_COOLDOWN_DIR = Path.home() / ".robot-resources"
_COOLDOWN_FILENAME = ".pricing_last_fetch"
_COOLDOWN_SECONDS: float = 300.0  # 5 minutes


def check_cooldown(*, base_dir: Path | None = None) -> float | None:
    """Check if cooldown is active.

    Returns seconds remaining if cooldown is active, or ``None`` if the
    fetch may proceed.  Missing, corrupt, or future-dated cooldown files
    are treated as *no cooldown* (never block on metadata).
    """
    path = (base_dir or _DEFAULT_COOLDOWN_DIR) / _COOLDOWN_FILENAME
    try:
        ts = float(path.read_text().strip())
    except (FileNotFoundError, OSError, ValueError):
        return None
    elapsed = time.time() - ts
    if elapsed < 0:  # clock skew
        return None
    remaining = _COOLDOWN_SECONDS - elapsed
    return remaining if remaining > 0 else None


def _write_cooldown(*, base_dir: Path | None = None) -> None:
    """Write current timestamp to cooldown file."""
    path = (base_dir or _DEFAULT_COOLDOWN_DIR) / _COOLDOWN_FILENAME
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(str(time.time()))


class PricingFetchError(Exception):
    """Raised when litellm pricing data cannot be fetched or parsed."""


@dataclass
class PricingChange:
    """A single field change for a model."""

    model_name: str
    field: str  # cost_per_1k_input | cost_per_1k_output | max_tokens
    old_value: float
    new_value: float


@dataclass
class NewModelEntry:
    """A model found in litellm but not in our catalog."""

    name: str
    provider: str
    cost_per_1k_input: float
    cost_per_1k_output: float
    max_tokens: int
    litellm_key: str


@dataclass
class UpdateResult:
    """Result of a pricing update operation."""

    changes: list[PricingChange] = field(default_factory=list)
    skipped: list[str] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)
    new_models: list[NewModelEntry] = field(default_factory=list)
    deprecated_models: list[str] = field(default_factory=list)
    fetched_at: str = ""
    cooldown_remaining: float | None = None


# ------------------------------------------------------------------
# Fetch
# ------------------------------------------------------------------


async def fetch_litellm_pricing(
    url: str = LITELLM_URL,
    *,
    timeout: float = 30.0,
) -> dict[str, Any]:
    """Fetch the litellm pricing JSON.

    Raises :class:`PricingFetchError` on HTTP or parse failure.
    """
    try:
        async with httpx.AsyncClient() as client:
            resp = await client.get(url, timeout=timeout, follow_redirects=True)
            resp.raise_for_status()
            data = resp.json()
    except httpx.HTTPStatusError as exc:
        raise PricingFetchError(
            f"HTTP {exc.response.status_code} fetching litellm pricing"
        ) from exc
    except httpx.TimeoutException as exc:
        raise PricingFetchError("Timeout fetching litellm pricing") from exc
    except httpx.HTTPError as exc:
        raise PricingFetchError(f"HTTP error fetching litellm pricing: {exc}") from exc
    except (json.JSONDecodeError, ValueError) as exc:
        raise PricingFetchError(f"Invalid JSON from litellm: {exc}") from exc

    if not isinstance(data, dict):
        raise PricingFetchError("litellm response is not a JSON object")

    return data


# ------------------------------------------------------------------
# Transform
# ------------------------------------------------------------------


def resolve_litellm_key(
    model_name: str,
    litellm_data: dict[str, Any],
) -> str | None:
    """Find the litellm key for one of our model names.

    Tries :data:`MODEL_NAME_MAP` candidates in order, returns first match.
    Falls back to the model name itself if there is no mapping entry.
    """
    candidates = MODEL_NAME_MAP.get(model_name, [model_name])
    for candidate in candidates:
        if candidate in litellm_data:
            return candidate
    return None


def extract_pricing(litellm_entry: dict[str, Any]) -> dict[str, float | None]:
    """Convert a single litellm entry to our schema fields.

    litellm stores cost **per token**; we store cost **per 1 000 tokens**.

    Returns a dict with keys ``cost_per_1k_input``, ``cost_per_1k_output``,
    and ``max_tokens``.  Values are ``None`` when the source field is missing.
    """
    # Skip non-chat entries (e.g. embedding, image_generation).
    mode = litellm_entry.get("mode")
    if mode is not None and mode != "chat":
        return {"cost_per_1k_input": None, "cost_per_1k_output": None, "max_tokens": None}

    result: dict[str, float | None] = {}

    raw_input = litellm_entry.get("input_cost_per_token")
    result["cost_per_1k_input"] = (
        round(raw_input * 1000, _PRECISION) if raw_input is not None else None
    )

    raw_output = litellm_entry.get("output_cost_per_token")
    result["cost_per_1k_output"] = (
        round(raw_output * 1000, _PRECISION) if raw_output is not None else None
    )

    raw_max = litellm_entry.get("max_input_tokens")
    result["max_tokens"] = int(raw_max) if raw_max is not None else None

    return result


# ------------------------------------------------------------------
# Merge
# ------------------------------------------------------------------


def merge_pricing(
    models_db: list[dict[str, Any]],
    litellm_data: dict[str, Any],
) -> UpdateResult:
    """Merge litellm pricing into an existing models DB **in place**.

    * Only updates existing models — never adds new ones.
    * Preserves ``capabilities``, ``provider``, and ``name``.
    * Sets ``last_updated`` to today for any model whose pricing changed.
    """
    today = datetime.date.today().isoformat()
    result = UpdateResult(fetched_at=today)

    for model in models_db:
        name = model.get("name", "<unnamed>")
        litellm_key = resolve_litellm_key(name, litellm_data)

        if litellm_key is None:
            result.skipped.append(name)
            logger.info("pricing_updater_skip", model=name, reason="not found in litellm")
            continue

        new_vals = extract_pricing(litellm_data[litellm_key])
        changed = False

        for field_name in ("cost_per_1k_input", "cost_per_1k_output", "max_tokens"):
            new_val = new_vals.get(field_name)
            if new_val is None:
                continue

            old_val = model.get(field_name)
            # Float comparison with rounding
            old_rounded = round(float(old_val), _PRECISION) if old_val is not None else None
            if old_rounded is not None and old_rounded == round(float(new_val), _PRECISION):
                continue

            result.changes.append(
                PricingChange(
                    model_name=name,
                    field=field_name,
                    old_value=float(old_val) if old_val is not None else 0.0,
                    new_value=float(new_val),
                )
            )
            model[field_name] = new_val
            changed = True

        if changed:
            model["last_updated"] = today

    if result.skipped:
        logger.warning(
            "pricing_updater_skipped_models",
            count=len(result.skipped),
            models=result.skipped,
        )

    # Detect new models from our providers that we don't have yet
    result.new_models = detect_new_models(models_db, litellm_data)
    if result.new_models:
        logger.info(
            "pricing_updater_new_models",
            count=len(result.new_models),
            models=[m.name for m in result.new_models],
        )

    # Detect deprecated models (in our DB but not in litellm)
    result.deprecated_models = detect_deprecated_models(models_db, litellm_data)
    if result.deprecated_models:
        logger.warning(
            "pricing_updater_deprecated_models",
            count=len(result.deprecated_models),
            models=result.deprecated_models,
        )

    return result


# ------------------------------------------------------------------
# New / Deprecated Model Detection
# ------------------------------------------------------------------

# litellm provider name → our provider name
_LITELLM_PROVIDER_MAP: dict[str, str] = {
    "openai": "openai",
    "anthropic": "anthropic",
    "vertex_ai": "google",
    "gemini": "google",
}

# Conservative capabilities for auto-added models.
# Low enough to avoid routing critical tasks to unknown models,
# high enough that the model is considered for cost savings.
_DEFAULT_CAPABILITIES: dict[str, float] = {
    "overall": 0.3,
    "coding": 0.3,
    "reasoning": 0.3,
    "analysis": 0.3,
    "simple_qa": 0.3,
    "creative": 0.3,
}


def detect_new_models(
    models_db: list[dict[str, Any]],
    litellm_data: dict[str, Any],
) -> list[NewModelEntry]:
    """Find chat models in litellm from our providers that we don't have."""
    existing_names = {m["name"] for m in models_db}
    new_models: list[NewModelEntry] = []

    for key, entry in litellm_data.items():
        if not isinstance(entry, dict):
            continue

        # Only chat models
        mode = entry.get("mode")
        if mode is not None and mode != "chat":
            continue

        # Map litellm provider to ours
        litellm_provider = entry.get("litellm_provider", "")
        our_provider = _LITELLM_PROVIDER_MAP.get(litellm_provider)
        if our_provider is None:
            continue

        # Strip provider prefix from key to get model name
        model_name = key.split("/", 1)[-1] if "/" in key else key

        # Skip if we already have it
        if model_name in existing_names:
            continue

        # Skip variants (fine-tuned, preview, dated versions of existing models)
        is_variant = any(
            model_name.startswith(e + "-") or model_name.startswith(e + ":")
            for e in existing_names
        )
        if is_variant:
            continue

        # Need pricing data
        pricing = extract_pricing(entry)
        if pricing["cost_per_1k_input"] is None or pricing["cost_per_1k_output"] is None:
            continue

        max_tokens = int(pricing["max_tokens"] or 128_000)

        new_models.append(NewModelEntry(
            name=model_name,
            provider=our_provider,
            cost_per_1k_input=pricing["cost_per_1k_input"],
            cost_per_1k_output=pricing["cost_per_1k_output"],
            max_tokens=max_tokens,
            litellm_key=key,
        ))

    return new_models


def detect_deprecated_models(
    models_db: list[dict[str, Any]],
    litellm_data: dict[str, Any],
) -> list[str]:
    """Find models in our DB that no longer exist in litellm."""
    deprecated: list[str] = []
    for model in models_db:
        name = model.get("name", "")
        if resolve_litellm_key(name, litellm_data) is None:
            deprecated.append(name)
    return deprecated


# ------------------------------------------------------------------
# Orchestrator
# ------------------------------------------------------------------


def _write_models_db(models_db: list[dict[str, Any]], db_path: Path) -> None:
    """Atomic write: temp file in same directory → os.replace."""
    content = json.dumps(models_db, indent=2, ensure_ascii=False) + "\n"
    fd, tmp_path = tempfile.mkstemp(
        dir=db_path.parent,
        prefix=".models_db_",
        suffix=".tmp",
    )
    closed = False
    try:
        os.write(fd, content.encode())
        os.close(fd)
        closed = True
        os.replace(tmp_path, db_path)
    except BaseException:
        if not closed:
            os.close(fd)
        if os.path.exists(tmp_path):
            os.unlink(tmp_path)
        raise


async def async_apply_update(
    *,
    dry_run: bool = False,
    force: bool = False,
    db_path: Path | None = None,
    url: str = LITELLM_URL,
    cooldown_dir: Path | None = None,
) -> UpdateResult:
    """Async entry point — safe to call from existing event loops (e.g. FastAPI).

    Args:
        dry_run: If ``True``, compute changes but do not write to disk.
        force: If ``True``, bypass cooldown check.
        db_path: Override path to ``models_db.json`` (default: package path).
        url: Override litellm URL (useful for testing).
        cooldown_dir: Override directory for cooldown file (default: ~/.robot-resources).

    Returns:
        :class:`UpdateResult` with all changes, skipped models, and errors.
    """
    if not force:
        remaining = check_cooldown(base_dir=cooldown_dir)
        if remaining is not None:
            elapsed = int(_COOLDOWN_SECONDS - remaining)
            logger.info("pricing_updater_cooldown", remaining_seconds=round(remaining))
            return UpdateResult(
                cooldown_remaining=remaining,
                errors=[f"Last fetch was {elapsed}s ago. Use --force to override."],
            )

    if db_path is None:
        db_path = _DB_PATH

    with open(db_path) as f:
        models_db: list[dict[str, Any]] = json.load(f)

    litellm_data = await fetch_litellm_pricing(url)

    result = merge_pricing(models_db, litellm_data)

    if result.changes and not dry_run:
        validate_models_db(models_db)
        _write_models_db(models_db, db_path)
        _write_cooldown(base_dir=cooldown_dir)
        logger.info(
            "pricing_updater_written",
            path=str(db_path),
            changes=len(result.changes),
        )

    return result


def apply_update(
    *,
    dry_run: bool = False,
    force: bool = False,
    db_path: Path | None = None,
    url: str = LITELLM_URL,
    cooldown_dir: Path | None = None,
) -> UpdateResult:
    """Sync entry point: wraps :func:`async_apply_update` via ``asyncio.run()``.

    Do NOT call from async context (e.g. FastAPI handlers, MCP tools) —
    use :func:`async_apply_update` instead.
    """
    return asyncio.run(
        async_apply_update(
            dry_run=dry_run, force=force, db_path=db_path, url=url,
            cooldown_dir=cooldown_dir,
        )
    )
