"""Capture Comfy Qwen layer-0 projection boundaries for direct parity checks."""

from pathlib import Path


path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
source = path.read_text(encoding="utf-8")
old = "        xq = self.q_proj(hidden_states)\n        xk = self.k_proj(hidden_states)\n        xv = self.v_proj(hidden_states)\n"
new = old + (
    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
    "        if capture_dir:\n"
    "            torch.save({\"q\": xq.detach().cpu(), \"k\": xk.detach().cpu(), \"v\": xv.detach().cpu()}, os.path.join(capture_dir, \"qwen0_qkv.pt\"))\n"
)
if source.count(old) != 1:
    raise RuntimeError("Unable to locate Qwen QKV projections.")
source = source.replace(old, new)

old = "        return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))\n"
new = (
    "        gate = self.gate_proj(x)\n"
    "        up = self.up_proj(x)\n"
    "        activated = self.activation(gate) * up\n"
    "        output = self.down_proj(activated)\n"
    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
    "        if capture_dir:\n"
    "            torch.save({\"gate\": gate.detach().cpu(), \"up\": up.detach().cpu(), \"activated\": activated.detach().cpu(), \"down\": output.detach().cpu()}, os.path.join(capture_dir, \"qwen0_mlp_projections.pt\"))\n"
    "        return output\n"
)
if source.count(old) != 1:
    raise RuntimeError("Unable to locate Qwen MLP projections.")
path.write_text(source.replace(old, new), encoding="utf-8")
print("Applied Qwen layer-0 projection capture patch.")