"""Add only block-0 sublayer captures to the minimal H3 reference hook."""

from pathlib import Path


model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
source = model.read_text(encoding="utf-8")

old = "        for i, block in enumerate(self.blocks):\n            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
new = "        for i, block in enumerate(self.blocks):\n            block._h3_capture_index = i\n            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
if source.count(old) != 1:
    raise RuntimeError("Unable to locate pristine H3 block loop.")
source = source.replace(old, new)

old = (
    "        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)\n"
    "        h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)\n"
    "        x = _mod_gate(x, gate_msa, self.attn(h, rope_freqs=rope_freqs, transformer_options=transformer_options), mod_segments)\n"
    "        h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)\n"
    "        return _mod_gate(x, gate_mlp, self.mlp(h), mod_segments)\n"
)
new = (
    "        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)\n"
    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_capture_index\", -1) == 0 and H3_CAPTURE_ACTIVE else None\n"
    "        h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)\n"
    "        if capture_dir: torch.save(h.detach().cpu(), os.path.join(capture_dir, \"block0_norm1.pt\"))\n"
    "        attention = self.attn(h, rope_freqs=rope_freqs, transformer_options=transformer_options)\n"
    "        if capture_dir: torch.save(attention.detach().cpu(), os.path.join(capture_dir, \"block0_attention.pt\"))\n"
    "        x = _mod_gate(x, gate_msa, attention, mod_segments)\n"
    "        if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"block0_post_attention.pt\"))\n"
    "        h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)\n"
    "        if capture_dir: torch.save(h.detach().cpu(), os.path.join(capture_dir, \"block0_norm2.pt\"))\n"
    "        mlp = self.mlp(h)\n"
    "        if capture_dir: torch.save(mlp.detach().cpu(), os.path.join(capture_dir, \"block0_mlp.pt\"))\n"
    "        x = _mod_gate(x, gate_mlp, mlp, mod_segments)\n"
    "        if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"block0_post_mlp.pt\"))\n"
    "        return x\n"
)
if source.count(old) != 1:
    raise RuntimeError("Unable to locate pristine H3 DiTBlock.forward.")
model.write_text(source.replace(old, new), encoding="utf-8")
print("Applied minimal H3 block-0 sublayer capture patch.")