"""Publish verified release components after redistribution rights are established.

This command changes Hugging Face state. It never accepts model access requests
or sends announcements. Components may be uploaded one at a time to conserve
scratch space. The main repository is finalized only after every component.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from h3_mlx.hub import read_json, sha256  # noqa: E402

PROJECT = "devin-lai/MiniMax-H3-Image-mlx"
RELEASE = "v0.4.0"
COMPONENTS = ("shared", "bf16", "packed12", "q8", "q4", "examples")


def write_json(path, value):
    path.write_text(json.dumps(value, indent=2) + "\n")


def verify_local(folder):
    manifest = read_json(folder / "files.json")
    if manifest.get("format") != "h3-mlx-files-v1":
        raise ValueError("Unknown file manifest")
    for name, spec in manifest["files"].items():
        from h3_mlx.hub import relative_path

        file = folder / relative_path(name)
        if not file.is_file() or file.stat().st_size != spec["bytes"]:
            raise ValueError(f"Missing or truncated local file: {file}")
        if sha256(file) != spec["sha256"]:
            raise ValueError(f"Local SHA-256 mismatch: {file}")
    quantization = folder / "model/core/quantization.json"
    if quantization.exists():
        files = read_json(quantization)["files"]
        if len(files) != 52:
            raise ValueError("Expected all fifty main blocks and two token refiners")
        for file in files:
            expected = manifest["files"].get("model/core/" + file["file"])
            if expected != {"bytes": file["bytes"], "sha256": file["sha256"]}:
                raise ValueError(f"Quantized export/manifest mismatch: {file['file']}")
    return manifest


def verify_remote(api, repo_id, manifest, *, repo_type="model"):
    from huggingface_hub import hf_hub_download

    info = api.repo_info(repo_id, repo_type=repo_type, files_metadata=True)
    files = {file.rfilename: file for file in info.siblings}
    for name, spec in manifest["files"].items():
        remote = files.get(name)
        if remote is None or remote.size != spec["bytes"]:
            raise ValueError(f"Remote file missing or wrong size: {repo_id}/{name}")
        lfs = getattr(remote, "lfs", None)
        digest = getattr(lfs, "sha256", None) if lfs is not None else None
        if isinstance(lfs, dict):
            digest = lfs.get("sha256")
        if digest is None:
            if remote.size > 10 * 1024**2:
                raise ValueError(
                    f"Cannot verify remote weight digest: {repo_id}/{name}"
                )
            digest = sha256(
                hf_hub_download(repo_id, name, revision=info.sha, repo_type=repo_type)
            )
        if digest != spec["sha256"]:
            raise ValueError(f"Remote SHA-256 mismatch: {repo_id}/{name}")
    return info.sha


def ensure_owned_repository(api, repo_id, folder, *, gated):
    from huggingface_hub import hf_hub_download

    marker = {"project": PROJECT, "release": RELEASE}
    if api.repo_exists(repo_id):
        try:
            existing = read_json(hf_hub_download(repo_id, "release-owner.json"))
        except Exception as error:
            raise RuntimeError(
                f"Existing repository has no matching release marker: {repo_id}"
            ) from error
        if existing != marker:
            raise RuntimeError(f"Refusing to modify unrelated repository: {repo_id}")
    else:
        api.create_repo(repo_id, private=True)
        # Establish ownership immediately so an interrupted upload is resumable.
        api.upload_file(
            repo_id=repo_id,
            path_in_repo="release-owner.json",
            path_or_fileobj=(json.dumps(marker) + "\n").encode(),
            commit_message=f"Initialize {RELEASE} release staging",
        )
    write_json(folder / "release-owner.json", marker)
    if gated:
        api.update_repo_settings(repo_id, gated="manual")


def refresh_small_file(folder, name):
    manifest = read_json(folder / "files.json")
    file = folder / name
    manifest["files"][name] = {"bytes": file.stat().st_size, "sha256": sha256(file)}
    write_json(folder / "files.json", manifest)


def publish_component(api, component, folder):
    repo_id = f"{PROJECT}-{component}"
    if component != "examples" and not (folder / "model").is_dir():
        raise ValueError(f"Model files are not prepared for {component}")
    if (
        component in ("q8", "q4")
        and not (folder / "model/core/quantization.json").is_file()
    ):
        raise ValueError(f"Quantized export incomplete: {component}")
    if (
        component == "packed12"
        and len(list((folder / "model/core").glob("blocks.*.safetensors"))) != 50
    ):
        raise ValueError("Expected all fifty packed blocks")
    card = (
        (folder / "README.md")
        .read_text()
        .replace(
            "**Preparation status:** this model repository is staged locally; publication\n"
            "and immutable commit IDs are pending confirmation of redistribution rights.",
            f"**Release:** {RELEASE}. Access requests are reviewed manually under the retained license.",
        )
        .replace(
            "**Preparation status:** staged locally; publication is pending confirmation\n"
            "of redistribution rights.",
            f"**Release:** {RELEASE}. Access requests are reviewed manually.",
        )
    )
    (folder / "README.md").write_text(card)
    refresh_small_file(folder, "README.md")
    manifest = verify_local(folder)
    ensure_owned_repository(api, repo_id, folder, gated=True)
    api.upload_large_folder(repo_id, folder, repo_type="model", num_workers=4)
    revision = verify_remote(api, repo_id, manifest)
    api.create_tag(repo_id, tag=RELEASE, revision=revision, exist_ok=True)
    if api.model_info(repo_id, revision=RELEASE).sha != revision:
        raise ValueError(f"Existing release tag points at another commit: {repo_id}")
    api.update_repo_settings(repo_id, gated="manual", private=False)
    info = api.model_info(repo_id)
    if info.private or info.gated != "manual":
        raise RuntimeError(f"Unexpected publication/access state: {repo_id}")
    return {
        "repo_id": repo_id,
        "revision": revision,
        "gated": "manual",
        "verified": True,
    }


def finalize_main(api, output, state):
    missing = [name for name in COMPONENTS if not state.get(name, {}).get("verified")]
    if missing:
        raise ValueError(
            f"Publish and verify these components first: {', '.join(missing)}"
        )
    from huggingface_hub import ModelCard

    catalog = read_json(ROOT / "h3_mlx/model_index.json")
    catalog["status"] = "published"
    catalog["shared"]["revision"] = state["shared"]["revision"]
    for name in catalog["variants"]:
        catalog["variants"][name]["revision"] = state[name]["revision"]
    main_folder = output / PROJECT.split("/")[-1]
    write_json(main_folder / "h3_mlx/model_index.json", catalog)
    collection = api.create_collection(
        "MiniMax H3 Image · Apple silicon / MLX",
        namespace="devin-lai",
        private=True,
        description="Local image generation and editing. BF16, lossless packed BF16, Q8 and Q4; selective downloads, shared components and measured comparisons.",
        exists_ok=True,
    )
    readme = (main_folder / "README.md").read_text()
    readme = readme.replace(
        "The current preparation has **not yet been\npublished** while redistribution rights are clarified. The commands below are\nthe prepared release workflow.",
        f"Release **{RELEASE}** is available with manual access review for model components.",
    )
    if f"[Collection]({collection.url})" not in readme:
        readme = readme.replace(
            "[Repository map](MODEL_INDEX.md)",
            f"[Collection]({collection.url}) · [Repository map](MODEL_INDEX.md)",
        )
    for name in ("q8", "q4"):
        files = read_json(output / f"MiniMax-H3-Image-mlx-{name}/files.json")["files"]
        size = (
            sum(
                spec["bytes"]
                for path, spec in files.items()
                if path.startswith("model/")
            )
            / 1024**3
        )
        readme = readme.replace(
            "About 20 GiB²" if name == "q8" else "About 11 GiB²", f"{size:.2f} GiB"
        )
    readme = readme.replace(
        " ² Estimates until the prepared export manifest is finalized.", ""
    )
    readme = readme.replace(
        " ² Q8 is an estimate until its export manifest is finalized.", ""
    )
    (main_folder / "README.md").write_text(readme)
    chinese = (
        (main_folder / "README.zh-CN.md")
        .read_text()
        .replace(
            "**发布状态：目前仅本地准备，尚未上传。**",
            f"**发布版本：{RELEASE}，模型组件采用人工审核访问。**",
        )
        .replace(
            "在确认发布方具有所需权限之前暂停公开发布。", "使用前请确认具备所需权限。"
        )
    )
    (main_folder / "README.zh-CN.md").write_text(chinese)
    release_file = main_folder / "RELEASE.md"
    release_notes = release_file.read_text()
    if "<!-- COLLECTION_LINK -->" in release_notes:
        release_notes = release_notes.replace(
            "<!-- COLLECTION_LINK -->", f"[Browse the collection]({collection.url})."
        )
    elif collection.url not in release_notes:
        release_notes += f"\n[Browse the collection]({collection.url}).\n"
    release_file.write_text(release_notes)
    ModelCard.load(main_folder / "README.md").validate()
    ensure_owned_repository(api, PROJECT, main_folder, gated=False)
    from scripts.prepare_hf_release import manifest as build_manifest

    build_manifest(main_folder, {})
    main_manifest = read_json(main_folder / "files.json")
    commit = api.upload_folder(
        repo_id=PROJECT,
        folder_path=main_folder,
        commit_message=f"Release {RELEASE}: selective MLX model downloads",
    )
    if verify_remote(api, PROJECT, main_manifest) != commit.oid:
        raise RuntimeError("Main repository changed during publication verification")
    api.create_tag(PROJECT, tag=RELEASE, revision=commit.oid, exist_ok=True)
    if api.model_info(PROJECT, revision=RELEASE).sha != commit.oid:
        raise ValueError("Main release tag points at another commit")
    api.update_repo_settings(PROJECT, private=False)
    notes = {
        "main": "Start here: runtime, quick start, model selection, authentic outputs and measured results.",
        "packed12": "Recommended exact storage choice: lossless BF16 packing, not quantization.",
        "bf16": "Unchanged original transformer files: native BF16/FP32 storage control.",
        "q8": "Smaller affine Q8 core; image changes and no established generation-speed advantage.",
        "q4": "Smallest affine Q4 core; stronger detail and lettering changes.",
        "shared": "VAE, native support, PDD heads, configurations and frozen six-step tables; reused across variants.",
        "examples": "Generated examples and visual comparisons, with access review under the retained license.",
    }
    for name in ("main", "packed12", "bf16", "q8", "q4", "shared", "examples"):
        repo_id = PROJECT if name == "main" else state[name]["repo_id"]
        api.add_collection_item(
            collection.slug, repo_id, "model", note=notes[name], exists_ok=True
        )
    api.update_collection_metadata(collection.slug, private=False, position=0)
    write_json(ROOT / "h3_mlx/model_index.json", catalog)
    return {
        "repo_id": PROJECT,
        "revision": commit.oid,
        "collection": collection.url,
        "verified": True,
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=ROOT / ".release/huggingface")
    parser.add_argument("--component", choices=(*COMPONENTS, "main"), required=True)
    parser.add_argument(
        "--rights-confirmed",
        action="store_true",
        help="Publisher has established the required redistribution rights",
    )
    args = parser.parse_args()
    if not args.rights_confirmed:
        parser.error("Establish redistribution rights first; no Hub changes were made")
    from huggingface_hub import HfApi

    api = HfApi()
    if api.whoami()["name"] != "devin-lai":
        raise RuntimeError("Expected the devin-lai publishing account")
    state_path = args.output.resolve().parent / "publication.json"
    state = read_json(state_path) if state_path.exists() else {}
    if args.component == "main":
        state["main"] = finalize_main(api, args.output.resolve(), state)
    else:
        folder = args.output / f"MiniMax-H3-Image-mlx-{args.component}"
        state[args.component] = publish_component(api, args.component, folder)
    write_json(state_path, state)
    print(json.dumps(state[args.component], indent=2))


if __name__ == "__main__":
    main()