#!/usr/bin/env python3
"""Synchronize llama.cpp models from Hugging Face into a managed local directory."""

from __future__ import annotations

import argparse
import configparser
import json
import os
import re
import shutil
import sys
from pathlib import Path
from typing import Any

from huggingface_hub import HfApi, snapshot_download


def _sanitize_slug(value: str) -> str:
    slug = re.sub(r"[^a-zA-Z0-9._-]+", "-", value.strip())
    slug = slug.strip("-._")
    return slug.lower() or "model"


def _resolve_repo_file(api: HfApi, model: dict[str, Any]) -> str:
    repo_files = api.list_repo_files(repo_id=model["model_id"], revision=model["revision"], repo_type="model")

    hf_file = model.get("hf_file")
    if hf_file:
        if hf_file in repo_files:
            return hf_file
        raise RuntimeError(
            f"Requested hf_file '{hf_file}' was not found in repo "
            f"{model['model_id']}@{model['revision']}"
        )

    quant_upper = (model.get("quant") or "").upper()
    ggufs = [item for item in repo_files if item.lower().endswith(".gguf")]
    matches = [item for item in ggufs if quant_upper in Path(item).name.upper()]
    if matches:
        return sorted(matches, key=lambda name: (len(Path(name).name), name))[0]

    raise RuntimeError(
        f"No GGUF file containing quant '{model.get('quant')}' found in "
        f"repo {model['model_id']}@{model['revision']}. Set hf_file explicitly if needed."
    )


def _load_models(models_file: Path) -> list[dict[str, Any]]:
    data = json.loads(models_file.read_text(encoding="utf-8"))
    if not isinstance(data, list):
        raise RuntimeError("models file must contain a JSON array")

    normalized: list[dict[str, Any]] = []
    for item in data:
        if not isinstance(item, dict):
            raise RuntimeError("each model entry must be an object")
        if item.get("enable", True) is False:
            continue

        model_id = str(item.get("model_id", "")).strip()
        quant = str(item.get("quant", "")).strip()
        revision = str(item.get("revision", "")).strip()
        hf_file = str(item.get("hf_file", "")).strip() or None
        name = str(item.get("name", "")).strip()
        preset = item.get("preset", {})

        if preset is None:
            preset = {}
        if not isinstance(preset, dict):
            raise RuntimeError(f"preset must be an object for model '{name or model_id or 'unknown'}'")

        if not name:
            raise RuntimeError(f"name is required for model '{model_id or 'unknown'}'")
        if not model_id:
            raise RuntimeError("model_id is required for all models")
        if not quant and not hf_file:
            raise RuntimeError(f"quant or hf_file is required for model '{model_id}'")
        if not revision:
            raise RuntimeError(f"revision is required for model '{model_id}'")

        normalized.append(
            {
                "model_id": model_id,
                "quant": quant,
                "revision": revision,
                "hf_file": hf_file,
                "name": name,
                "preset": preset,
            }
        )

    return normalized


def _load_preset_global(preset_global_file: Path | None) -> dict[str, Any]:
    if preset_global_file is None:
        return {}

    if not preset_global_file.exists():
        return {}

    data = json.loads(preset_global_file.read_text(encoding="utf-8"))
    if data is None:
        return {}
    if not isinstance(data, dict):
        raise RuntimeError("global preset options must be a JSON object")
    return data


def _to_preset_value(value: Any) -> str:
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, (int, float, str)):
        return str(value)
    return json.dumps(value, separators=(",", ":"))


def _normalize_preset_options(options: dict[str, Any], scope: str) -> dict[str, str]:
    normalized: dict[str, str] = {}
    for key, value in options.items():
        key_str = str(key).strip()
        if not key_str:
            raise RuntimeError(f"{scope} preset option keys must be non-empty")
        if value is None:
            continue
        normalized[key_str] = _to_preset_value(value)
    return normalized


def _download_model(model: dict[str, Any], target_dir: Path, repo_file: str, dry_run: bool) -> None:
    if dry_run:
        return

    snapshot_download(
        repo_id=model["model_id"],
        revision=model["revision"],
        local_dir=str(target_dir),
        allow_patterns=[repo_file],
        local_dir_use_symlinks=False,
    )


def _write_preset(preset_file: Path, entries: list[dict[str, Any]], global_options: dict[str, Any]) -> None:
    parser = configparser.ConfigParser(interpolation=None)
    parser.optionxform = str
    parser["*"] = _normalize_preset_options(global_options, "global")

    for entry in entries:
        model_options = _normalize_preset_options(entry.get("preset", {}), f"model '{entry['name']}'")
        model_options["model"] = entry["container_model_path"]
        parser[entry["name"]] = model_options

    preset_file.parent.mkdir(parents=True, exist_ok=True)
    with preset_file.open("w", encoding="utf-8") as fh:
        # Router preset format supports top-level version key.
        fh.write("version = 1\n\n")
        parser.write(fh)


def _prune_unmanaged(managed_dir: Path, links_dir: Path, keep_slugs: set[str], dry_run: bool) -> tuple[list[str], list[str]]:
    pruned_dirs: list[str] = []
    pruned_links: list[str] = []

    if managed_dir.exists():
        for child in managed_dir.iterdir():
            if not child.is_dir():
                continue
            if child.name == ".router":
                continue
            if child.name not in keep_slugs:
                pruned_dirs.append(child.name)
                if not dry_run:
                    shutil.rmtree(child)

    if links_dir.exists():
        for child in links_dir.iterdir():
            if child.suffix != ".gguf":
                continue
            slug = child.stem
            if slug not in keep_slugs:
                pruned_links.append(child.name)
                if not dry_run:
                    child.unlink(missing_ok=True)

    return pruned_dirs, pruned_links


def main() -> int:
    parser = argparse.ArgumentParser(description="Sync Hugging Face GGUF models for llama.cpp router mode")
    parser.add_argument("--models-file", required=True)
    parser.add_argument("--managed-dir", required=True)
    parser.add_argument("--links-dir", required=True)
    parser.add_argument("--manifest-file", required=True)
    parser.add_argument("--preset-file", required=True)
    parser.add_argument("--preset-global-file")
    parser.add_argument("--container-links-dir", required=True)
    parser.add_argument("--prune", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    models_file = Path(args.models_file)
    managed_dir = Path(args.managed_dir)
    links_dir = Path(args.links_dir)
    manifest_file = Path(args.manifest_file)
    preset_file = Path(args.preset_file)
    preset_global_file = Path(args.preset_global_file) if args.preset_global_file else None

    managed_dir.mkdir(parents=True, exist_ok=True)
    links_dir.mkdir(parents=True, exist_ok=True)
    manifest_file.parent.mkdir(parents=True, exist_ok=True)
    preset_file.parent.mkdir(parents=True, exist_ok=True)

    models = _load_models(models_file)
    global_options = _load_preset_global(preset_global_file)
    api = HfApi()
    entries: list[dict[str, Any]] = []

    for model in models:
        slug = _sanitize_slug(model["name"])
        target_dir = managed_dir / slug
        target_dir.mkdir(parents=True, exist_ok=True)
        repo_file = _resolve_repo_file(api, model)

        _download_model(model, target_dir, repo_file, args.dry_run)
        selected_file = target_dir / repo_file
        if not args.dry_run and not selected_file.exists():
            raise RuntimeError(f"Expected downloaded file not found: {selected_file}")

        link_path = links_dir / f"{slug}.gguf"
        container_model_path = f"{args.container_links_dir.rstrip('/')}/{slug}.gguf"

        if not args.dry_run:
            if link_path.exists() or link_path.is_symlink():
                link_path.unlink()
            # Keep symlink targets relative so they remain valid inside the
            # container-mounted /models tree.
            relative_target = os.path.relpath(selected_file, start=link_path.parent)
            link_path.symlink_to(relative_target)

        entries.append(
            {
                "name": model["name"],
                "slug": slug,
                "model_id": model["model_id"],
                "revision": model["revision"],
                "quant": model["quant"],
                "hf_file": model.get("hf_file"),
                "preset": model.get("preset", {}),
                "repo_file": repo_file,
                "selected_file": str(selected_file),
                "link_path": str(link_path),
                "container_model_path": container_model_path,
            }
        )

    pruned_dirs: list[str] = []
    pruned_links: list[str] = []
    if args.prune:
        keep_slugs = {entry["slug"] for entry in entries}
        pruned_dirs, pruned_links = _prune_unmanaged(managed_dir, links_dir, keep_slugs, args.dry_run)

    result = {
        "models": entries,
        "dry_run": args.dry_run,
        "prune": args.prune,
        "pruned_dirs": pruned_dirs,
        "pruned_links": pruned_links,
    }

    if not args.dry_run:
        manifest_file.write_text(json.dumps(result, indent=2), encoding="utf-8")
        _write_preset(preset_file, entries, global_options)

    print(json.dumps(result))
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        raise SystemExit(1)