#!/usr/bin/env python3
"""
FLUX Image Generator — Desktop GUI
Supports: FLUX.1 Schnell (fast) and FLUX.2 Klein Heretic (unrestricted)
Run: python3 gui.py
No external dependencies — stdlib + tkinter only.
"""

import json
import random
import threading
import time
import urllib.request
from pathlib import Path
import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox

COMFYUI = "http://localhost:8188"
OUTPUT_DIR = Path.home() / "Pictures" / "mcp-generated"
WORKFLOWS_DIR = Path(__file__).parent / "src/workflows"

# ── Model definitions ────────────────────────────────────────────────────────
# Each entry defines how to patch the workflow for that model.
# node_pos / node_neg: CLIPTextEncode node IDs for positive/negative prompts
# node_latent: latent image node (for width/height)
# node_seed: where to write the seed value
# node_save: SaveImage node (for filename_prefix)
# node_steps: dict of {node_id: field} for steps — None if not patchable
MODELS = {
    "FLUX.2 Klein Heretic (unrestricted)": {
        "workflow": "flux2_klein_heretic.json",
        "default_steps": 20,
        "node_pos":    ("2",  "text"),
        "node_neg":    ("3",  "text"),
        "node_latent": ("6",  "width", "height"),
        "node_seed":   ("10", "noise_seed"),
        "node_steps":  ("7",  "steps"),
        "node_save":   ("13", "filename_prefix"),
        "description": "FLUX.2 Klein 4B + Heretic abliterated encoder. ~50s/image. No refusals.",
    },
    "FLUX.1 Schnell (fast)": {
        "workflow": "flux_schnell.json",
        "default_steps": 4,
        "node_pos":    ("6",  "text"),
        "node_neg":    ("33", "text"),
        "node_latent": ("27", "width", "height"),
        "node_seed":   ("13", "seed"),       # KSampler has seed directly
        "node_steps":  ("13", "steps"),
        "node_save":   ("9",  "filename_prefix"),
        "description": "FLUX.1 Schnell — fast (~5s/image), standard quality. Has safety filter.",
    },
}

PRESETS = {
    "Square 1024":    (1024, 1024),
    "Landscape 16:9": (1280, 720),
    "Portrait 9:16":  (720, 1280),
    "Wide 3:2":       (1536, 1024),
    "Tall 2:3":       (1024, 1536),
    "BFL Wide 7:4":   (1344, 768),
    "BFL Tall 4:7":   (768, 1344),
}


def load_workflow(filename):
    with open(WORKFLOWS_DIR / filename) as f:
        return json.load(f)


def patch_workflow(wf, spec, prompt, neg, width, height, steps, seed, name):
    """Apply generation parameters to a workflow dict in-place."""
    node_pos_id, node_pos_field = spec["node_pos"]
    node_neg_id, node_neg_field = spec["node_neg"]
    lat_id, lat_w, lat_h = spec["node_latent"]
    seed_id, seed_field = spec["node_seed"]
    save_id, save_field = spec["node_save"]

    wf[node_pos_id]["inputs"][node_pos_field] = prompt
    wf[node_neg_id]["inputs"][node_neg_field] = neg
    wf[lat_id]["inputs"][lat_w] = width
    wf[lat_id]["inputs"][lat_h] = height
    wf[seed_id]["inputs"][seed_field] = seed
    wf[save_id]["inputs"][save_field] = name

    if spec["node_steps"]:
        steps_id, steps_field = spec["node_steps"]
        wf[steps_id]["inputs"][steps_field] = steps

    return wf


def submit_prompt(workflow):
    data = json.dumps({"prompt": workflow}).encode()
    req = urllib.request.Request(
        f"{COMFYUI}/prompt", data=data,
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read())["prompt_id"]


def wait_for_image(prompt_id, timeout=300):
    start = time.time()
    while time.time() - start < timeout:
        with urllib.request.urlopen(f"{COMFYUI}/history/{prompt_id}", timeout=10) as resp:
            history = json.loads(resp.read())
        if prompt_id in history:
            for node_out in history[prompt_id].get("outputs", {}).values():
                if "images" in node_out:
                    return node_out["images"][0]
            return None
        time.sleep(2)
    raise TimeoutError("Timed out waiting for image")


def download_image(filename, subfolder=""):
    url = f"{COMFYUI}/view?filename={filename}&subfolder={subfolder}&type=output"
    with urllib.request.urlopen(url, timeout=30) as resp:
        return resp.read()


class App(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("FLUX Image Generator")
        self.resizable(True, True)
        self.minsize(760, 640)
        self._current_image = None
        self._build_ui()

    def _build_ui(self):
        main = ttk.Frame(self, padding=12)
        main.pack(fill="both", expand=True)
        main.columnconfigure(0, weight=1)
        main.columnconfigure(1, weight=2)

        # ── LEFT PANEL ───────────────────────────────────────────────────────
        left = ttk.Frame(main)
        left.grid(row=0, column=0, sticky="nsew", padx=(0, 8))
        left.columnconfigure(0, weight=1)
        row = 0

        # Model selector
        ttk.Label(left, text="Model", font=("", 10, "bold")).grid(
            row=row, column=0, sticky="w"); row += 1
        self.model_var = tk.StringVar(value=list(MODELS.keys())[0])
        model_cb = ttk.Combobox(left, textvariable=self.model_var,
                                 values=list(MODELS.keys()), state="readonly", width=40)
        model_cb.grid(row=row, column=0, sticky="ew", pady=(2, 2)); row += 1
        model_cb.bind("<<ComboboxSelected>>", self._on_model_change)
        self.model_desc = ttk.Label(left, text="", foreground="gray",
                                     wraplength=300, font=("", 8))
        self.model_desc.grid(row=row, column=0, sticky="w", pady=(0, 8)); row += 1
        # description updated after all widgets are created (see end of _build_ui)

        # Prompt
        ttk.Label(left, text="Prompt", font=("", 10, "bold")).grid(
            row=row, column=0, sticky="w"); row += 1
        self.prompt_txt = scrolledtext.ScrolledText(left, height=6, wrap="word", font=("", 10))
        self.prompt_txt.grid(row=row, column=0, sticky="ew", pady=(2, 8)); row += 1

        # Negative prompt
        ttk.Label(left, text="Negative Prompt (optional)").grid(
            row=row, column=0, sticky="w"); row += 1
        self.neg_txt = scrolledtext.ScrolledText(left, height=3, wrap="word", font=("", 9))
        self.neg_txt.grid(row=row, column=0, sticky="ew", pady=(2, 8)); row += 1

        # Size preset
        ttk.Label(left, text="Size Preset").grid(row=row, column=0, sticky="w"); row += 1
        self.preset_var = tk.StringVar(value="Square 1024")
        preset_cb = ttk.Combobox(left, textvariable=self.preset_var,
                                  values=list(PRESETS.keys()), state="readonly")
        preset_cb.grid(row=row, column=0, sticky="ew", pady=(2, 2)); row += 1
        preset_cb.bind("<<ComboboxSelected>>", self._apply_preset)

        # Width / Height
        wh = ttk.Frame(left)
        wh.grid(row=row, column=0, sticky="ew", pady=(4, 8)); row += 1
        wh.columnconfigure(1, weight=1)
        wh.columnconfigure(3, weight=1)
        ttk.Label(wh, text="W").grid(row=0, column=0, padx=(0, 4))
        self.width_var = tk.IntVar(value=1024)
        ttk.Spinbox(wh, from_=256, to=4096, increment=8, textvariable=self.width_var,
                    width=6).grid(row=0, column=1, sticky="ew")
        ttk.Label(wh, text="H", padding=(8, 0, 4, 0)).grid(row=0, column=2)
        self.height_var = tk.IntVar(value=1024)
        ttk.Spinbox(wh, from_=256, to=4096, increment=8, textvariable=self.height_var,
                    width=6).grid(row=0, column=3, sticky="ew")

        # Steps
        steps_row = ttk.Frame(left)
        steps_row.grid(row=row, column=0, sticky="ew", pady=(0, 4)); row += 1
        steps_row.columnconfigure(1, weight=1)
        ttk.Label(steps_row, text="Steps").grid(row=0, column=0, padx=(0, 8))
        self.steps_var = tk.IntVar(value=20)
        self.steps_lbl = ttk.Label(steps_row, text="20", width=3)
        self.steps_lbl.grid(row=0, column=2, padx=(6, 0))
        ttk.Scale(steps_row, from_=1, to=60, variable=self.steps_var,
                  orient="horizontal",
                  command=lambda v: self.steps_lbl.config(text=str(int(float(v))))
                  ).grid(row=0, column=1, sticky="ew")

        # Count
        count_row = ttk.Frame(left)
        count_row.grid(row=row, column=0, sticky="ew", pady=(0, 8)); row += 1
        count_row.columnconfigure(1, weight=1)
        ttk.Label(count_row, text="Count").grid(row=0, column=0, padx=(0, 8))
        self.count_var = tk.IntVar(value=1)
        ttk.Spinbox(count_row, from_=1, to=20, textvariable=self.count_var,
                    width=4).grid(row=0, column=1, sticky="w")

        # Seed
        seed_row = ttk.Frame(left)
        seed_row.grid(row=row, column=0, sticky="ew", pady=(0, 8)); row += 1
        seed_row.columnconfigure(1, weight=1)
        ttk.Label(seed_row, text="Seed").grid(row=0, column=0, padx=(0, 8))
        self.seed_var = tk.StringVar(value="-1")
        ttk.Entry(seed_row, textvariable=self.seed_var, width=12).grid(row=0, column=1, sticky="w")
        ttk.Button(seed_row, text="🎲", width=3,
                   command=lambda: self.seed_var.set(str(random.randint(0, 2**32 - 1)))
                   ).grid(row=0, column=2, padx=(4, 0))
        ttk.Label(seed_row, text="(-1 = random)", foreground="gray"
                  ).grid(row=0, column=3, padx=(4, 0))

        # Name
        name_row = ttk.Frame(left)
        name_row.grid(row=row, column=0, sticky="ew", pady=(0, 12)); row += 1
        name_row.columnconfigure(1, weight=1)
        ttk.Label(name_row, text="Name").grid(row=0, column=0, padx=(0, 8))
        self.name_var = tk.StringVar(value="img")
        ttk.Entry(name_row, textvariable=self.name_var).grid(row=0, column=1, sticky="ew")

        # Generate button
        self.gen_btn = ttk.Button(left, text="⚡ Generate", command=self._start_generation)
        self.gen_btn.grid(row=row, column=0, sticky="ew", pady=(0, 8)); row += 1

        # Status + progress
        self.status_var = tk.StringVar(value="Ready")
        ttk.Label(left, textvariable=self.status_var, foreground="gray",
                  wraplength=300).grid(row=row, column=0, sticky="w"); row += 1
        self.progress = ttk.Progressbar(left, mode="indeterminate")
        self.progress.grid(row=row, column=0, sticky="ew", pady=(4, 0)); row += 1

        # ── RIGHT PANEL — preview ────────────────────────────────────────────
        right = ttk.LabelFrame(main, text="Preview", padding=8)
        right.grid(row=0, column=1, sticky="nsew")
        right.columnconfigure(0, weight=1)
        right.rowconfigure(0, weight=1)

        self.preview_lbl = ttk.Label(right, text="No image yet", anchor="center",
                                      background="#1a1a1a", foreground="#888")
        self.preview_lbl.grid(row=0, column=0, sticky="nsew")

        self.path_lbl = ttk.Label(right, text="", foreground="gray", font=("", 8))
        self.path_lbl.grid(row=1, column=0, sticky="w", pady=(4, 0))

        ttk.Button(right, text="Open folder",
                   command=self._open_folder).grid(row=2, column=0, sticky="e", pady=(4, 0))

        # Init model description + steps default now that all widgets exist
        self._on_model_change()

    def _on_model_change(self, _=None):
        spec = MODELS[self.model_var.get()]
        self.model_desc.config(text=spec["description"])
        self.steps_var.set(spec["default_steps"])
        self.steps_lbl.config(text=str(spec["default_steps"]))

    def _apply_preset(self, _=None):
        w, h = PRESETS[self.preset_var.get()]
        self.width_var.set(w)
        self.height_var.set(h)

    def _start_generation(self):
        prompt = self.prompt_txt.get("1.0", "end").strip()
        if not prompt:
            messagebox.showwarning("No prompt", "Please enter a prompt.")
            return
        self.gen_btn.config(state="disabled")
        self.progress.start(10)
        self.status_var.set("Generating…")
        t = threading.Thread(target=self._run_generation, args=(prompt,), daemon=True)
        t.start()

    def _run_generation(self, prompt):
        try:
            neg = self.neg_txt.get("1.0", "end").strip()
            steps = int(self.steps_var.get())
            width = int(self.width_var.get())
            height = int(self.height_var.get())
            count = int(self.count_var.get())
            name = self.name_var.get().strip() or "img"
            seed_str = self.seed_var.get().strip()
            base_seed = int(seed_str) if seed_str else -1

            model_name = self.model_var.get()
            spec = MODELS[model_name]

            OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

            for i in range(count):
                seed = (base_seed if base_seed == -1 else base_seed + i)
                if seed == -1:
                    seed = random.randint(0, 2**32 - 1)

                label = f"{name}_{i+1:02d}" if count > 1 else name
                self.after(0, self.status_var.set,
                           f"[{i+1}/{count}] {model_name.split('(')[0].strip()} · seed {seed}…")

                wf = load_workflow(spec["workflow"])
                patch_workflow(wf, spec, prompt, neg, width, height, steps, seed, label)

                prompt_id = submit_prompt(wf)
                img_info = wait_for_image(prompt_id)

                if img_info:
                    img_data = download_image(img_info["filename"], img_info.get("subfolder", ""))
                    out_path = OUTPUT_DIR / f"{label}_{seed}.png"
                    out_path.write_bytes(img_data)
                    self.after(0, self._show_preview, out_path)

            self.after(0, self.status_var.set,
                       f"✅ Done — {count} image(s) saved to ~/Pictures/mcp-generated/")
        except Exception as exc:
            self.after(0, self.status_var.set, f"❌ Error: {exc}")
        finally:
            self.after(0, self.progress.stop)
            self.after(0, lambda: self.gen_btn.config(state="normal"))

    def _show_preview(self, path):
        try:
            photo = tk.PhotoImage(file=str(path))
            pw, ph = photo.width(), photo.height()
            subsample = 1
            while pw // subsample > 600 or ph // subsample > 600:
                subsample += 1
            if subsample > 1:
                photo = photo.subsample(subsample, subsample)
            self.preview_lbl.config(image=photo, text="")
            self._current_image = photo
            self.path_lbl.config(text=str(path))
        except Exception as e:
            self.status_var.set(f"Preview error: {e}")

    def _open_folder(self):
        import subprocess
        subprocess.Popen(["xdg-open", str(OUTPUT_DIR)])


if __name__ == "__main__":
    app = App()
    app.mainloop()