"""Shared helpers for the ExecuTorch WebGPU export recipes in this directory.

Each recipe builds an eager model from a pinned upstream revision and calls
`export_vision` (or its own LLM path), which writes:

  <out>/<name>_webgpu.pte    program lowered with WebGPUPartitioner (delegate id "VulkanBackend")
  <out>/<name>_portable.pte  same exported graph without a delegate, for CPU validation only
  <out>/recipe.json          upstream id + revision, versions, shapes, dtype, command line
  <out>/validation.json      eager vs exported-program vs portable .pte (and goldens, if given)
"""

import hashlib
import io
import json
import platform
import shlex
import sys
import urllib.request
import zipfile
from pathlib import Path

import numpy as np
import torch

PROXY_HINT = (
    "set https_proxy if the machine needs a proxy for huggingface.co / pypi.org"
)


def versions():
    import executorch
    import torchao
    from executorch.version import git_version

    out = {
        "python": platform.python_version(),
        "torch": torch.__version__,
        "torchao": torchao.__version__,
        "executorch": getattr(executorch, "__version__", None)
        or __import__("executorch.version").version.__version__,
        "executorch_git": git_version,
    }
    try:
        import transformers

        out["transformers"] = transformers.__version__
    except ImportError:
        pass
    try:
        import huggingface_hub

        out["huggingface_hub"] = huggingface_hub.__version__
    except ImportError:
        pass
    return out


def sha256_file(path):
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def hf_file(repo_id, filename, revision, expected_sha256=None, cache_dir=None):
    from huggingface_hub import hf_hub_download

    path = hf_hub_download(
        repo_id=repo_id, filename=filename, revision=revision, cache_dir=cache_dir
    )
    digest = sha256_file(path)
    if expected_sha256 and digest != expected_sha256:
        sys.exit(
            f"{repo_id}@{revision}:{filename} sha256 {digest} != {expected_sha256}"
        )
    return Path(path), digest


def pypi_wheel_member(url, expected_sha256, member):
    """Read one file out of a PyPI wheel, verifying the wheel's sha256."""
    with urllib.request.urlopen(url, timeout=120) as resp:
        blob = resp.read()
    digest = hashlib.sha256(blob).hexdigest()
    if digest != expected_sha256:
        sys.exit(f"{url}: sha256 {digest} != {expected_sha256} ({PROXY_HINT})")
    with zipfile.ZipFile(io.BytesIO(blob)) as zf:
        return zf.read(member).decode()


def scrub(text, out_dir):
    """Replace machine-specific paths so recipe.json can ship with the model."""
    text = text.replace(str(Path(out_dir).resolve()), "<out>")
    text = text.replace(str(Path(__file__).resolve().parent), "recipes")
    return text.replace(str(Path.home()), "~")


def command_line(out_dir):
    argv = [f"recipes/{Path(sys.argv[0]).name}", *sys.argv[1:]]
    args = []
    for a in argv:
        if a.startswith("/") or a.startswith("."):
            a = scrub(str(Path(a).resolve()), out_dir)
        args.append(a)
    return scrub(" ".join(["python", *(shlex.quote(a) for a in args)]), out_dir)


def lower_webgpu(ep):
    from executorch.backends.webgpu.partitioner import WebGPUPartitioner
    from executorch.exir import to_edge_transform_and_lower

    return to_edge_transform_and_lower(
        ep, partitioner=[WebGPUPartitioner()]
    ).to_executorch()


def lower_portable(ep):
    from executorch.exir import to_edge_transform_and_lower

    return to_edge_transform_and_lower(ep).to_executorch()


def delegate_summary(pte_bytes):
    """Non-delegated kernels and per-delegate op histograms of a serialized program."""
    import collections

    from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
        extract_vk_flatbuffer,
        flatbuffer_to_vk_graph,
    )
    from executorch.exir._serialize._program import deserialize_pte_binary

    pte = deserialize_pte_binary(pte_bytes)
    prog = pte.program if hasattr(pte, "program") else pte
    seg_base = int.from_bytes(pte_bytes[24:32], "little")
    plan = prog.execution_plan[0]
    ops = [f"{o.name}.{o.overload}" for o in plan.operators]
    kernels = collections.Counter()
    for chain in plan.chains:
        for ins in chain.instructions:
            if type(ins.instr_args).__name__ == "KernelCall":
                kernels[ops[ins.instr_args.op_index]] += 1
    delegates = []
    for d in plan.delegates:
        proc = d.processed
        if proc.location.name == "INLINE":
            blob = prog.backend_delegate_data[proc.index].data
        else:
            seg = prog.segments[proc.index]
            blob = pte_bytes[seg_base + seg.offset : seg_base + seg.offset + seg.size]
        g = flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob))
        delegates.append(
            {
                "id": d.id,
                "ops": dict(collections.Counter(c.name for c in g.chain)),
                "inputs": len(g.input_ids),
                "outputs": len(g.output_ids),
            }
        )
    return {"non_delegated_kernels": dict(kernels), "delegates": delegates}


def compare(a, b):
    a = np.asarray(a, dtype=np.float64).ravel()
    b = np.asarray(b, dtype=np.float64).ravel()
    n = min(a.size, b.size)
    a, b = a[:n], b[:n]
    cos = float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30))
    return {
        "max_abs": float(np.max(np.abs(a - b))),
        "mean_abs": float(np.mean(np.abs(a - b))),
        "cosine": cos,
    }


def run_pte(pte_bytes, inputs):
    from executorch.extension.pybindings.portable_lib import (
        _load_for_executorch_from_buffer,
    )

    m = _load_for_executorch_from_buffer(pte_bytes)
    return [
        o.detach().numpy() if isinstance(o, torch.Tensor) else o
        for o in m.forward(inputs)
    ]


def try_run_webgpu_pte_on_cpu(pte_bytes, inputs):
    try:
        return run_pte(pte_bytes, inputs), None
    except Exception as e:  # noqa: BLE001 - the error text is the result
        return None, f"{type(e).__name__}: {str(e).splitlines()[0][:300]}"


def export_vision(
    *,
    name,
    model,
    example_input,
    out_dir,
    upstream,
    extra_recipe,
    golden=None,
    output_index=None,
):
    """Export, lower, write both .pte files, recipe.json, and validation.json.

    `golden` is an optional (input ndarray, output ndarray) pair; when its input
    shape matches `example_input`, all paths are also compared against it.
    `output_index` picks one tensor when the eager model returns a tuple.
    """
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    torch.manual_seed(0)
    model = model.eval()

    class Wrapped(torch.nn.Module):
        def __init__(self, m):
            super().__init__()
            self.m = m

        def forward(self, x):
            y = self.m(x)
            return y[output_index] if output_index is not None else y

    wrapped = Wrapped(model).eval()
    with torch.no_grad():
        ep = torch.export.export(wrapped, (example_input,), strict=False)
        webgpu = lower_webgpu(ep)
        portable = lower_portable(ep)

    webgpu_path = out_dir / f"{name}_webgpu.pte"
    portable_path = out_dir / f"{name}_portable.pte"
    webgpu_path.write_bytes(webgpu.buffer)
    portable_path.write_bytes(portable.buffer)

    summary = delegate_summary(webgpu.buffer)
    out_shape = list(ep.module()(example_input).shape)
    recipe = {
        "model": name,
        "upstream": upstream,
        "versions": versions(),
        "export": {
            "api": "torch.export.export(strict=False) -> to_edge_transform_and_lower(partitioner=[WebGPUPartitioner()]) -> to_executorch()",
            "partitioner": "executorch.backends.webgpu.partitioner.WebGPUPartitioner() (default compile options; wraps VulkanPartitioner)",
            "inputs": [
                {
                    "shape": list(example_input.shape),
                    "dtype": "float32",
                    "layout": "NCHW",
                }
            ],
            "outputs": [{"shape": out_shape, "dtype": "float32"}],
            "dtype": "float32",
            "quantization": "none (fp32)",
            **extra_recipe,
        },
        "files": {
            webgpu_path.name: {
                "bytes": webgpu_path.stat().st_size,
                "sha256": sha256_file(webgpu_path),
            },
            portable_path.name: {
                "bytes": portable_path.stat().st_size,
                "sha256": sha256_file(portable_path),
                "note": "CPU-only validation artifact (portable kernels, no delegate); not for release",
            },
        },
        "graph": summary,
        "command": command_line(out_dir),
    }
    (out_dir / "recipe.json").write_text(json.dumps(recipe, indent=2) + "\n")

    # Validation on CPU.
    gen = torch.Generator().manual_seed(1234)
    rand_input = torch.rand(example_input.shape, generator=gen)
    cases = {"random": rand_input}
    if golden is not None and tuple(golden[0].shape) == tuple(example_input.shape):
        cases["golden"] = torch.from_numpy(golden[0].copy())
    report = {
        "cpu_runtime": "executorch pybindings (portable_lib) from the pip wheel",
        "cases": {},
    }
    webgpu_cpu_error = None
    for label, x in cases.items():
        with torch.no_grad():
            eager = wrapped(x).numpy()
            exported = ep.module()(x).numpy()
        port = run_pte(portable.buffer, (x,))[0]
        res = {
            "exported_vs_eager": compare(exported, eager),
            "portable_pte_vs_eager": compare(port, eager),
        }
        wg, err = try_run_webgpu_pte_on_cpu(webgpu.buffer, (x,))
        if wg is not None:
            res["webgpu_pte_vs_eager"] = compare(wg[0], eager)
        else:
            webgpu_cpu_error = err
        if label == "golden":
            ref = golden[1]
            res["internal_golden_vs_eager"] = compare(eager, ref)
            res["internal_golden_vs_portable_pte"] = compare(port, ref)
            res["golden_note"] = (
                "reference output from the internal goldens.manifest.json"
            )
        report["cases"][label] = res
    report["webgpu_pte_on_cpu"] = (
        "ran"
        if webgpu_cpu_error is None
        else f"not runnable on CPU pybindings: {webgpu_cpu_error}"
    )
    (out_dir / "validation.json").write_text(json.dumps(report, indent=2) + "\n")
    print(
        json.dumps(
            {"recipe": recipe["files"], "graph": summary, "validation": report},
            indent=2,
        )
    )
    return ep, webgpu, portable


def load_golden(manifest_path, key):
    man = json.loads(Path(manifest_path).read_text())[key]["tensors"]
    base = Path(manifest_path).parent

    def ld(t):
        return np.fromfile(base / t["file"], dtype=np.float32).reshape(t["shape"])

    inp = man.get("input")
    return ld(inp), ld(man["output"])