#!/usr/bin/env python3
"""Llama 3.2 1B Instruct -> ExecuTorch WebGPU (4-bit groupwise weights, KV cache in program).

Same export path as the Qwen recipes (_llm.py): export_llm with the Vulkan partitioner (which
WebGPUPartitioner wraps unchanged), qmode=4w group 64, 4-bit embedding, fp32 activations and fp32
KV buffers, prefill up to 512 tokens, 8960-token context. Program contract (same as the previous
artifact): forward(tokens[1, 1..512] int64, input_pos[1] int64) -> logits[1, 128256] float32 for
the last position.

Weights: the Meta-format checkpoint `original/consolidated.00.pth` + `original/params.json` of
meta-llama/Llama-3.2-1B-Instruct at a pinned revision (Llama 3.2 Community License). The upstream
repo is gated. `--source auto` (default) downloads from it when the token has access, otherwise
from a public mirror. Either way every file is checked against the hashes Hugging Face publishes
for the *official* repo at the pinned revision (LFS sha256 for large files, git blob sha1 for small
ones; both are visible without access), so the bytes used are the official bytes.
"""

import argparse
import hashlib
import json
import os
import shlex
import shutil
import subprocess
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import _common  # noqa: E402
import _llm  # noqa: E402

OFFICIAL = {
    "repo": "meta-llama/Llama-3.2-1B-Instruct",
    "revision": "9213176726f574b556790deb65791e0c5aa438b6",
}
# Public mirror whose files are byte-identical to OFFICIAL at the pinned revision (checked below).
MIRROR = {
    "repo": "LIF1014/ptdbench-Llama-3.2-1B-Instruct",
    "revision": "828e5f37c833a02415416260d5a93780fd41d10e",
}
# Hashes Hugging Face reports for OFFICIAL@revision (HfApi.model_info(files_metadata=True)).
# lfs: sha256 of the file; blob: git blob sha1 (sha1(b"blob <len>\0" + content)).
FILES = {
    "original/consolidated.00.pth": {
        "lfs": "fc17d497df5e4175b3a8acb4f5865b26f7fc1b009b25bef814b95fde10e8a1f3"
    },
    "original/params.json": {"blob": "9cd8dbdf2dc6f4d8abb60bdb5ce64f4bec2fdfd9"},
    "tokenizer.json": {"blob": "5cc5f00a5b203e90a27a3bd60d1ec393b07971e8"},
    "tokenizer_config.json": {"blob": "4ff488a165e900e5129cda7c20ab32d568d2a475"},
    "special_tokens_map.json": {"blob": "02ee80b6196926a5ad790a004d9efd6ab1ba6542"},
    "config.json": {"blob": "3e3aaf51a035cb5092d9f6827a0dc074657ba88c"},
    "generation_config.json": {"blob": "75ae08310d6d23df373ee2644b497192b3cce6d8"},
    "LICENSE.txt": {"blob": "085b47c1575cb889b7024030e60b78f54f0b8c9e"},
    "USE_POLICY.md": {"blob": "ac3c5f21b9779e3da0677d6d3c587778fe3a331e"},
}
# Shipped next to the .pte for the playground's tokenizer (transformers.js) and license notice.
TOKENIZER_FILES = [
    "tokenizer.json",
    "tokenizer_config.json",
    "special_tokens_map.json",
    "config.json",
    "generation_config.json",
    "LICENSE.txt",
    "USE_POLICY.md",
]
NAME = "llama3_2_1b_instruct_q4gsw"
MODEL_CLASS = "llama3_2"


def git_blob_sha1(path):
    data = Path(path).read_bytes()
    return hashlib.sha1(b"blob %d\0" % len(data) + data).hexdigest()


def check(path, want):
    if "lfs" in want:
        got = _common.sha256_file(path)
        return got == want["lfs"], {"sha256": got}
    got = git_blob_sha1(path)
    return got == want["blob"], {"git_blob_sha1": got}


def official_metadata():
    """Cross-check the pinned hashes against what the Hub reports now (no access needed)."""
    from huggingface_hub import HfApi

    info = HfApi().model_info(
        OFFICIAL["repo"], revision=OFFICIAL["revision"], files_metadata=True
    )
    live = {s.rfilename: s for s in info.siblings}
    for f, want in FILES.items():
        s = live[f]
        if "lfs" in want and (s.lfs is None or s.lfs.sha256 != want["lfs"]):
            sys.exit(f"official {f}: Hub reports {s.lfs} != pinned {want}")
        if "blob" in want and s.blob_id != want["blob"]:
            sys.exit(f"official {f}: Hub reports blob {s.blob_id} != pinned {want}")
    return info.sha


def fetch(source):
    from huggingface_hub import hf_hub_download
    from huggingface_hub.errors import GatedRepoError

    def get(repo):
        return {
            f: Path(hf_hub_download(repo["repo"], f, revision=repo["revision"]))
            for f in FILES
        }

    used = None
    if source in ("auto", "official"):
        try:
            paths, used = get(OFFICIAL), OFFICIAL
        except GatedRepoError as e:
            if source == "official":
                raise
            print(
                f"official repo not accessible ({str(e).splitlines()[-1][:120]}); using mirror",
                flush=True,
            )
    if used is None:
        paths, used = get(MIRROR), MIRROR
    verified = {}
    for f, p in paths.items():
        ok, got = check(p, FILES[f])
        if not ok:
            sys.exit(
                f"{used['repo']}@{used['revision']}:{f} {got} does not match official {FILES[f]}"
            )
        verified[f] = got
    return paths, used, verified


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    _llm.add_args(ap, max_context=8960)
    ap.add_argument(
        "--embedding-group-size",
        type=int,
        default=None,
        help="group size of the 4-bit embedding (default: --group-size; the previous artifact used 32)",
    )
    ap.add_argument("--source", choices=["auto", "official", "mirror"], default="auto")
    args = ap.parse_args()

    out = Path(args.out).resolve()
    work = out / "_work"
    work.mkdir(parents=True, exist_ok=True)
    official_sha = official_metadata()
    paths, used, verified = fetch(args.source)
    params = json.loads(paths["original/params.json"].read_text())

    sys.path.insert(0, str(_llm.SHIMS))
    env = {
        **os.environ,
        "PYTHONPATH": os.pathsep.join(
            filter(None, [str(_llm.SHIMS), os.environ.get("PYTHONPATH")])
        ),
    }
    ckpt = work / "checkpoint.pth"
    if ckpt.is_symlink() or ckpt.exists():
        ckpt.unlink()
    ckpt.symlink_to(paths["original/consolidated.00.pth"])
    params_path = work / "params.json"
    shutil.copyfile(paths["original/params.json"], params_path)

    emb_group = args.embedding_group_size or args.group_size
    output_name = f"{NAME}.pte"
    overrides = [
        f"base.model_class={MODEL_CLASS}",
        f"base.checkpoint={ckpt}",
        f"base.params={params_path}",
        "model.use_kv_cache=true",
        "model.use_sdpa_with_kv_cache=true",
        "model.enable_dynamic_shape=true",
        "model.dtype_override=fp32",
        f"export.max_seq_length={args.prefill_max + 1}",
        f"export.max_context_length={args.max_context}",
        f"export.output_dir={out}",
        f"export.output_name={out / output_name}",
        "quantization.qmode=4w",
        f"quantization.group_size={args.group_size}",
        f"quantization.embedding_quantize='4,{emb_group}'",
        "backend.vulkan.enabled=true",
    ]
    cmd = [
        sys.executable,
        "-m",
        "executorch.extension.llm.export.export_llm",
        *overrides,
    ]
    print("+", " ".join(shlex.quote(c) for c in cmd), flush=True)
    subprocess.run(cmd, check=True, env=env)

    pte = out / output_name
    tok_dir = out / "tokenizer"
    tok_dir.mkdir(exist_ok=True)
    tokenizer_files = {}
    for f in TOKENIZER_FILES:
        shutil.copyfile(paths[f], tok_dir / f)
        tokenizer_files[f] = _common.sha256_file(tok_dir / f)
    shutil.rmtree(work)

    data = pte.read_bytes()
    recipe = {
        "model": NAME,
        "upstream": {
            "model": OFFICIAL["repo"],
            "revision": OFFICIAL["revision"],
            "official": True,
            "license": "Llama 3.2 Community License (tokenizer/LICENSE.txt, tokenizer/USE_POLICY.md); Built with Llama",
            "files": {f: {**FILES[f], **verified[f]} for f in FILES},
            "hub_revision_checked": official_sha,
            "downloaded_from": {
                **used,
                "note": "all files verified byte-identical to the official hashes above",
            },
        },
        "versions": _common.versions(),
        "export": {
            "api": "python -m executorch.extension.llm.export.export_llm (hydra overrides below)",
            "overrides": [
                _common.scrub(o, out).replace(str(work.resolve()), "<work>")
                for o in overrides
            ],
            "checkpoint": "original/consolidated.00.pth (Meta format, bf16; no conversion), params = original/params.json",
            "params": params,
            "rope": "llama3 scaled RoPE, rope_scale_factor=32 (set by ExecuTorch for model_class llama3_2 when use_scaled_rope)",
            "partitioner": "VulkanPartitioner via backend.vulkan.enabled (WebGPUPartitioner is a pass-through wrapper of it; delegate id VulkanBackend)",
            "inputs": [
                {
                    "name": "tokens",
                    "shape": [1, f"1..{args.prefill_max}"],
                    "dtype": "int64",
                },
                {"name": "input_pos", "shape": [1], "dtype": "int64"},
            ],
            "outputs": [
                {
                    "name": "logits",
                    "shape": [1, params["vocab_size"]],
                    "dtype": "float32",
                    "note": "last position only",
                }
            ],
            "prefill_max": args.prefill_max,
            "max_seq_len_metadata": args.prefill_max + 1,
            "max_context_length": args.max_context,
            "dtype": "float32 activations",
            "quantization": f"4-bit groupwise weight-only linears (torchao IntxWeightOnlyConfig int4, group {args.group_size}) and 4-bit groupwise embedding (group {emb_group}); output head quantized too",
            "kv_cache": "fp32 buffers in the program; fp16 KV is a runtime option of the WebGPU runtime (enable_f16_kv_cache / et_set_kv_f16), not an export flag",
        },
        "files": {
            pte.name: {"bytes": len(data), "sha256": _common.sha256_file(pte)},
            **{f"tokenizer/{k}": {"sha256": v} for k, v in tokenizer_files.items()},
        },
        "graph": _common.delegate_summary(data),
        "command": _common.command_line(out),
    }
    (out / "recipe.json").write_text(json.dumps(recipe, indent=2) + "\n")
    print(json.dumps({"files": recipe["files"], "graph": recipe["graph"]}, indent=2))


if __name__ == "__main__":
    main()