"""Capture Qwen decoder layer-0 intermediates for direct parity debugging."""

from pathlib import Path


path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
source = path.read_text(encoding="utf-8")
if "import os\n" not in source:
    source = "import os\n" + source
old = (
    "        # Self Attention\n"
    "        residual = x\n"
    "        x = self.input_layernorm(x)\n"
    "        x, present_key_value = self.self_attn(\n"
    "            hidden_states=x,\n"
    "            attention_mask=attention_mask,\n"
    "            freqs_cis=freqs_cis,\n"
    "            optimized_attention=optimized_attention,\n"
    "            past_key_value=past_key_value,\n"
    "        )\n"
    "        x = residual + x\n"
    "\n"
    "        # MLP\n"
    "        residual = x\n"
    "        x = self.post_attention_layernorm(x)\n"
    "        x = self.mlp(x)\n"
    "        x = residual + x\n"
)
new = (
    "        # Self Attention\n"
    "        residual = x\n"
    "        x = self.input_layernorm(x)\n"
    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
    "        if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"qwen0_norm1.pt\"))\n"
    "        attention, present_key_value = self.self_attn(\n"
    "            hidden_states=x,\n"
    "            attention_mask=attention_mask,\n"
    "            freqs_cis=freqs_cis,\n"
    "            optimized_attention=optimized_attention,\n"
    "            past_key_value=past_key_value,\n"
    "        )\n"
    "        if capture_dir: torch.save(attention.detach().cpu(), os.path.join(capture_dir, \"qwen0_attention.pt\"))\n"
    "        x = residual + attention\n"
    "        if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"qwen0_post_attention.pt\"))\n"
    "\n"
    "        # MLP\n"
    "        residual = x\n"
    "        x = self.post_attention_layernorm(x)\n"
    "        if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"qwen0_norm2.pt\"))\n"
    "        mlp = self.mlp(x)\n"
    "        if capture_dir: torch.save(mlp.detach().cpu(), os.path.join(capture_dir, \"qwen0_mlp.pt\"))\n"
    "        x = residual + mlp\n"
    "        if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"qwen0_output.pt\"))\n"
)
if source.count(old) != 1:
    raise RuntimeError("Unable to locate Qwen TransformerBlock.forward.")
source = source.replace(old, new)

old = "        for i, layer in enumerate(self.layers):\n"
new = "        for i, layer in enumerate(self.layers):\n            layer._h3_trace_index = i\n"
if source.count(old) != 1:
    raise RuntimeError("Unable to locate Qwen decoder layer loop.")
path.write_text(source.replace(old, new), encoding="utf-8")
print("Applied Qwen layer-0 sublayer capture patch.")