"""Build reviewable Hugging Face upload folders without changing Hub state.

Run from the source checkout with its tested MLX environment. Existing native
and packed weights are linked, never duplicated or removed. Quantized outputs
are generated only when explicitly selected with --quantize.
"""

from __future__ import annotations

import argparse
import json
import re
import shutil
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

IMAGE_REVISION = "6534ea1c979a5f8cb00c2bd1025d5ca4ee8d8901"
ENCODER_REVISION = "0919fb684102ef115126612f3fbb84d1027e7317"


def write_json(path, value):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2) + "\n")


def component_config(component, catalog):
    data = {
        "format": "h3-mlx-hub-component-v1",
        "component": component,
        "library_name": "mlx",
        "runtime_repository": "devin-lai/MiniMax-H3-Image-mlx",
        "base_model": catalog["source"]["repo_id"],
        "source_revision": catalog["source"]["revision"],
        "manifest": "files.json",
    }
    if component in catalog["variants"]:
        entry = catalog["variants"][component]
        data.update(
            core_format=entry["core_format"],
            bits=entry.get("bits"),
            group_size=entry.get("group_size"),
        )
    elif component == "runtime":
        data["variants"] = sorted(catalog["variants"])
    return data


def link(source, destination):
    source = Path(source).resolve(strict=True)
    destination.parent.mkdir(parents=True, exist_ok=True)
    if destination.is_symlink() and destination.resolve() == source:
        return
    if destination.exists() or destination.is_symlink():
        raise FileExistsError(destination)
    destination.symlink_to(source)


def model_metadata(variant=None):
    import yaml

    data = {
        "library_name": "mlx",
        "pipeline_tag": "text-to-image",
        "license": "other",
        "license_name": "minimax-h3-image-community-license-agreement",
        "license_link": "LICENSE",
        "base_model": "MiniMaxAI/MiniMax-H3-Image",
        "language": ["en"],
        "inference": False,
        "tags": [
            "mlx",
            "apple-silicon",
            "macos",
            "minimax-h3",
            "image-generation",
            "image-editing",
            "image-to-image",
            "safetensors",
        ],
        "extra_gated_heading": "Request access under the MiniMax license",
        "extra_gated_description": "Manual review. Read LICENSE and the territorial, commercial and acceptable-use conditions before requesting access.",
        "extra_gated_fields": {
            "Country of use": "country",
            "I accept the MiniMax-H3-Image Community License Agreement including Exhibit A": "checkbox",
            "I am authorized to receive and use these materials in my territory": "checkbox",
        },
        "extra_gated_eu_disallowed": True,
    }
    if variant in ("q8", "q4"):
        data["base_model_relation"] = "quantized"
        data["tags"].append("quantized")
    return (
        "---\n" + yaml.safe_dump(data, sort_keys=False, allow_unicode=True) + "---\n\n"
    )


def licenses(destination):
    import requests

    for repo, revision, filename in (
        ("MiniMaxAI/MiniMax-H3-Image", IMAGE_REVISION, "LICENSE"),
        ("MiniMaxAI/MiniMax-H3", ENCODER_REVISION, "LICENSE-text-encoder"),
    ):
        response = requests.get(
            f"https://huggingface.co/{repo}/raw/{revision}/LICENSE", timeout=30
        )
        response.raise_for_status()
        (destination / filename).write_text(response.text)
    notice = (
        "MiniMax-H3-Image is licensed under the MiniMax-H3-Image Community License Agreement, "
        "Copyright © 2026 MiniMax. All Rights Reserved.\n\n"
        "Community MLX adaptation and packaging by Devin Lai. This is an independent release.\n"
        "Modified files: MLX affine Q8/Q4 core shards, lossless palette12 shards, "
        "repacked support tensors and evaluated six-step modulation tables. "
        "No training or fine-tuning was performed. See MODIFICATIONS.md.\n"
        "The separate text encoder retains its MiniMax H3 license.\n"
    )
    (destination / "NOTICE").write_text(notice)
    shutil.copy2(ROOT / "docs/MODIFICATIONS.md", destination / "MODIFICATIONS.md")


def manifest(destination, cache):
    files = {}
    for file in sorted(destination.rglob("*")):
        if not file.is_file() or file.name == "files.json" or ".cache" in file.parts:
            continue
        stat = file.stat()
        key = f"{file.resolve()}:{stat.st_size}:{stat.st_mtime_ns}"
        if key not in cache:
            print("Hashing", file.relative_to(destination), flush=True)
            cache[key] = sha256(file)
        files[file.relative_to(destination).as_posix()] = {
            "bytes": stat.st_size,
            "sha256": cache[key],
        }
    write_json(
        destination / "files.json", {"format": "h3-mlx-files-v1", "files": files}
    )
    return sum(
        spec["bytes"] for name, spec in files.items() if name.startswith("model/")
    )


def shared_components(model, packed, destination):
    import mlx.core as mx

    from h3_mlx.config import load_sigmas
    from h3_mlx.dit import DiT
    from h3_mlx.runtime import configure_memory
    from h3_mlx.weights import WeightStore

    configure_memory()
    target = destination / "model"
    (target / "transformer").mkdir(parents=True, exist_ok=True)
    support = target / "transformer/support.safetensors"
    if not support.exists():
        store = WeightStore(model / "transformer")
        names = []
        for key in store.index:
            if key.startswith("blocks."):
                if key.endswith(
                    (
                        "norm1.weight",
                        "norm2.weight",
                        "attn.q_norm.weight",
                        "attn.k_norm.weight",
                    )
                ):
                    names.append(key)
            elif ".adaln_proj." not in key:
                names.append(key)
        print("Exporting native support tensors", flush=True)
        tensors = {key: store.get(key) for key in names}
        mx.save_safetensors(
            str(support),
            tensors,
            metadata={
                "modification": "Repacked native support tensors for the community MLX port by Devin Lai; original tensor values retained",
                "source": f"MiniMaxAI/MiniMax-H3-Image@{IMAGE_REVISION}",
            },
        )
        del tensors, store
        mx.clear_cache()
    link(model / "transformer/config.json", target / "transformer/config.json")
    for name in ("pdd_config.json", "pdd_fused_heads.safetensors", "model_index.json"):
        link(model / name, target / name)
    for file in (model / "vae").iterdir():
        if file.is_file() and not file.name.startswith("."):
            link(file, target / "vae" / file.name)
    modulation = target / "modulation"
    modulation.mkdir(exist_ok=True)
    if not all(
        (modulation / f"{name}.safetensors").exists()
        for name in ("generation", "editing")
    ):
        dit = DiT(model, cache_dir=ROOT / "cache", cache_gb=0)
        for name, edit in (("generation", False), ("editing", True)):
            dit.prepare_modulation(load_sigmas(model), edit=edit)
            data, metadata = mx.load(
                str(ROOT / "cache" / f"adaln-{dit.modulation_key}.safetensors"),
                return_metadata=True,
            )
            metadata["modification"] = (
                "Original timestep projections evaluated for the six released PDD steps by Devin Lai; no interpolation or quantization"
            )
            mx.save_safetensors(
                str(modulation / f"{name}.safetensors"), data, metadata=metadata
            )
            del data
        dit.clear()
    if len(list(packed.glob("blocks.*.safetensors"))) != 50:
        raise ValueError("Expected the validated 50-block packed12 cache")


def stage_main(destination):
    for name in (
        "README.md",
        "README.zh-CN.md",
        "USAGE.md",
        "MODEL_INDEX.md",
        "BENCHMARKS.md",
        "RELEASE.md",
        "LICENSE",
        "NOTICE",
        "pyproject.toml",
        "requirements-macos.lock",
        ".gitignore",
    ):
        shutil.copy2(ROOT / name, destination / name)
    for name in (
        "h3_mlx",
        "tests",
        "docs",
        "assets",
        "examples",
        "scripts",
        "provenance",
        "reports",
        "comparisons",
        "experiments",
        "optimization",
    ):
        shutil.copytree(
            ROOT / name,
            destination / name,
            dirs_exist_ok=True,
            ignore=shutil.ignore_patterns("__pycache__", "*.pyc", ".DS_Store"),
        )
    shutil.copy2(ROOT / "LICENSE-model", destination / "LICENSE-model")
    shutil.copy2(ROOT / "LICENSE-text-encoder", destination / "LICENSE-text-encoder")
    write_json(
        destination / "config.json",
        component_config("runtime", read_json(ROOT / "h3_mlx/model_index.json")),
    )


def stage_gallery(main_folder, destination, cache):
    """Keep generated images and visual comparisons under manual access review."""
    destination.mkdir(parents=True, exist_ok=True)
    licenses(destination)
    write_json(
        destination / "config.json",
        component_config("examples", read_json(ROOT / "h3_mlx/model_index.json")),
    )
    for name in ("examples/README.md", "comparisons/README.md"):
        source = main_folder / name
        target = destination / name
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(source, target)
    for source in (main_folder / "examples").glob("*.json"):
        shutil.copy2(source, destination / "examples" / source.name)
    comparison_index = destination / "comparisons/README.md"
    comparison_index.write_text(
        comparison_index.read_text().replace(
            "../reports/RESULTS.md",
            "https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx/blob/v0.4.0/reports/RESULTS.md",
        )
    )
    moved = set()
    for source in list(main_folder.rglob("*")):
        if source.is_file() and source.suffix.lower() in (".png", ".jpg", ".jpeg", ".html"):
            relative = source.relative_to(main_folder)
            target = destination / relative
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.move(str(source), str(target))
            moved.add(relative.as_posix())

    gallery_url = "https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx-examples"
    pattern = re.compile(r"(!?)\[([^\]]*)\]\(([^)\n]+)\)")
    for document in main_folder.rglob("*.md"):
        text = document.read_text()
        text = re.sub(r"^thumbnail: .*\n", "", text, flags=re.MULTILINE)

        def replace_link(match, base=document.parent):
            target = match.group(3)
            if "://" in target or target.startswith("#"):
                return match.group(0)
            try:
                relative = (base / target).resolve().relative_to(
                    main_folder.resolve()
                ).as_posix()
            except ValueError:
                return match.group(0)
            if relative not in moved:
                return match.group(0)
            return f"[{match.group(2)}]({gallery_url}/blob/v0.4.0/{relative})"

        text = pattern.sub(replace_link, text)
        document.write_text(text)
    (destination / "README.md").write_text(
        model_metadata()
        + "# MiniMax H3 Image MLX · example gallery\n\n"
        "Generated images and visual comparisons from the community MLX port by Devin Lai. "
        "Powered by MiniMax-H3-Image. These images are AI-generated.\n\n"
        "Examples are subject to the MiniMax-H3-Image Community License Agreement, "
        "including its territorial and acceptable-use restrictions. Request access "
        "and review the license before viewing or using the files. The standard grant "
        "excludes the EU, UK, Republic of Korea and US.\n\n"
        "- [Showcase](assets/showcase.jpg)\n"
        "- [Generated examples, prompts and seeds](examples/README.md)\n"
        "- [Visual precision comparisons](comparisons/README.md)\n"
        "- [Runtime and installation](https://huggingface.co/devin-lai/MiniMax-H3-Image-mlx)\n\n"
        "See [LICENSE](LICENSE), [NOTICE](NOTICE) and [modifications](MODIFICATIONS.md). "
        "Access review does not grant additional rights.\n"
    )
    manifest(destination, cache)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True, type=Path)
    parser.add_argument("--packed-cache", required=True, type=Path)
    parser.add_argument("--output", type=Path, default=ROOT / ".release/huggingface")
    parser.add_argument("--quantize", nargs="*", choices=("q8", "q4"), default=[])
    args = parser.parse_args()
    model, packed = args.model.resolve(), args.packed_cache.resolve()
    output = args.output.resolve()
    output.mkdir(parents=True, exist_ok=True)
    catalog = read_json(ROOT / "h3_mlx/model_index.json")
    provenance = read_json(ROOT / "provenance/sources.json")
    if (
        sha256(model / "transformer/config.json")
        != provenance["files"]["models/MiniMax-H3-Image/transformer/config.json"]
    ):
        raise ValueError("Image model configuration differs from the pinned source")
    cache_path = output.parent / "hash-cache.json"
    cache = read_json(cache_path) if cache_path.exists() else {}
    summary = {"status": "prepared", "publication": "not started", "repositories": {}}
    shared = output / "MiniMax-H3-Image-mlx-shared"
    shared.mkdir(exist_ok=True)
    licenses(shared)
    write_json(shared / "config.json", component_config("shared", catalog))
    shared_components(model, packed, shared)
    (ROOT / "LICENSE-model").write_text((shared / "LICENSE").read_text())
    (ROOT / "LICENSE-text-encoder").write_text(
        (shared / "LICENSE-text-encoder").read_text()
    )
    (shared / "README.md").write_text(
        model_metadata() + (ROOT / "docs/cards/shared.md").read_text()
    )
    summary["repositories"][shared.name] = {
        "ready": True,
        "model_bytes": manifest(shared, cache),
    }
    write_json(cache_path, cache)
    for variant, entry in catalog["variants"].items():
        destination = output / entry["repo_id"].split("/")[-1]
        destination.mkdir(exist_ok=True)
        licenses(destination)
        write_json(destination / "config.json", component_config(variant, catalog))
        write_json(
            destination / "variant.json",
            {"variant": variant, **entry, "source": catalog["source"]},
        )
        (destination / "README.md").write_text(
            model_metadata(variant) + (ROOT / f"docs/cards/{variant}.md").read_text()
        )
        if variant == "bf16":
            for file in (model / "transformer").iterdir():
                if file.is_file() and not file.name.startswith("."):
                    link(file, destination / "model/transformer" / file.name)
        elif variant == "packed12":
            for file in sorted(packed.glob("blocks.*.safetensors")):
                link(file, destination / "model/core" / file.name)
        elif variant in args.quantize:
            from h3_mlx.convert import convert

            core = destination / "model/core"
            if not (core / "quantization.json").exists():
                convert(model, core, bits=entry["bits"], group_size=64)
        ready = (destination / "model").is_dir() and (
            entry["core_format"] != "quantized"
            or (destination / "model/core/quantization.json").is_file()
        )
        if ready and entry["core_format"] == "quantized":
            metadata = read_json(destination / "model/core/quantization.json")
            ready = len(metadata["files"]) == 52 and all(
                (destination / "model/core" / file["file"]).is_file()
                and (destination / "model/core" / file["file"]).stat().st_size
                == file["bytes"]
                for file in metadata["files"]
            )
        size = manifest(destination, cache)
        summary["repositories"][destination.name] = {
            "ready": ready,
            "model_bytes": size,
        }
        write_json(cache_path, cache)
    main_folder = output / "MiniMax-H3-Image-mlx"
    main_folder.mkdir(exist_ok=True)
    stage_main(main_folder)
    gallery = output / "MiniMax-H3-Image-mlx-examples"
    stage_gallery(main_folder, gallery, cache)
    summary["repositories"][gallery.name] = {"ready": True}
    write_json(cache_path, cache)
    summary["repositories"][main_folder.name] = {"ready": True}
    write_json(output.parent / "preparation.json", summary)
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()