#!/usr/bin/env python3
"""
PuLID-Flux ComfyUI v0.25 Compatibility Patch
=============================================

PuLID's forward_orig in pulidflux.py doesn't pass attn_mask and
transformer_options to its internal block forward calls. ComfyUI v0.25
needs transformer_options for fp8 quantization patches. Without this fix,
generation crashes (OOM / exit 137) within seconds.

Usage (on the HF Space, via SSH):
    # Upload via base64 pipe (SCP doesn't work on HF Spaces)
    cat patch_pulid_v025.py | base64 | ssh SPACE@ssh.hf.space \\
        "base64 -d > /tmp/patch_pulid.py && python3 /tmp/patch_pulid.py"

    # Then restart ComfyUI to reload the patched module:
    ssh SPACE@ssh.hf.space "pkill -f 'python3 main.py'; cd /ComfyUI && python3 main.py --listen 0.0.0.0 --port 8188 --output-directory /tmp/output &"

IMPORTANT: This patch is LIVE-FILESYSTEM ONLY. It reverts on factory rebuild.
Re-run after every container rebuild, or bake into Dockerfile/start.sh.

Verified: ComfyUI v0.25.0, PyTorch 2.5.1+cu121, a10g-large (46GB RAM).
"""

import re
import sys

PULID_FILE = "/home/ubuntu/ComfyUI/custom_nodes/ComfyUI-PuLID-Flux/pulidflux.py"

def patch():
    with open(PULID_FILE, "r") as f:
        content = f.read()

    original = content
    changes = []

    # 1. Extract attn_mask and transformer_options from kwargs at the top of forward_orig
    #    Look for the kwargs extraction pattern after **kwargs in the signature
    old_extract = "attn_mask = kwargs.get('attn_mask')"
    if old_extract not in content:
        # Add extraction if not present
        # Find the kwargs parameter in forward_orig signature
        content = re.sub(
            r"(def forward_orig\(self.*?\*\*kwargs,?\s*\):)",
            r"\1\n        attn_mask = kwargs.get('attn_mask')\n        transformer_options = kwargs.get('transformer_options', {})",
            content,
            count=1,
        )
        changes.append("Added attn_mask/transformer_options extraction from kwargs")
    else:
        # Ensure transformer_options is also extracted
        if "transformer_options = kwargs.get('transformer_options'" not in content:
            content = content.replace(
                old_extract,
                old_extract + "\n        transformer_options = kwargs.get('transformer_options', {})",
            )
            changes.append("Added transformer_options extraction")

    # 2. Patch double-block calls to pass attn_mask and transformer_options
    #    Pattern: img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
    old_double = "img, txt = block(img=img, txt=txt, vec=vec, pe=pe)"
    new_double = "img, txt = block(img=img, txt=txt, vec=vec, pe=pe, attn_mask=attn_mask, transformer_options=transformer_options)"
    if old_double in content and new_double not in content:
        content = content.replace(old_double, new_double)
        changes.append("Patched double-block call with attn_mask + transformer_options")

    # Also handle alternative double-block patterns (varies by PuLID version)
    old_double_alt = "img, txt = block(img=img, txt=txt, vec=vec, pe=pe,"
    if old_double_alt in content and "attn_mask=attn_mask" not in content.split(old_double_alt)[1].split(")")[0]:
        # Already partially has extra kwargs - just add ours
        pass  # The simple replace above should handle the common case

    # 3. Patch single-block calls to pass attn_mask and transformer_options
    #    Pattern: img = block(img, vec=vec, pe=pe)
    old_single = "img = block(img, vec=vec, pe=pe)"
    new_single = "img = block(img, vec=vec, pe=pe, attn_mask=attn_mask, transformer_options=transformer_options)"
    if old_single in content and new_single not in content:
        content = content.replace(old_single, new_single)
        changes.append("Patched single-block call with attn_mask + transformer_options")

    if content == original:
        print("✅ No changes needed — patch already applied or pattern not found.")
        print("   Verify manually:")
        print(f"   grep -n 'attn_mask' {PULID_FILE}")
        sys.exit(0)

    # Write patched file
    with open(PULID_FILE, "w") as f:
        f.write(content)

    print(f"✅ Patched {PULID_FILE}")
    for c in changes:
        print(f"   - {c}")
    print("\n⚠️  Restart ComfyUI to reload the patched module:")
    print("   pkill -f 'python3 main.py'; cd /ComfyUI && python3 main.py --listen 0.0.0.0 --port 8188 &")
    print("\n⚠️  This patch is LIVE-FILESYSTEM ONLY — reverts on factory rebuild.")


if __name__ == "__main__":
    patch()