#!/usr/bin/env python3
import argparse
import json
from pathlib import Path
import re
import subprocess
import tempfile
import xml.etree.ElementTree as ET

ROOT = Path(__file__).resolve().parents[1]


def run(*args, **kwargs):
    return subprocess.run(args, check=True, text=True, capture_output=True, **kwargs).stdout


def main():
    parser = argparse.ArgumentParser(description="Compare Rust layouts with matching Khronos headers")
    parser.add_argument("--headers", type=Path, required=True, help="Vulkan-Headers checkout")
    args = parser.parse_args()
    include = args.headers.resolve() / "include"
    headers = ["vulkan/vulkan.h"] + [
        str(path.relative_to(include)) for path in sorted((include / "vk_video").glob("*.h"))
    ]
    preamble = "#define VK_ENABLE_BETA_EXTENSIONS\n" + "".join(
        f"#include <{header}>\n" for header in headers
    )
    declarations = run("cc", "-E", "-P", "-x", "c", "-I", str(include), "-", input=preamble)
    rust_source = (ROOT / "volcanic/src/lib.rs").read_text()
    generated = set(re.findall(r"pub (?:struct|union|type) (\w+)", rust_source))
    c_checks, rust_checks, rust_aliases = [], [], []
    seen = set()

    def check(label, c_expr, rust_expr):
        c_checks.append(f'printf("{label} %zu\\n", (size_t)({c_expr}));')
        rust_checks.append(f'println!("{label} {{}}", {rust_expr});')

    for registry in ["vk.xml", "video.xml"]:
        for node in ET.parse(ROOT / "spec" / registry).findall("./types/type"):
            name = node.get("name") or node.findtext("name")
            if not name or name in seen or name not in generated:
                continue
            if "vulkan" not in node.get("api", "vulkan").split(","):
                continue
            if node.get("category") not in {"struct", "union", "enum", "handle", "bitmask", "basetype"}:
                continue
            if not re.search(rf"\b{re.escape(name)}\b", declarations):
                continue
            seen.add(name)
            if alias := node.get("alias"):
                rust_aliases.append(f"let _: fn(&mut {name}) -> &mut {alias} = |value| value;")
            check(f"{name}.size", f"sizeof({name})", f"size_of::<{name}>()")
            check(f"{name}.align", f"_Alignof({name})", f"align_of::<{name}>()")
            for member in node.findall("member"):
                if "vulkan" not in member.get("api", "vulkan").split(","):
                    continue
                field = member.findtext("name")
                tail = member.find("name").tail or ""
                if ":" in tail:
                    continue
                ident = "r#type" if field == "type" else field
                check(f"{name}.{field}", f"offsetof({name}, {field})", f"offset_of!({name}, {ident})")

    with tempfile.TemporaryDirectory(prefix="volcanic-abi-") as directory:
        work = Path(directory)
        (work / "layout.c").write_text(
            preamble + "#include <stddef.h>\n#include <stdio.h>\nint main(void) {\n"
            + "\n".join(c_checks) + "\nreturn 0;\n}\n"
        )
        (work / "main.rs").write_text(
            "use volcanic::*;\nuse std::mem::{size_of, align_of, offset_of};\nfn main() {\n"
            + "\n".join(rust_aliases + rust_checks) + "\n}\n"
        )
        (work / "Cargo.toml").write_text(
            '[package]\nname = "volcanic-abi"\nversion = "0.0.0"\nedition = "2024"\n'
            '[[bin]]\nname = "volcanic-abi"\npath = "main.rs"\n'
            '[dependencies]\nvolcanic = { path = ' + json.dumps(str(ROOT / "volcanic"))
            + ', features = ["everything"] }\n'
        )
        run("cc", "-std=c11", "-I", str(include), str(work / "layout.c"), "-o", str(work / "layout"))
        c_layout = run(str(work / "layout")).splitlines()
        rust_layout = run("cargo", "run", "--quiet", "--offline", "--manifest-path", str(work / "Cargo.toml")).splitlines()
        mismatches = [(c, r) for c, r in zip(c_layout, rust_layout, strict=True) if c != r]
        if mismatches:
            for c, rust in mismatches:
                print(f"C: {c}; Rust: {rust}")
            raise SystemExit(f"{len(mismatches)} ABI mismatches")
        print(f"Matched {len(c_checks)} sizes, alignments, and field offsets across {len(seen)} types; verified {len(rust_aliases)} type aliases")


if __name__ == "__main__":
    try:
        main()
    except subprocess.CalledProcessError as error:
        raise SystemExit(error.stderr) from error