"""Shared export path for the Qwen LLM recipes: pinned HF snapshot -> ExecuTorch checkpoint
-> export_llm (Vulkan partitioner, which WebGPUPartitioner wraps unchanged) -> .pte.

Graph contract (matches the internal WebGPU LLM artifacts): forward(tokens[1, <=prefill_max]
int64, input_pos[1] int64) -> logits[1, vocab] float32 for the last position, KV cache inside the
program (fp32 buffers, max_context_length positions), 4-bit groupwise weight-only linears
(et_vk.linear_q4gsw) and 4-bit groupwise embedding (et_vk.embedding_q4gsw).
"""

import importlib
import json
import os
import shlex
import shutil
import subprocess
import sys
from pathlib import Path

import _common

SHIMS = Path(__file__).resolve().parent / "_shims"

TOKENIZER_FILES = [
    "tokenizer.json",
    "tokenizer_config.json",
    "vocab.json",
    "merges.txt",
    "generation_config.json",
    "config.json",
    "LICENSE",
]


def export(
    *, name, upstream_repo, revision, model_class, params, converter, max_context, args
):
    from huggingface_hub import snapshot_download

    out = Path(args.out).resolve()
    work = out / "_work"
    work.mkdir(parents=True, exist_ok=True)
    snap = Path(
        snapshot_download(
            upstream_repo,
            revision=revision,
            allow_patterns=["*.safetensors", "*.json", "merges.txt", "LICENSE"],
        )
    )

    sys.path.insert(0, str(SHIMS))
    env = {**os.environ, "PYTHONPATH": os.pathsep.join(filter(None, [str(SHIMS), os.environ.get("PYTHONPATH")]))}
    ckpt = work / "checkpoint.pth"
    if not ckpt.exists():
        importlib.import_module(converter).convert_weights(str(snap), str(ckpt))
    params_path = work / "params.json"
    params_path.write_text(json.dumps(params, indent=2) + "\n")

    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={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,{args.group_size}'",
        "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:
        if (snap / f).exists():
            shutil.copyfile(snap / 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": upstream_repo,
            "revision": revision,
            "official": True,
            "safetensors_sha256": {
                p.name: _common.sha256_file(p)
                for p in sorted(snap.glob("*.safetensors"))
            },
        },
        "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_conversion": f"{converter}.convert_weights(<snapshot>, checkpoint.pth)",
            "import_shim": "recipes/_shims/torchtune (re-exports executorch.examples.models.checkpoint.get_mapped_key; torchtune itself is not installed)",
            "params": params,
            "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": 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 {args.group_size}); output head quantized too",
            "kv_cache": "fp32 buffers in the program; fp16 KV is a runtime option of the WebGPU runtime (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))


def add_args(ap, max_context):
    ap.add_argument("--out", required=True, help="output directory")
    ap.add_argument(
        "--prefill-max", type=int, default=512, help="max tokens per forward call"
    )
    ap.add_argument(
        "--max-context", type=int, default=max_context, help="KV cache length"
    )
    ap.add_argument("--group-size", type=int, default=64)