#!/usr/bin/env python3
"""Inspection helpers for the SAM 2 recipe (sam2-hiera-tiny.py). CPU only, read-only on inputs.

_sam2.py ops A.pte [B.pte]          every delegated op call with its argument types/shapes; with B,
                                    the op-signature set difference (ops/arity/ranks in A not in B)
_sam2.py consts A.pte B.pte         how many delegate constants of A are byte-identical in B
_sam2.py io A.pte                   method names, program inputs/outputs, non-delegated kernels
"""

import collections
import hashlib
import json
import sys


def _program(data):
    from executorch.exir._serialize._program import deserialize_pte_binary

    pte = deserialize_pte_binary(data)
    return pte, (pte.program if hasattr(pte, "program") else pte)


def _blobs(data):
    pte, prog = _program(data)
    seg_base = int.from_bytes(data[24:32], "little")
    out = []
    for plan in prog.execution_plan:
        for d in plan.delegates:
            proc = d.processed
            if proc.location.name == "INLINE":
                out.append(
                    (
                        plan.name,
                        d.id,
                        bytes(prog.backend_delegate_data[proc.index].data),
                    )
                )
            else:
                seg = prog.segments[proc.index]
                out.append(
                    (
                        plan.name,
                        d.id,
                        data[seg_base + seg.offset : seg_base + seg.offset + seg.size],
                    )
                )
    nd = getattr(pte, "named_data", None)
    resolve = (
        (
            lambda k: bytes(nd.buffers[nd.pte_data[k].buffer_index])
            if k in nd.pte_data
            else None
        )
        if nd
        else (lambda k: None)
    )
    return out, resolve


def vk_graph(blob):
    from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
        extract_vk_flatbuffer,
        flatbuffer_to_vk_graph,
    )

    return flatbuffer_to_vk_graph(extract_vk_flatbuffer(blob))


def describe(g, vid):
    v = g.values[vid].value
    n = type(v).__name__
    if n == "VkTensor":
        kind = "const" if v.constant_id >= 0 else "tensor"
        return f"{kind}{list(v.dims)}:{v.datatype.name if hasattr(v.datatype, 'name') else v.datatype}"
    if n == "Int":
        return f"Int={v.int_val}"
    if n == "Double":
        return f"Double={v.double_val:g}"
    if n == "Bool":
        return f"Bool={v.bool_val}"
    if n in ("IntList", "DoubleList", "BoolList"):
        return f"{n}={list(v.items)}"
    if n == "ValueList":
        return "ValueList[" + ", ".join(describe(g, i) for i in v.items) + "]"
    if n == "String":
        return f"String={v.string_val!r}"
    return n


def ops(path):
    data = open(path, "rb").read()
    blobs, _ = _blobs(data)
    calls = []
    for plan, did, blob in blobs:
        g = vk_graph(blob)
        for c in g.chain:
            calls.append(
                {"plan": plan, "op": c.name, "args": [describe(g, a) for a in c.args]}
            )
    return calls


def signature(call):
    """Op name + arity + per-argument kind/rank (shape-agnostic) for set comparisons."""
    parts = []
    for a in call["args"]:
        if a.startswith(("tensor[", "const[")):
            dims = json.loads(a[a.index("[") : a.index("]") + 1])
            parts.append(f"{a.split('[')[0]}{len(dims)}d:{a.rsplit(':', 1)[1]}")
        else:
            parts.append(a.split("=")[0])
    return f"{call['op']}({', '.join(parts)})"


def consts(path):
    from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
        VulkanDelegateHeader,
    )

    data = open(path, "rb").read()
    blobs, resolve = _blobs(data)
    out = []
    for plan, did, blob in blobs:
        h = VulkanDelegateHeader.from_bytes(
            blob[: VulkanDelegateHeader.EXPECTED_LENGTH]
        )
        g = vk_graph(blob)
        raw = blob[h.bytes_offset : h.bytes_offset + h.bytes_size]
        for v in g.values:
            t = v.value
            if type(t).__name__ != "VkTensor" or t.constant_id < 0:
                continue
            c = g.constants[t.constant_id]
            b = (
                resolve(c.named_key)
                if c.named_key
                else raw[c.offset : c.offset + c.length]
            )
            out.append(
                {
                    "dims": list(t.dims),
                    "sha": hashlib.sha256(b).hexdigest(),
                    "bytes": len(b),
                }
            )
    return out


def io(path):
    data = open(path, "rb").read()
    _, prog = _program(data)
    res = {"file": path, "bytes": len(data), "methods": {}}
    for plan in prog.execution_plan:

        def d(i):
            v = plan.values[i].val
            return (
                f"{list(v.sizes)}:{v.scalar_type.name}"
                if hasattr(v, "sizes")
                else type(v).__name__
            )

        names = [f"{o.name}.{o.overload}" for o in plan.operators]
        kernels = collections.Counter(
            names[ins.instr_args.op_index]
            for ch in plan.chains
            for ins in ch.instructions
            if type(ins.instr_args).__name__ == "KernelCall"
        )
        res["methods"][plan.name] = {
            "inputs": [d(i) for i in plan.inputs],
            "outputs": [d(i) for i in plan.outputs],
            "delegates": [x.id for x in plan.delegates],
            "non_delegated_kernels": dict(kernels),
        }
    return res


OFFICIAL = {
    "sam2.1": ("facebook/sam2.1-hiera-tiny", "de431c4043854a71d8101e17995dfe596bf101a5", "sam2.1_hiera_tiny.pt",
               "7402e0d864fa82708a20fbd15bc84245c2f26dff0eb43a4b5b93452deb34be69"),
}


def verify_official(weights="sam2.1"):
    import importlib.util
    import pathlib
    import tempfile

    import torch
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file

    here = pathlib.Path(__file__).resolve().parent
    spec = importlib.util.spec_from_file_location("sam2_recipe", here / "sam2-hiera-tiny.py")
    recipe = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(recipe)
    snap, digest = recipe.snapshot(weights)
    hf = load_file(str(snap / "model.safetensors"))
    repo, rev, name, want = OFFICIAL[weights]
    with tempfile.TemporaryDirectory() as tmp:
        p = hf_hub_download(repo, name, revision=rev, local_dir=tmp)
        got = hashlib.sha256(open(p, "rb").read()).hexdigest()
        if got != want:
            sys.exit(f"{name}: sha256 {got} != {want}")
        orig = torch.load(p, map_location="cpu", weights_only=True)["model"]
    h = lambda t: hashlib.sha256(t.detach().contiguous().numpy().tobytes()).hexdigest()  # noqa: E731
    pool = {h(v): k for k, v in orig.items()}
    matched, unmatched = 0, []
    for k, v in hf.items():
        if h(v) in pool:
            matched += 1
        else:
            unmatched.append(k)
    notes = {}
    if "prompt_encoder.point_embed.weight" in unmatched:
        pe = hf["prompt_encoder.point_embed.weight"]
        if all(torch.equal(pe[i], orig[f"sam_prompt_encoder.point_embeddings.{i}.weight"][0]) for i in range(pe.shape[0])):
            unmatched.remove("prompt_encoder.point_embed.weight")
            matched += 1
            notes["prompt_encoder.point_embed.weight"] = "rows == sam_prompt_encoder.point_embeddings.{0..3}.weight"
    return {"hf": f"{repo}@{rev}/model.safetensors", "hf_sha256": digest, "official": f"{repo}@{rev}/{name}",
            "official_sha256": want, "hf_tensors": len(hf), "official_tensors": len(orig),
            "hf_tensors_bit_identical_to_official": matched, "unmatched": unmatched, "notes": notes}


def main():
    cmd, args = sys.argv[1], sys.argv[2:]
    if cmd == "io":
        print(json.dumps(io(args[0]), indent=1))
    elif cmd == "ops":
        a = ops(args[0])
        if len(args) == 1:
            for c in a:
                print(c["op"], ", ".join(c["args"]))
            return
        b = ops(args[1])
        sa, sb = (
            collections.Counter(map(signature, a)),
            collections.Counter(map(signature, b)),
        )
        print(
            json.dumps(
                {
                    "ops_count": [sum(sa.values()), sum(sb.values())],
                    "op_types_only_in_a": sorted(
                        {c["op"] for c in a} - {c["op"] for c in b}
                    ),
                    "op_types_only_in_b": sorted(
                        {c["op"] for c in b} - {c["op"] for c in a}
                    ),
                    "signatures_only_in_a": {
                        k: v for k, v in sorted(sa.items()) if k not in sb
                    },
                    "signatures_only_in_b": {
                        k: v for k, v in sorted(sb.items()) if k not in sa
                    },
                },
                indent=1,
            )
        )
    elif cmd == "verify-official":
        print(json.dumps(verify_official(*(args or ["sam2.1"])), indent=1))
    elif cmd == "consts":
        a, b = consts(args[0]), consts(args[1])
        hb, ha = (
            collections.Counter(c["sha"] for c in b),
            collections.Counter(c["sha"] for c in a),
        )
        print(
            json.dumps(
                {
                    "a_constants": len(a),
                    "b_constants": len(b),
                    "a_found_in_b": sum(c["sha"] in hb for c in a),
                    "b_found_in_a": sum(c["sha"] in ha for c in b),
                    "a_missing": [c["dims"] for c in a if c["sha"] not in hb],
                    "b_missing": [c["dims"] for c in b if c["sha"] not in ha],
                },
                indent=1,
            )
        )
    else:
        sys.exit(__doc__)


if __name__ == "__main__":
    main()