"""The Metal type table against the upstream code it was transcribed from.

`gguf_metal/ggml_dispatch.mm` re-implements, for the two ops here, what ggml's Metal backend does in
`ggml_metal_library_get_pipeline_mul_mv` and `ggml_metal_op_mul_mat`: which types have a gemv, how
much threadgroup memory each wants, and which of the two grids it is dispatched on. Those files are
vendored, so `vendor.py --rev` replaces them, and a pin bump can change any of it.

Most of that drift is caught by the compiler -- the nsg/nr0 values are used through upstream's own
`N_SG_*`/`N_R0_*` macros, so a rename or removal fails the build. What is not caught is a changed
`smem` expression or a changed grid rule: both still compile and both silently return wrong numbers.
This parses upstream and fails loudly instead.

No device and no built extension needed, so it runs anywhere.
"""

import re
from pathlib import Path

import pytest


ROOT = Path(__file__).resolve().parent.parent
DEVICE_CPP = ROOT / "vendor/src/ggml-metal/ggml-metal-device.cpp"
OPS_CPP = ROOT / "vendor/src/ggml-metal/ggml-metal-ops.cpp"
DISPATCH_MM = ROOT / "gguf_metal/ggml_dispatch.mm"

pytestmark = pytest.mark.skipif(not DEVICE_CPP.exists(), reason="vendor/ is not checked out")


def upstream_mul_mv_table():
    """{type name: smem expression} from ggml_metal_library_get_pipeline_mul_mv's switch."""
    body = DEVICE_CPP.read_text()
    start = body.index("ggml_metal_library_get_pipeline_mul_mv(")
    end = body.index("ggml_metal_library_get_pipeline_mul_mm_id_map0", start)
    switch = body[start:end]

    table = {}
    for case in re.finditer(
        r"case GGML_TYPE_(\w+):\s*\{(.*?)\}\s*break;", switch, re.S
    ):
        name, arm = case.group(1), case.group(2)
        smem = re.search(r"smem\s*=\s*([^;]+);", arm)
        table[name] = " ".join(smem.group(1).split()) if smem else "0"
    return table


def ours():
    """{type name: smem expression} from our table, keyed the way upstream names its types."""
    body = DISPATCH_MM.read_text()
    start = body.index("const std::unordered_map<int, TypeInfo> &type_table()")
    end = body.index("const TypeInfo *lookup(", start)

    table = {}
    for entry in re.finditer(
        r"\{GGML_(\w+),\s*\{\"[\w]+\",\s*\d+,\s*\d+,\s*N_SG_\w+,\s*N_R0_\w+,\s*(true|false),\s*([^}]+)\}\}",
        body[start:end],
    ):
        name, reduce_flag, smem = entry.group(1), entry.group(2), entry.group(3)
        table[name] = (" ".join(smem.split()).rstrip(","), reduce_flag == "true")
    return table


def normalize(expr: str) -> str:
    """`32 * sizeof(float) * N_R0_Q8_0` and `32*sizeof(float)*N_R0_Q8_0` are the same expression."""
    return expr.replace(" ", "").replace("(float)", "(float)")


def test_every_type_we_claim_still_exists_upstream():
    upstream = upstream_mul_mv_table()
    missing = sorted(set(ours()) - set(upstream))
    assert not missing, (
        f"{missing} are in our table but no longer in upstream's mul_mv switch; the pin moved under "
        "gguf_metal/ggml_dispatch.mm"
    )


def test_threadgroup_memory_matches_upstream():
    upstream = upstream_mul_mv_table()
    mismatched = {
        name: (smem, upstream[name])
        for name, (smem, _) in ours().items()
        if normalize(smem) != normalize(upstream[name])
    }
    assert not mismatched, (
        "upstream changed the threadgroup memory for these types (ours, theirs): "
        f"{mismatched}. Too little is a fault, too much silently costs occupancy."
    )


def test_grid_rule_still_singles_out_q8_0():
    """The one place a type's grid divisor is decided, and the one field we cannot get from a macro.

    Upstream dispatches the dense types and q8_0 on `(ne01 + nr0 - 1)/nr0` and everything else on
    `(ne01 + nr0*nsg - 1)/(nr0*nsg)`. Our `mv_reduce_across_sgs` is exactly that predicate, so if
    upstream ever moves a quantized type across, this catches it.
    """
    body = OPS_CPP.read_text()
    start = body.index("ggml_metal_library_get_pipeline_mul_mv(lib, op)")
    branch = body[start : body.index("return 1;", start)]

    condition = re.search(r"if \(op->src\[0\]->type ==(.*?)\) \{", branch, re.S)
    assert condition, "upstream's mul_mv grid branch is no longer shaped as a type comparison"
    quantized_on_nr0_grid = {
        t for t in re.findall(r"GGML_TYPE_(\w+)", condition.group(1))
    } - {"F32", "F16", "BF16"}

    ours_on_nr0_grid = {name for name, (_, reduce_flag) in ours().items() if reduce_flag}
    assert ours_on_nr0_grid == quantized_on_nr0_grid, (
        f"upstream now puts {quantized_on_nr0_grid} on the undivided grid, we assume "
        f"{ours_on_nr0_grid}; mv_reduce_across_sgs in gguf_metal/ggml_dispatch.mm has to follow"
    )