#!/usr/bin/env python3
"""Dry-run by default uploader for the FLUX.1-dev ConvRot Hugging Face repo."""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import subprocess
import sys
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
MANIFEST_PATH = ROOT / "release_manifest.json"
VERIFY_SCRIPT = ROOT / "scripts/verify_release.py"
VERIFICATION_PATH = ROOT / "release_verification.json"
CHUNK_BYTES = 16 * 1024 * 1024


def parse_args() -> argparse.Namespace:
    manifest = json.loads(MANIFEST_PATH.read_text(encoding="utf-8"))
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo-id", default=manifest["suggested_repo_id"])
    parser.add_argument(
        "--model-root",
        type=Path,
        default=Path(os.environ.get("COMFYUI_MODELS_ROOT", "/mnt/d/comfyui/models")),
    )
    parser.add_argument("--revision", default="main")
    parser.add_argument("--execute", action="store_true", help="Perform network writes; default is dry-run.")
    parser.add_argument(
        "--support-only",
        action="store_true",
        help="Upload documentation/evidence only; do not re-upload model artifacts.",
    )
    parser.add_argument(
        "--create-private",
        action="store_true",
        help="Create a new private model repository before upload.",
    )
    parser.add_argument(
        "--accept-flux-license",
        action="store_true",
        help="Confirm that the uploader reviewed and accepts the upstream FLUX.1-dev license.",
    )
    return parser.parse_args()


def load_manifest() -> dict[str, Any]:
    value = json.loads(MANIFEST_PATH.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise RuntimeError("release manifest must be a JSON object")
    return value


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


def existing_verification_is_current(manifest: dict[str, Any], model_root: Path) -> bool:
    if not VERIFICATION_PATH.is_file():
        return False
    verification = json.loads(VERIFICATION_PATH.read_text(encoding="utf-8"))
    if verification.get("status") != "PASS" or not verification.get("model_hashes_verified"):
        return False
    if verification.get("artifact_total_bytes") != manifest.get("artifact_total_bytes"):
        return False
    support = {record["path"]: record for record in verification.get("support_files", [])}
    for relative in manifest["required_repository_files"]:
        record = support.get(relative)
        path = ROOT / relative
        if record is None or not path.is_file() or path.is_symlink():
            return False
        if path.stat().st_size != record["bytes"] or sha256_file(path) != record["sha256"]:
            return False
    verified_artifacts = {
        record["repo_path"]: record for record in verification.get("artifacts", [])
    }
    verification_mtime_ns = VERIFICATION_PATH.stat().st_mtime_ns
    for artifact in manifest["artifacts"]:
        record = verified_artifacts.get(artifact["repo_path"])
        source = model_root / artifact["source_relative_path"]
        if record is None or not record.get("hash_verified") or not source.is_file():
            return False
        if record.get("sha256") != artifact["sha256"]:
            return False
        if source.stat().st_size != artifact["bytes"] or source.stat().st_mtime_ns > verification_mtime_ns:
            return False
    return True


def main() -> int:
    args = parse_args()
    manifest = load_manifest()
    model_root = args.model_root.expanduser().resolve(strict=True)
    plan = {
        "mode": "execute" if args.execute else "dry-run",
        "repo_id": args.repo_id,
        "revision": args.revision,
        "create_private": args.create_private,
        "support_only": args.support_only,
        "support_root": str(ROOT),
        "artifact_total_bytes": manifest["artifact_total_bytes"],
        "artifacts": [
            {
                "source": str(model_root / artifact["source_relative_path"]),
                "repo_path": artifact["repo_path"],
                "bytes": artifact["bytes"],
                "sha256": artifact["sha256"],
            }
            for artifact in manifest["artifacts"]
        ],
    }
    print(json.dumps(plan, indent=2))
    if not args.execute:
        print("DRY RUN: no repository was created and no files were uploaded.")
        return 0
    if not args.accept_flux_license:
        raise SystemExit("Refusing upload: pass --accept-flux-license after reviewing LICENSE_NOTICE.md")

    if existing_verification_is_current(manifest, model_root):
        print("Reusing current full-hash release_verification.json")
    else:
        subprocess.run(
            [sys.executable, str(VERIFY_SCRIPT), "--model-root", str(model_root), "--models"],
            check=True,
        )

    from huggingface_hub import HfApi

    api = HfApi()
    if args.create_private:
        api.create_repo(repo_id=args.repo_id, repo_type="model", private=True, exist_ok=False)

    support_paths = [
        ".gitattributes",
        "LICENSE_NOTICE.md",
        "README.md",
        "assets/*",
        "manifests/*",
        "release_manifest.json",
        "release_verification.json",
        "scripts/*",
        "workflows/*",
    ]
    api.upload_folder(
        repo_id=args.repo_id,
        repo_type="model",
        folder_path=str(ROOT),
        path_in_repo=".",
        revision=args.revision,
        allow_patterns=support_paths,
        ignore_patterns=["**/__pycache__/**", "**/*.pyc"],
        commit_message="Add model card, workflows, and verified release evidence",
    )
    if args.support_only:
        print(f"Support-file upload complete: https://huggingface.co/{args.repo_id}")
        return 0
    for artifact in manifest["artifacts"]:
        source = model_root / artifact["source_relative_path"]
        api.upload_file(
            repo_id=args.repo_id,
            repo_type="model",
            path_or_fileobj=str(source),
            path_in_repo=artifact["repo_path"],
            revision=args.revision,
            commit_message=f"Add {artifact['repo_path']}",
        )
    print(f"Upload complete: https://huggingface.co/{args.repo_id}")
    print("If publication is intended, configure gating and review visibility in Hugging Face settings before making the repository public.")
    return 0


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