"""Capture MiniMax Qwen layer-50 output before H3 token refinement."""

from pathlib import Path


path = Path("/opt/ComfyUI/comfy/text_encoders/minimax.py")
source = path.read_text(encoding="utf-8")
if "import os\n" not in source:
    source = source.replace("import math\n", "import math\nimport os\n")
old = (
    "        return super().forward(input_ids, attention_mask=attention_mask, embeds=embeds,\n"
    "                               num_tokens=num_tokens, intermediate_output=intermediate_output,\n"
    "                               final_layer_norm_intermediate=final_layer_norm_intermediate,\n"
    "                               dtype=dtype, embeds_info=embeds_info, **kwargs)\n"
)
new = (
    "        output = super().forward(input_ids, attention_mask=attention_mask, embeds=embeds,\n"
    "                               num_tokens=num_tokens, intermediate_output=intermediate_output,\n"
    "                               final_layer_norm_intermediate=final_layer_norm_intermediate,\n"
    "                               dtype=dtype, embeds_info=embeds_info, **kwargs)\n"
    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
    "        if capture_dir:\n"
    "            os.makedirs(capture_dir, exist_ok=True)\n"
    "            torch.save(input_ids.detach().cpu() if input_ids is not None else torch.empty(0, dtype=torch.long), os.path.join(capture_dir, \"qwen_input_ids.pt\"))\n"
    "            torch.save(embeds.detach().cpu(), os.path.join(capture_dir, \"qwen_input_embeds.pt\"))\n"
    "            torch.save(output[0].detach().cpu(), os.path.join(capture_dir, \"qwen_layer50.pt\"))\n"
    "        return output\n"
)
if source.count(old) != 1:
    raise RuntimeError("Unable to locate MiniMax Qwen forward return.")
path.write_text(source.replace(old, new), encoding="utf-8")
print("Applied MiniMax Qwen layer-50 capture patch.")