import os
os.environ["MPLCONFIGDIR"] = "/tmp"  # matplotlib sans ~/.config

import gradio as gr
from pathlib import Path
import subprocess, json, tempfile, sys

# Répertoires
COMFY_DIR = Path("ComfyUI")
WORKFLOW_JSON = Path("workflow_api.json")
LOGS_DIR = Path(tempfile.gettempdir()) / "comfy_logs"
LOGS_DIR.mkdir(parents=True, exist_ok=True)

def patch_workflow(wf, pic_path, face_path, prompt):
    for node in wf.values():
        if node.get("_meta", {}).get("title") == "Picture to swap":
            node["inputs"]["image"] = str(pic_path)
        if node.get("_meta", {}).get("title") == "Input face":
            node["inputs"]["image"] = str(face_path)
        if node["class_type"] == "easy positive":
            node["inputs"]["text"] = prompt
    return wf

def run(picture, face, positive):
    tmpdir = Path(tempfile.mkdtemp())
    pic_path = tmpdir / "pic.png"
    face_path = tmpdir / "face.png"
    picture.save(pic_path)
    face.save(face_path)

    with open(WORKFLOW_JSON) as f:
        wf = json.load(f)
    wf = patch_workflow(wf, pic_path, face_path, positive)

    tmp_wf = tmpdir / "wf.json"
    with open(tmp_wf, "w") as f:
        json.dump(wf, f)

    result = subprocess.run([
        sys.executable, "main.py",
        "--disable-auto-launch",
        "--input", str(tmpdir),
        "--output", str(tmpdir),
        "--workflow", str(tmp_wf)
    ], cwd=COMFY_DIR, capture_output=True, text=True)

    log_file = LOGS_DIR / "comfy.txt"
    with open(log_file, "w") as f:
        f.write("=== STDOUT ===\n")
        f.write(result.stdout)
        f.write("\n=== STDERR ===\n")
        f.write(result.stderr)

    if result.returncode != 0:
        raise RuntimeError("ComfyUI failed to run — voir comfy.txt dans Files.")

    out = list(tmpdir.glob("*.png"))
    if not out:
        raise RuntimeError("No output image generated.")

    final_img = LOGS_DIR / "output.png"
    out[0].replace(final_img)

    return final_img, log_file

demo = gr.Interface(
    fn=run,
    inputs=[
        gr.Image(type="pil", label="Picture to swap"),
        gr.Image(type="pil", label="Face input"),
        gr.Textbox(label="Positive prompt", lines=3,
                   placeholder="portrait of a woman, detailed face...")
    ],
    outputs=[
        gr.Image(type="filepath", label="Swapped"),
        gr.File(label="ComfyUI logs")
    ],
    title="ComfyUI InstantID Face Swap",
    allow_flagging="never"
)

demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)