#!/usr/bin/env python3
"""Generate a reduced LiteLLM model capability table for pipeline inference.

Source data:
  https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json

LiteLLM is MIT licensed. This generator keeps only the fields needed by
csghub-lite to infer local pipeline tags and model input/output modalities.
"""

from __future__ import annotations

import json
import pathlib
import re
import urllib.request


SOURCE_URL = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
OUT = pathlib.Path("internal/server/provider_pipeline_rules_gen.go")

MODE_ALLOWLIST = {
    "image_generation",
    "audio_transcription",
    "audio_speech",
}

CAPABILITY_KEYS = (
    "supports_vision",
    "supports_video_input",
    "supports_audio_input",
    "supports_audio_output",
)


def normalize_key(value: str) -> str:
    return re.sub(r"\s+", "", value.strip().lower())


def go_string(value: str) -> str:
    return json.dumps(value, ensure_ascii=False)


def go_bool(value: bool) -> str:
    return "true" if value else "false"


def should_keep(info: dict) -> bool:
    mode = str(info.get("mode") or "").strip()
    if mode in MODE_ALLOWLIST:
        return True
    return any(bool(info.get(key)) for key in CAPABILITY_KEYS)


def main() -> None:
    with urllib.request.urlopen(SOURCE_URL, timeout=60) as resp:
        data = json.load(resp)

    rules: dict[str, dict] = {}
    for model_id, info in data.items():
        if model_id == "sample_spec" or not isinstance(info, dict) or not should_keep(info):
            continue
        provider = str(info.get("litellm_provider") or "").strip()
        mode = str(info.get("mode") or "").strip()
        rule = {
            "mode": mode,
            "supports_vision": bool(info.get("supports_vision")),
            "supports_video_input": bool(info.get("supports_video_input")),
            "supports_audio_input": bool(info.get("supports_audio_input")),
            "supports_audio_output": bool(info.get("supports_audio_output")),
        }
        keys = {normalize_key(model_id)}
        if provider:
            keys.add(normalize_key(f"{provider}/{model_id}"))
        for key in keys:
            if key:
                rules[key] = rule

    lines = [
        "package server",
        "",
        "// Code generated by scripts/update_litellm_pipeline_rules.py; DO NOT EDIT.",
        "// Source: LiteLLM model_prices_and_context_window.json (MIT License).",
        "",
        "type liteLLMPipelineRule struct {",
        "\tMode string",
        "\tSupportsVision bool",
        "\tSupportsVideoInput bool",
        "\tSupportsAudioInput bool",
        "\tSupportsAudioOutput bool",
        "}",
        "",
        f"var liteLLMPipelineRules = map[string]liteLLMPipelineRule{{",
    ]
    for key in sorted(rules):
        rule = rules[key]
        lines.append(
            f"\t{go_string(key)}: {{Mode: {go_string(rule['mode'])}, "
            f"SupportsVision: {go_bool(rule['supports_vision'])}, "
            f"SupportsVideoInput: {go_bool(rule['supports_video_input'])}, "
            f"SupportsAudioInput: {go_bool(rule['supports_audio_input'])}, "
            f"SupportsAudioOutput: {go_bool(rule['supports_audio_output'])}}},"
        )
    lines.append("}")
    lines.append("")

    OUT.write_text("\n".join(lines), encoding="utf-8")
    print(f"wrote {OUT} with {len(rules)} rules")


if __name__ == "__main__":
    main()
