#!/usr/bin/env python3
"""Render OMP's two runtime files from one private catalog source."""

from __future__ import annotations

import argparse
import copy
import sys
from pathlib import Path
from typing import Any

try:
    import yaml
except ModuleNotFoundError:  # pragma: no cover - fresh-install convenience
    sys.exit("PyYAML is required: pip install pyyaml")


BUILTIN_PROVIDERS = {
    "google-antigravity",
    "openai-codex",
    "opencode-zen",
}


def load_yaml(path: Path) -> dict[str, Any]:
    value = yaml.safe_load(path.read_text())
    if not isinstance(value, dict):
        raise ValueError(f"{path} must contain a mapping")
    return value


def load_catalog_dir(dir_path: Path) -> dict[str, Any]:
    catalog: dict[str, Any] = {"version": 1}

    for name in ["settings.yml", "settings.yaml"]:
        p = dir_path / name
        if p.is_file():
            data = load_yaml(p)
            if "settings" in data:
                catalog["settings"] = data["settings"]
                if "version" in data:
                    catalog["version"] = data["version"]
            else:
                catalog["settings"] = data
            break

    for name in ["routing.yml", "routing.yaml"]:
        p = dir_path / name
        if p.is_file():
            data = load_yaml(p)
            if "routing" in data:
                catalog["routing"] = data["routing"]
            else:
                catalog["routing"] = data
            break

    for name in ["providers.yml", "providers.yaml"]:
        p = dir_path / name
        if p.is_file():
            data = load_yaml(p)
            if "providers" in data:
                catalog["providers"] = data["providers"]
            else:
                catalog["providers"] = data
            break

    if "routing" not in catalog or "providers" not in catalog:
        raise ValueError(f"Catalog directory {dir_path} must contain routing.yml and providers.yml")

    return catalog


def load_catalog(path: Path) -> dict[str, Any]:
    if path.is_dir():
        return load_catalog_dir(path)

    sibling_dir = path.parent / "catalog"
    if sibling_dir.is_dir() and path.name in ("catalog.yml", "catalog.yaml"):
        return load_catalog_dir(sibling_dir)

    return load_yaml(path)


def expand(value: Any, tiers: dict[str, Any], pools: dict[str, Any], path: str) -> Any:
    if isinstance(value, str):
        if value.startswith("@tier."):
            name = value.removeprefix("@tier.")
            if name not in tiers or not isinstance(tiers[name], str):
                raise ValueError(f"{path}: unknown tier {value}")
            return expand(tiers[name], tiers, pools, path)
        if value.startswith("@pool."):
            name = value.removeprefix("@pool.")
            if name not in pools or not isinstance(pools[name], list):
                raise ValueError(f"{path}: unknown pool {value}")
            return [expand(item, tiers, pools, f"{path}.pool[{index}]") for index, item in enumerate(pools[name])]
        if value.startswith("@account."):
            reference = value.removeprefix("@account.")
            account_ref, separator, model_id = reference.partition("/")
            if not separator or "." not in account_ref:
                raise ValueError(f"{path}: invalid account reference {value}")
            provider, account_id = account_ref.rsplit(".", 1)
            return f"{provider}-{account_id}/{model_id}"
        return value
    if isinstance(value, list):
        result = []
        for index, item in enumerate(value):
            expanded = expand(item, tiers, pools, f"{path}[{index}]")
            result.extend(expanded if isinstance(expanded, list) else [expanded])
        return result
    if isinstance(value, dict):
        return {
            key: expand(item, tiers, pools, f"{path}.{key}")
            for key, item in value.items()
        }
    return value


def provider_of(selector: str) -> str:
    if "/" not in selector:
        raise ValueError(f"invalid model selector: {selector}")
    return selector.split("/", 1)[0]


def expand_providers(catalog: dict[str, Any]) -> dict[str, Any]:
    rendered = {}
    for name, spec in catalog["providers"].items():
        if not isinstance(spec, dict):
            raise ValueError(f"providers.{name} must be a mapping")
        accounts = spec.get("accounts")
        if not accounts:
            rendered[name] = copy.deepcopy(spec)
            continue
        if not isinstance(accounts, list):
            raise ValueError(f"providers.{name}.accounts must be a list")
        base = {key: value for key, value in spec.items() if key != "accounts"}
        for account in accounts:
            if not isinstance(account, dict) or not account.get("id") or not account.get("apiKey"):
                raise ValueError(f"providers.{name}.accounts entries require id and apiKey")
            provider = copy.deepcopy(base)
            provider["apiKey"] = account["apiKey"]
            provider_id = f"{name}-{account['id']}"
            rendered[provider_id] = provider
    return rendered


def validate_catalog(catalog: dict[str, Any], config: dict[str, Any], models: dict[str, Any]) -> None:
    routing = catalog.get("routing")
    providers = expand_providers(catalog)
    if not isinstance(routing, dict) or not isinstance(providers, dict):
        raise ValueError("catalog requires routing and providers mappings")

    tiers = routing.get("tiers", {})
    pools = routing.get("pools", {})
    if not isinstance(tiers, dict) or not tiers:
        raise ValueError("routing.tiers must not be empty")

    resolved = expand(routing.get("modelRoles", {}), tiers, pools, "routing.modelRoles")
    resolved_chains = expand(routing.get("fallbackChains", {}), tiers, pools, "routing.fallbackChains")
    selectors = list(resolved.values()) if isinstance(resolved, dict) else []
    for chain in resolved_chains.values():
        selectors.extend(chain)
    for target in resolved_chains:
        if "/" in target:
            selectors.append(target)

    known_providers = set(providers) | BUILTIN_PROVIDERS
    for selector in selectors:
        provider = provider_of(selector)
        if provider.endswith("*"):
            provider = provider.removesuffix("*").rstrip("/")
        if provider not in known_providers:
            raise ValueError(f"selector references unknown provider: {selector}")

    for name, provider in providers.items():
        if not isinstance(provider, dict):
            raise ValueError(f"providers.{name} must be a mapping")
        if not provider.get("apiKey"):
            raise ValueError(f"providers.{name}.apiKey must name an env var")
        if not provider.get("baseUrl"):
            raise ValueError(f"providers.{name}.baseUrl is required")

    # Generated files must never retain the source-only tier section.
    if "fallbackTiers" in config or "fallbackTiers" in models:
        raise ValueError("generated OMP files must not contain fallbackTiers")


def render(catalog: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
    routing = catalog["routing"]
    tiers = routing["tiers"]
    pools = routing.get("pools", {})
    config = copy.deepcopy(catalog.get("settings", {}))
    config["modelRoles"] = expand(routing["modelRoles"], tiers, pools, "routing.modelRoles")
    config["cycleOrder"] = copy.deepcopy(routing.get("cycleOrder", ["default", "slow"]))
    retry = copy.deepcopy(config.get("retry", {}))
    retry["fallbackChains"] = expand(routing["fallbackChains"], tiers, pools, "routing.fallbackChains")
    config["retry"] = retry
    return config, {"providers": expand_providers(catalog)}


def dump_yaml(path: Path, value: dict[str, Any]) -> None:
    rendered = yaml.safe_dump(
        value,
        allow_unicode=True,
        default_flow_style=False,
        sort_keys=False,
        width=120,
    )
    path.write_text(
        "# GENERATED FILE. Edit catalog/, then run make omp-config.\n" + rendered
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--models", type=Path, required=True)
    args = parser.parse_args()

    try:
        catalog = load_catalog(args.source)
        config, models = render(catalog)
        validate_catalog(catalog, config, models)
        dump_yaml(args.config, config)
        dump_yaml(args.models, models)
    except (OSError, ValueError, yaml.YAMLError) as exc:
        print(f"OMP catalog render failed: {exc}", file=sys.stderr)
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main())