from __future__ import annotations

import argparse
import hashlib
import io
import json
import tarfile
from datetime import UTC, datetime
from pathlib import Path


def digest(path: Path) -> str:
    value = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(1024 * 1024):
            value.update(chunk)
    return value.hexdigest()


def main() -> None:
    parser = argparse.ArgumentParser(description="Build a versioned ImageFind model bundle")
    parser.add_argument("source", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--version", default="1")
    args = parser.parse_args()
    source = args.source.resolve()
    output = args.output.resolve()
    for required in ("visual/image", "visual/text"):
        if not (source / required).is_dir():
            parser.error(f"missing {required}")
    optional_components = {
        "ocr": ("det.onnx", "rec.onnx", "cls.onnx"),
        "faces": ("detector.xml", "detector.bin", "reidentification.xml", "reidentification.bin"),
    }
    for component, required_files in optional_components.items():
        root = source / component
        if root.exists():
            missing = [name for name in required_files if not (root / name).is_file()]
            if missing:
                parser.error(f"incomplete {component}: missing {', '.join(missing)}")
    audio_root = source / "audio"
    if audio_root.exists() and (not (audio_root / "config.json").is_file() or not list(audio_root.glob("*.xml"))):
        parser.error("incomplete audio: missing config.json or OpenVINO XML files")
    try:
        output.relative_to(source)
    except ValueError:
        pass
    else:
        parser.error("output must be outside the source model directory")

    files = sorted(
        path for path in source.rglob("*") if path.is_file() and path.relative_to(source).as_posix() != "manifest.json"
    )
    created_at = datetime.now(UTC)
    manifest = {
        "format_version": 2,
        "version": args.version,
        "created_at": created_at.isoformat(),
        "source": "imagefind-build-model-bundle",
        "components": {
            "visual": {"version": args.version},
            **({"ocr": {"version": args.version}} if (source / "ocr").is_dir() else {}),
            **({"faces": {"version": args.version}} if (source / "faces").is_dir() else {}),
            **({"audio": {"version": args.version}} if audio_root.is_dir() else {}),
        },
        "files": {path.relative_to(source).as_posix(): digest(path) for path in files},
    }
    manifest_bytes = (json.dumps(manifest, ensure_ascii=False, indent=2) + "\n").encode()
    output.parent.mkdir(parents=True, exist_ok=True)
    with tarfile.open(output, "w:gz") as archive:
        for path in sorted(source.rglob("*")):
            if path.relative_to(source).as_posix() == "manifest.json":
                continue
            archive.add(path, path.relative_to(source).as_posix(), recursive=False)
        info = tarfile.TarInfo("manifest.json")
        info.size = len(manifest_bytes)
        info.mode = 0o644
        info.mtime = int(created_at.timestamp())
        archive.addfile(info, io.BytesIO(manifest_bytes))
    print(f"bundle={output} sha256={digest(output)}")


if __name__ == "__main__":
    main()