#!/usr/bin/env -S uv run python
import argparse
import hashlib
import json
import shutil
import urllib.request
from pathlib import Path

MODEL_ID = "PayamFard123/dermaintel-face-parsing/resnet18.onnx"
SOURCE_REVISION = "c1c5fd399cc72b45e3613019879fc6c17a7d8f6a"
EXPECTED_SHA256 = "0d9bd318e46987c3bdbfacae9e2c0f461cae1c6ac6ea6d43bbe541a91727e33f"
PRIMARY_URL = f"https://huggingface.co/PayamFard123/dermaintel-face-parsing/resolve/{SOURCE_REVISION}/resnet18.onnx"
FALLBACK_URL = "https://github.com/yakhyo/face-parsing/releases/download/weights/resnet18.onnx"
DEFAULT_CACHE_DIR = Path.home() / ".cache" / "chromaframe-vision" / "face-parser"
MODEL_FILENAME = "resnet18.onnx"
MANIFEST_FILENAME = "resnet18.manifest.json"

LABELS = {
    "background": 0,
    "skin": 1,
    "l_brow": 2,
    "r_brow": 3,
    "l_eye": 4,
    "r_eye": 5,
    "eye_g": 6,
    "l_ear": 7,
    "r_ear": 8,
    "ear_r": 9,
    "nose": 10,
    "mouth": 11,
    "u_lip": 12,
    "l_lip": 13,
    "neck": 14,
    "neck_l": 15,
    "cloth": 16,
    "hair": 17,
    "hat": 18,
}


def main() -> int:
    parser = argparse.ArgumentParser(description="Install ChromaFrame baseline face parser outside the repository.")
    parser.add_argument("--cache-dir", type=Path, default=DEFAULT_CACHE_DIR)
    parser.add_argument("--source-file", type=Path, help="Copy an existing local resnet18.onnx instead of downloading.")
    parser.add_argument("--skip-download", action="store_true", help="Use an existing cache file and only rewrite the manifest.")
    args = parser.parse_args()

    result = install_face_parser(args.cache_dir, args.source_file, args.skip_download)
    print(json.dumps(result, indent=2))
    return 0


def install_face_parser(cache_dir: Path, source_file: Path | None = None, skip_download: bool = False) -> dict:
    cache_dir.mkdir(parents=True, exist_ok=True)
    model_path = cache_dir / MODEL_FILENAME
    manifest_path = cache_dir / MANIFEST_FILENAME

    if source_file is not None:
        if not source_file.is_file():
            raise SystemExit("source_file_not_readable")
        shutil.copyfile(source_file, model_path)
    elif not skip_download:
        download_model(model_path)
    elif not model_path.is_file():
        raise SystemExit("cached_model_not_readable")

    digest = sha256_file(model_path)
    if digest != EXPECTED_SHA256:
        model_path.unlink(missing_ok=True)
        manifest_path.unlink(missing_ok=True)
        raise SystemExit("model_checksum_mismatch")

    manifest_path.write_text(json.dumps(manifest(digest), indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return {
        "model": str(model_path),
        "manifest": str(manifest_path),
        "sha256": digest,
        "export": {
            "CHROMAFRAME_FACE_PARSING_ONNX_PATH": str(model_path),
            "CHROMAFRAME_VISION_ASSET_MANIFEST": str(manifest_path),
        },
    }


def download_model(model_path: Path) -> None:
    for url in (PRIMARY_URL, FALLBACK_URL):
        try:
            request = urllib.request.Request(url, headers={"User-Agent": "ChromaFrame face-parser setup"})
            with urllib.request.urlopen(request, timeout=120) as response:
                model_path.write_bytes(response.read())
            return
        except Exception:
            continue
    raise SystemExit("model_download_failed")


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def manifest(digest: str) -> dict:
    return {
        "face_parsing_onnx": {
            "manifest_schema_version": 1,
            "model_id": MODEL_ID,
            "sha256": digest,
            "expected_sha256": EXPECTED_SHA256,
            "license": "MIT model card; weights licensed separately from ChromaFrame code",
            "license_id": "MIT",
            "approved_for_use": True,
            "model_type": "baseline_face_parser",
            "source_url": PRIMARY_URL,
            "source_revision": SOURCE_REVISION,
            "fallback_source": FALLBACK_URL,
            "provenance": "PayamFard123/dermaintel-face-parsing resnet18.onnx; yakhyo/face-parsing weights",
            "labels": LABELS,
            "input": {"width": 512, "height": 512},
            "normalization": {"mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225]},
            "output": {"type": "logits_argmax"},
        }
    }


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