# ASCII-safe patch for Portrait Generator V5
# Run with: python PATCH_V5_FINAL_CANDIDATE_ASCII.py

from pathlib import Path
from datetime import datetime
import re
import shutil

root = Path(__file__).resolve().parent
candidates = [root / "_system" / "portrait_generator.py", root / "portrait_generator.py"]
app = next((p for p in candidates if p.exists()), None)
if app is None:
    print("ERROR: portrait_generator.py not found. Put this file next to INSTALLER.bat / LANCER.bat.")
    input("Press Enter to close...")
    raise SystemExit(1)

print(f"Patching: {app}")
backup = app.with_name(app.name + ".bak_" + datetime.now().strftime("%Y%m%d_%H%M%S"))
shutil.copy2(app, backup)
print(f"Backup created: {backup}")

s = app.read_text(encoding="utf-8")
changed = False

# 1) Auto-save all generated images immediately after generation.
if "autosaved_paths" not in s:
    pattern = re.compile(
        r'(current_params\["seeds"\] = seeds\s+'
        r'current_params\["seed"\] = current_seed\s+'
        r'current_params\["generation_mode"\] = generation_mode\s+'
        r'current_params\["num_images"\] = num_images\s+)'
        r'(progress\(1\.0, desc=.*?\))',
        re.DOTALL,
    )
    autosave = r'''\1
        # Auto-save: save every generated image immediately.
        autosaved_paths = []
        for auto_idx, auto_img in enumerate(images):
            try:
                auto_params = deepcopy(current_params)
                auto_params["autosave_index"] = auto_idx + 1
                if auto_idx < len(seeds):
                    auto_params["autosave_seed"] = seeds[auto_idx]
                autosaved_paths.append(save_portrait(auto_img, positive_prompt, auto_params))
            except Exception as save_err:
                print(f"Auto-save image {auto_idx+1} failed: {save_err}")
        current_params["autosaved_paths"] = autosaved_paths

        progress(1.0, desc="Generation complete and images auto-saved")'''
    s2, n = pattern.subn(autosave, s, count=1)
    if n:
        s = s2
        changed = True
        print("OK: auto-save added")
    else:
        print("WARNING: auto-save anchor not found")
else:
    print("Auto-save already present")

# 2) Add OpenPose ignored note for FLUX2 in info text.
if "openpose_note" not in s:
    pattern = re.compile(
        r'(mode_note = .*?\n)'
        r'(\s*info_text = \(f".*?\n\s*f".*?"\))',
        re.DOTALL,
    )
    repl = r'''\1        openpose_note = " | OpenPose ignored in FLUX.2 Klein" if (is_flux2 and pose_image is not None) else ""
        info_text = (f"OK {len(images)} image(s) auto-saved | Mode: {mode_note} | Final: {width}x{height} | LLM: {llm_used} | "
                     f"Style: {art_style} | Steps: {steps} | CFG: {cfg}{openpose_note}")'''
    s2, n = pattern.subn(repl, s, count=1)
    if n:
        s = s2
        changed = True
        print("OK: OpenPose/Flux2 info added")
    else:
        print("WARNING: info_text anchor not found")
else:
    print("OpenPose note already present")

# 3) Favorite button label, ASCII-safe.
s2, n = re.subn(
    r'final_btn = gr\.Button\([^\n]*variant="primary", interactive=False, scale=1\)',
    'final_btn = gr.Button("Save selected as favorite copy", variant="primary", interactive=False, scale=1)',
    s,
    count=1,
)
if n:
    s = s2
    changed = True
    print("OK: favorite button label patched")

# 4) OpenPose label and help text, ASCII-safe.
s2, n = re.subn(
    r'pose_image = gr\.Image\(label=.*?type="pil", height=180\)',
    'pose_image = gr.Image(label="Pose reference image (OpenPose - SDXL only)", type="pil", height=180)',
    s,
    count=1,
)
if n:
    s = s2
    changed = True
    print("OK: OpenPose image label patched")

s2, n = re.subn(
    r'pose_strength = gr\.Slider\(0\.0, 1\.5, value=0\.8, step=0\.05, label="Force OpenPose", info=.*?\)',
    'pose_strength = gr.Slider(0.0, 1.5, value=0.8, step=0.05, label="Force OpenPose", info="OpenPose works with SDXL fast/quality only. It is ignored in FLUX.2 Klein because SDXL ControlNet is not compatible with FLUX.2.")',
    s,
    count=1,
)
if n:
    s = s2
    changed = True
    print("OK: OpenPose help text patched")

# 5) Final info label.
s2 = s.replace('final_info = gr.Textbox(label="Statut", interactive=False)', 'final_info = gr.Textbox(label="Status / Favorite", interactive=False)')
if s2 != s:
    s = s2
    changed = True
    print("OK: status label patched")

if changed:
    app.write_text(s, encoding="utf-8")
    print("Patch applied successfully.")
else:
    print("No changes applied. File may already be patched or anchors differ.")

print("Restart LANCER.bat to use the patched app.")
input("Press Enter to close...")