"""Cache pinned official adapter/base snapshots before loading model code."""

import json
from pathlib import Path

from huggingface_hub import HfApi, snapshot_download

from visual_doc_rag.data import write_json

MODEL_ID = "ModernVBERT/colmodernvbert"


def cache_model(data_dir: Path, revision: str = "main", base_revision: str | None = None) -> dict:
    root = data_dir / "models"
    manifest_path = root / "manifest.json"
    if manifest_path.exists():
        cached = json.loads(manifest_path.read_text(encoding="utf-8"))
        if revision in {"main", cached["adapter_revision"]} and base_revision in {
            None,
            cached["base_revision"],
        }:
            return cached
    api = HfApi()
    adapter_sha = api.model_info(MODEL_ID, revision=revision).sha
    adapter_dir = root / "adapter" / adapter_sha
    snapshot_download(
        MODEL_ID,
        revision=adapter_sha,
        local_dir=adapter_dir,
        cache_dir=root / "hub",
        allow_patterns=["*.json", "*.safetensors", "*.jinja", "README.md", "LICENSE*"],
        max_workers=2,
    )
    adapter = json.loads((adapter_dir / "adapter_config.json").read_text(encoding="utf-8"))
    base_id = adapter["base_model_name_or_path"]
    base_sha = api.model_info(
        base_id, revision=base_revision or adapter.get("revision") or "main"
    ).sha
    base_dir = root / "base" / base_sha
    snapshot_download(
        base_id,
        revision=base_sha,
        local_dir=base_dir,
        cache_dir=root / "hub",
        allow_patterns=["*.json", "*.safetensors", "*.jinja", "README.md", "LICENSE*"],
        max_workers=2,
    )
    manifest = {
        "model_id": MODEL_ID,
        "adapter_revision": adapter_sha,
        "base_model_id": base_id,
        "base_revision": base_sha,
        "adapter_path": str(adapter_dir.relative_to(root)),
        "base_path": str(base_dir.relative_to(root)),
        "license": "MIT",
    }
    write_json(manifest_path, manifest)
    return manifest