from __future__ import annotations

import json

from common import ROOT, iter_model_card_paths, read_yaml_mapping


INDEX_PATH = ROOT / "registry" / "index.json"


def simplify_card(path: Path, data: dict) -> dict:
    deployment = data.get("deployment") or {}
    return {
        "id": data.get("id"),
        "name": data.get("name"),
        "type": data.get("type"),
        "status": data.get("status"),
        "task": data.get("task") or [],
        "language": data.get("language") or [],
        "source": data.get("source"),
        "upstream_id": data.get("upstream_id"),
        "license": data.get("license"),
        "storage_path": data.get("storage_path"),
        "format": data.get("format"),
        "quantization": data.get("quantization"),
        "runtime": deployment.get("runtime"),
        "served_model_name": deployment.get("served_model_name"),
        "card_path": str(path.relative_to(ROOT)).replace("\\", "/"),
    }


def main() -> int:
    models = []
    for path in iter_model_card_paths():
        models.append(simplify_card(path, read_yaml_mapping(path)))

    index = {
        "schema_version": 1,
        "models": models,
    }

    INDEX_PATH.write_text(
        json.dumps(index, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(f"Wrote {INDEX_PATH.relative_to(ROOT)} with {len(models)} model(s).")
    return 0


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