#!/usr/bin/env python3
"""Patch _allocate_kv_cache_tensors to print the layer name mismatch."""
import sys

def patch(path):
    with open(path, 'r') as f:
        content = f.read()

    if "CLAWMINE_DEBUG_LAYERS" in content:
        print("Already patched, skipping")
        return

    old = """        assert layer_names == set(kv_cache_raw_tensors.keys()), (
            "Some layers are not correctly initialized"
        )"""

    new = """        # CLAWMINE_DEBUG_LAYERS: print mismatch instead of asserting
        missing = layer_names - set(kv_cache_raw_tensors.keys())
        extra = set(kv_cache_raw_tensors.keys()) - layer_names
        if missing or extra:
            print(f"CLAWMINE DEBUG: missing layers ({len(missing)}): {sorted(missing)[:20]}")
            print(f"CLAWMINE DEBUG: extra layers ({len(extra)}): {sorted(extra)[:20]}")
            print(f"CLAWMINE DEBUG: expected ({len(layer_names)}), got ({len(kv_cache_raw_tensors.keys())})")
        assert layer_names == set(kv_cache_raw_tensors.keys()), (
            "Some layers are not correctly initialized"
        )"""

    if old not in content:
        print("ERROR: Could not find the code to patch")
        sys.exit(1)

    content = content.replace(old, new)

    with open(path, 'w') as f:
        f.write(content)
    print("Patched gpu_model_runner.py for debug layer names")

if __name__ == "__main__":
    patch(sys.argv[1])