#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
open_image_ai_gui.py
Single-file, fully open-source, offline-capable image AI with bilingual GUI (FA/EN).
- No paid APIs, no accounts. Uses local open-source models (diffusers).
- Supports SDXL and SD 1.5 if available; falls back to local illusion generator.
- Prompt cycling, caching, post-processing (10 filters), and 10 anti-error mechanisms.
- Build into one-file executable with PyInstaller if desired.

Run:
  python open_image_ai_gui.py

Build (optional):
  python -m pip install pyinstaller
  pyinstaller --onefile open_image_ai_gui.py
"""

import os
import sys
import json
import time
import math
import hashlib
import random
import threading
from datetime import datetime
from pathlib import Path

# Lazy installer for single-file friendliness
def _lazy_install(pkg):
    try:
        __import__(pkg)
    except ImportError:
        os.system(f"{sys.executable} -m pip install --quiet {pkg}")

for pkg in ("torch", "diffusers", "accelerate", "safetensors", "Pillow"):
    _lazy_install(pkg)

import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import torch
from PIL import Image, ImageOps, ImageFilter, ImageDraw
from diffusers import StableDiffusionPipeline, DiffusionPipeline

# ---------------------------
# Config & Directories
# ---------------------------
APP_NAME = "Open Image AI"
OUTPUT_DIR = Path(os.environ.get("OIA_OUT", "outputs")).resolve()
META_DIR = OUTPUT_DIR / "meta"
CACHE_DIR = OUTPUT_DIR / "cache"
TMP_DIR = OUTPUT_DIR / "tmp"
for d in (OUTPUT_DIR, META_DIR, CACHE_DIR, TMP_DIR):
    d.mkdir(parents=True, exist_ok=True)

# Defaults
DEFAULT_STEPS = int(os.environ.get("OIA_STEPS", "30"))
DEFAULT_GUIDANCE = float(os.environ.get("OIA_GUIDANCE", "7.5"))
MAX_QUEUE = int(os.environ.get("OIA_MAX_QUEUE", "64"))
MAX_PARALLEL = int(os.environ.get("OIA_MAX_PARALLEL", "1"))  # GUI-friendly
SEED_AUTOMATIC = -1

# Available open models (pullable via diffusers). You can place local copies too.
OPEN_MODELS = {
    "SDXL (base)": "stabilityai/stable-diffusion-xl-base-1.0",
    "SD 1.5 (classic)": "runwayml/stable-diffusion-v1-5",
    # You can add other open models here if supported by diffusers:
    # "Stable Cascade": "stabilityai/stable-cascade",
    # "OpenJourney (SD1.5 finetune)": "prompthero/openjourney-v4",
}

# ---------------------------
# Prompt utilities
# ---------------------------
OPTICAL_ILLUSION_PRESETS = [
    "high-contrast black and white optical illusion, impossible geometry, Penrose stairs, concentric lines, parallax shifts",
    "moiré patterns, nested grids, Escher-inspired recursion, tilt-shift depth cues, perspective ambiguity",
    "rotational symmetry, non-Euclidean corridor, variable line thickness, alternating diagonals, dizzying parallax",
    "spiral tunnels, zigzag pathways, impossible cube, nested frames, sharp edges, stark contrast",
]

WEIGHTS = [
    "hyper-detailed", "vector-sharp edges", "ultra high resolution",
    "minimal noise", "geometric precision", "dynamic parallax",
    "photorealistic lighting"
]

MODS = [
    "isometric view", "45-degree inclination", "multi-angle perception",
    "viewer-dependent illusion", "top-down perspective"
]

def vary_prompt(base: str) -> str:
    extra = ", ".join(random.sample(OPTICAL_ILLUSION_PRESETS, k=random.randint(1, 3)))
    hint = ", ".join(random.sample(WEIGHTS, k=random.randint(2, 4)))
    mod = random.choice(MODS)
    return f"{base}, {extra}, {hint}, {mod}"

def prompt_hash(p: str) -> str:
    return hashlib.sha256(p.encode("utf-8")).hexdigest()[:16]

# ---------------------------
# Cache & Meta
# ---------------------------
def cache_path_for(prompt: str) -> Path:
    return CACHE_DIR / f"{prompt_hash(prompt)}.png"

def cached(prompt: str) -> Path | None:
    p = cache_path_for(prompt)
    return p if p.exists() else None

def write_meta(meta: dict, fname: Path):
    meta["timestamp"] = datetime.utcnow().isoformat()
    with open(META_DIR / (fname.stem + ".json"), "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)

# ---------------------------
# Anti-error mechanisms (10)
# ---------------------------
class AntiError:
    def __init__(self):
        self.fail_count = 0
        self.last_minute_calls = []
        self.circuit_open = False
        self.warned_no_model = False

    # 1) Retries cap
    def allow_retry(self, attempt, max_retries=3):
        return attempt < max_retries

    # 2) Backoff incremental
    def backoff(self, attempt):
        time.sleep(1.2 * attempt + random.random())

    # 3) Rate limiting per minute
    def rate_limit_ok(self, rate=60):  # 60 calls/min cap default
        now = time.time()
        self.last_minute_calls = [t for t in self.last_minute_calls if now - t < 60]
        if len(self.last_minute_calls) >= rate:
            return False
        self.last_minute_calls.append(now)
        return True

    # 4) Circuit breaker
    def circuit_should_open(self):
        return self.fail_count >= 5

    def circuit_reset(self):
        self.fail_count = 0
        self.circuit_open = False

    # 5) Cache short-circuit
    def cache_hit(self, prompt):
        return cached(prompt)

    # 6) Dedup by prompt hash
    def dedup_name(self, prompt):
        return prompt_hash(prompt)

    # 7) Health check (model presence)
    def has_model(self, model_loaded):
        return model_loaded is not None

    # 8) Seed normalization
    def normalize_seed(self, seed):
        return None if seed is None or seed < 0 else seed

    # 9) Logging
    def log(self, s):
        print(s)

    # 10) Fallback flag when model missing
    def warn_model_missing(self):
        if not self.warned_no_model:
            self.log("[!] No model loaded. Using local illusion generator as fallback.")
            self.warned_no_model = True

anti = AntiError()

# ---------------------------
# Local illusion generator (offline fallback)
# ---------------------------
def local_illusion(size=1024) -> Image.Image:
    img = Image.new("RGB", (size, size), "white")
    draw = ImageDraw.Draw(img)

    # Concentric square rings
    step = 6
    for r in range(24, size//2, step):
        val = 0 if (r//step) % 2 == 0 else 255
        # top and bottom lines
        for x in range(size):
            if abs(x - size//2) == r:
                draw.line([(x, 0), (x, size)], fill=(val, val, val))
        # left and right lines
        for y in range(size):
            if abs(y - size//2) == r:
                draw.line([(0, y), (size, y)], fill=(val, val, val))

    # Diagonal hatch
    hatch = Image.new("L", (size, size), color=255)
    hpx = hatch.load()
    for x in range(size):
        for y in range(size):
            if (x + y) % 18 == 0:
                hpx[x, y] = 0

    merged = Image.merge("RGB", (ImageOps.autocontrast(hatch), ImageOps.autocontrast(hatch), ImageOps.autocontrast(hatch)))
    merged = Image.blend(img, merged, 0.35)
    merged = merged.filter(ImageFilter.SHARPEN)
    merged = merged.rotate(random.choice([4, -6, 8, -10]), resample=Image.BICUBIC, expand=False)
    return merged

# ---------------------------
# Post-processing filters (10)
# ---------------------------
POST_MODES = [
    "illusion_boost", "bw_high", "moire_mix", "invert", "edge_halo",
    "soft_glow", "posterize4", "contrast_max", "rotate_slight", "grain_fine"
]

def post_process(img: Image.Image, mode: str) -> Image.Image:
    if mode == "illusion_boost":
        out = ImageOps.autocontrast(img).filter(ImageFilter.UnsharpMask(radius=2, percent=180, threshold=3))
    elif mode == "bw_high":
        out = ImageOps.autocontrast(ImageOps.grayscale(img)).convert("RGB")
    elif mode == "moire_mix":
        out = ImageOps.posterize(img.filter(ImageFilter.DETAIL).filter(ImageFilter.SHARPEN), bits=4)
    elif mode == "invert":
        out = ImageOps.invert(img)
    elif mode == "edge_halo":
        out = img.filter(ImageFilter.UnsharpMask(radius=3, percent=240, threshold=2))
    elif mode == "soft_glow":
        blur = img.filter(ImageFilter.GaussianBlur(radius=2))
        out = Image.blend(img, blur, alpha=0.3)
    elif mode == "posterize4":
        out = ImageOps.posterize(img, bits=4)
    elif mode == "contrast_max":
        out = ImageOps.autocontrast(img)
    elif mode == "rotate_slight":
        out = img.rotate(random.choice([-3, 3, -2, 2]), resample=Image.BICUBIC, expand=False)
    elif mode == "grain_fine":
        out = img.filter(ImageFilter.SMOOTH_MORE)
    else:
        out = img
    return out

# ---------------------------
# Model holder and generator
# ---------------------------
class ModelHub:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.pipe = None
        self.model_id = None

    def load(self, model_id: str, fp16=True):
        self.model_id = model_id
        dtype = torch.float16 if (fp16 and self.device == "cuda") else torch.float32
        try:
            if "stable-diffusion-xl" in model_id:
                self.pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype, variant="fp16" if dtype==torch.float16 else None)
            else:
                self.pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
            self.pipe = self.pipe.to(self.device)
            if self.device == "cuda":
                self.pipe.enable_attention_slicing()
            anti.circuit_reset()
            return True
        except Exception as e:
            anti.log(f"[!] Model load error: {e}")
            self.pipe = None
            return False

    def generate(self, prompt: str, steps: int, guidance: float, seed: int | None, height=1024, width=1024) -> Image.Image:
        if self.pipe is None:
            anti.warn_model_missing()
            return local_illusion(size=min(height, width))
        generator = torch.Generator(device=self.device)
        if seed is not None:
            generator = generator.manual_seed(seed)
        try:
            if hasattr(self.pipe, "generate"):
                # DiffusionPipeline custom generate (rare)
                out = self.pipe.generate(prompt, height=height, width=width, num_inference_steps=steps, guidance_scale=guidance, generator=generator)
                img = out.images[0] if hasattr(out, "images") else out[0]
            else:
                out = self.pipe(prompt, height=height, width=width, num_inference_steps=steps, guidance_scale=guidance, generator=generator)
                img = out.images[0]
            return img
        except Exception as e:
            anti.log(f"[!] Generate error: {e}")
            anti.fail_count += 1
            if anti.circuit_should_open():
                anti.circuit_open = True
            return local_illusion(size=min(height, width))

hub = ModelHub()

# ---------------------------
# Orchestrator (queue-based)
# ---------------------------
class Orchestrator:
    def __init__(self):
        self.queue = []
        self.running = False

    def enqueue(self, prompt, count, vary, steps, guidance, seed, size, post_mode):
        for _ in range(count):
            if len(self.queue) >= MAX_QUEUE:
                break
            self.queue.append({
                "prompt": prompt,
                "vary": vary,
                "steps": steps,
                "guidance": guidance,
                "seed": seed,
                "size": size,
                "post_mode": post_mode
            })

    def _do_item(self, item):
        p = item["prompt"]
        if item["vary"]:
            p = vary_prompt(p)
        # Cache check
        cp = cached(p)
        img = None
        if cp:
            img = Image.open(cp).convert("RGB")
        else:
            seed = anti.normalize_seed(item["seed"])
            img = hub.generate(
                prompt=p,
                steps=item["steps"],
                guidance=item["guidance"],
                seed=seed,
                height=item["size"],
                width=item["size"],
            )
            # save to cache
            try:
                cache_path = cache_path_for(p)
                img.save(cache_path, "PNG", optimize=True)
            except Exception as e:
                anti.log(f"[!] Cache save error: {e}")

        if item["post_mode"]:
            img = post_process(img, item["post_mode"])

        # final save
        fname = OUTPUT_DIR / f"{datetime.utcnow().strftime('%Y%m%d_%H%M%S_%f')}_{prompt_hash(p)}.png"
        img.save(fname, "PNG", optimize=True)
        write_meta({"prompt": p, "model": hub.model_id or "local_illusion", "post_mode": item["post_mode"]}, fname)
        return fname

    def run(self, on_log):
        if self.running:
            return
        self.running = True
        try:
            while self.queue:
                if not anti.rate_limit_ok(rate=60):
                    time.sleep(1.0)
                item = self.queue.pop(0)
                attempt = 0
                result_path = None
                while anti.allow_retry(attempt, max_retries=3):
                    try:
                        result_path = self._do_item(item)
                        anti.circuit_reset()
                        break
                    except Exception as e:
                        anti.fail_count += 1
                        on_log(f"    x Error: {e}")
                        attempt += 1
                        anti.backoff(attempt)
                if result_path:
                    on_log(f"    -> Saved: {result_path}")
                else:
                    on_log("    x Failed after retries")
        finally:
            self.running = False

orch = Orchestrator()

# ---------------------------
# GUI (Tkinter) bilingual
# ---------------------------
class App:
    def __init__(self, root):
        self.root = root
        root.title(APP_NAME)
        root.geometry("980x720")

        self.lang = tk.StringVar(value="FA")
        self.prompt = tk.StringVar()
        self.count = tk.IntVar(value=3)
        self.vary = tk.BooleanVar(value=True)
        self.steps = tk.IntVar(value=DEFAULT_STEPS)
        self.guidance = tk.DoubleVar(value=DEFAULT_GUIDANCE)
        self.seed = tk.IntVar(value=SEED_AUTOMATIC)
        self.size = tk.IntVar(value=768)
        self.post_mode = tk.StringVar(value="")
        self.model_choice = tk.StringVar(value=list(OPEN_MODELS.keys())[0])

        self._build_ui()
        self._set_texts()
        self._log_header()

    def _build_ui(self):
        top = ttk.Frame(self.root)
        top.pack(fill="x", padx=10, pady=10)

        ttk.Label(top, text="Language / زبان").pack(side="left")
        ttk.Combobox(top, textvariable=self.lang, values=["FA", "EN"], width=5, state="readonly").pack(side="left", padx=8)
        ttk.Button(top, text="Apply/اعمال", command=self._set_texts).pack(side="left", padx=8)

        model_frame = ttk.LabelFrame(self.root, text="Model / مدل")
        model_frame.pack(fill="x", padx=10, pady=6)
        ttk.Label(model_frame, text="Select model:").grid(row=0, column=0, sticky="w", padx=6, pady=6)
        ttk.Combobox(model_frame, textvariable=self.model_choice, values=list(OPEN_MODELS.keys()), width=28, state="readonly").grid(row=0, column=1, sticky="w", padx=6, pady=6)
        ttk.Button(model_frame, text="Load model", command=self.load_model).grid(row=0, column=2, padx=6, pady=6)

        mid = ttk.LabelFrame(self.root, text="Controls / کنترل‌ها")
        mid.pack(fill="x", padx=10, pady=6)

        # Row 0
        self.lbl_prompt = ttk.Label(mid, text="")
        self.lbl_prompt.grid(row=0, column=0, sticky="w", padx=6, pady=6)
        ttk.Entry(mid, textvariable=self.prompt, width=72).grid(row=0, column=1, sticky="we", padx=6, pady=6)
        # Row 1
        self.lbl_count = ttk.Label(mid, text="")
        self.lbl_count.grid(row=1, column=0, sticky="w", padx=6, pady=6)
        ttk.Spinbox(mid, from_=1, to=50, textvariable=self.count, width=6).grid(row=1, column=1, sticky="w", padx=6, pady=6)
        # Row 2
        self.lbl_vary = ttk.Label(mid, text="")
        self.lbl_vary.grid(row=2, column=0, sticky="w", padx=6, pady=6)
        ttk.Checkbutton(mid, text="", variable=self.vary).grid(row=2, column=1, sticky="w", padx=6, pady=6)
        # Row 3
        self.lbl_steps = ttk.Label(mid, text="")
        self.lbl_steps.grid(row=3, column=0, sticky="w", padx=6, pady=6)
        ttk.Spinbox(mid, from_=5, to=100, textvariable=self.steps, width=6).grid(row=3, column=1, sticky="w", padx=6, pady=6)
        # Row 4
        self.lbl_guidance = ttk.Label(mid, text="")
        self.lbl_guidance.grid(row=4, column=0, sticky="w", padx=6, pady=6)
        ttk.Spinbox(mid, from_=0.0, to=20.0, increment=0.5, textvariable=self.guidance, width=6).grid(row=4, column=1, sticky="w", padx=6, pady=6)
        # Row 5
        self.lbl_seed = ttk.Label(mid, text="")
        self.lbl_seed.grid(row=5, column=0, sticky="w", padx=6, pady=6)
        ttk.Spinbox(mid, from_=-1, to=2**31-1, textvariable=self.seed, width=12).grid(row=5, column=1, sticky="w", padx=6, pady=6)
        # Row 6
        self.lbl_size = ttk.Label(mid, text="")
        self.lbl_size.grid(row=6, column=0, sticky="w", padx=6, pady=6)
        ttk.Spinbox(mid, from_=256, to=1024, increment=64, textvariable=self.size, width=6).grid(row=6, column=1, sticky="w", padx=6, pady=6)
        # Row 7
        self.lbl_pp = ttk.Label(mid, text="")
        self.lbl_pp.grid(row=7, column=0, sticky="w", padx=6, pady=6)
        ttk.Combobox(mid, textvariable=self.post_mode, values=[""] + POST_MODES, width=20, state="readonly").grid(row=7, column=1, sticky="w", padx=6, pady=6)

        btns = ttk.Frame(mid)
        btns.grid(row=8, column=0, columnspan=2, sticky="we", padx=6, pady=6)
        self.btn_run = ttk.Button(btns, text="", command=self.run_queue)
        self.btn_run.pack(side="left", padx=6)
        ttk.Button(btns, text="Open outputs", command=self.open_outputs).pack(side="left", padx=6)
        ttk.Button(btns, text="Illusion offline test", command=self.make_offline_illusion).pack(side="left", padx=6)

        self.log = tk.Text(self.root, height=18, wrap="word")
        self.log.pack(fill="both", expand=True, padx=10, pady=10)

        status_bar = ttk.Frame(self.root)
        status_bar.pack(fill="x", padx=10, pady=5)
        self.status = tk.StringVar(value="Ready")
        self.lbl_status = ttk.Label(status_bar, textvariable=self.status)
        self.lbl_status.pack(side="left")

    def _set_texts(self):
        if self.lang.get() == "FA":
            self.lbl_prompt.config(text="متن راهنمای تصویر (پرومپت):")
            self.lbl_count.config(text="تعداد تولید:")
            self.lbl_vary.config(text="تنوع‌دهی خودکار پرومپت:")
            self.lbl_steps.config(text="تعداد گام‌ها:")
            self.lbl_guidance.config(text="راهنمایی (Guidance):")
            self.lbl_seed.config(text="Seed (برای تکرارپذیری، -1 خودکار):")
            self.lbl_size.config(text="اندازه (پیکسل مربع):")
            self.lbl_pp.config(text="پس‌پردازش:")
            self.btn_run.config(text="شروع تولید")
            self.root.title(f"{APP_NAME} - رابط فارسی")
        else:
            self.lbl_prompt.config(text="Image prompt:")
            self.lbl_count.config(text="Count:")
            self.lbl_vary.config(text="Auto prompt variation:")
            self.lbl_steps.config(text="Steps:")
            self.lbl_guidance.config(text="Guidance:")
            self.lbl_seed.config(text="Seed (-1 automatic):")
            self.lbl_size.config(text="Size (square px):")
            self.lbl_pp.config(text="Post-process:")
            self.btn_run.config(text="Start")
            self.root.title(f"{APP_NAME} - English UI")

    def _log_header(self):
        self.append_log(f"[{APP_NAME}] Ready. Select a model, enter prompt, and start.")

    def append_log(self, s):
        self.log.insert("end", s + "\n")
        self.log.see("end")

    def open_outputs(self):
        path = str(OUTPUT_DIR)
        try:
            if os.name == "nt":
                os.startfile(path)
            elif sys.platform == "darwin":
                os.system(f'open "{path}"')
            else:
                os.system(f'xdg-open "{path}"')
        except Exception:
            messagebox.showinfo("Info", f"Outputs at: {path}")

    def load_model(self):
        name = self.model_choice.get()
        model_id = OPEN_MODELS.get(name)
        if not model_id:
            messagebox.showerror("Error", "Model id not found.")
            return
        self.status.set("Loading model...")
        self.append_log(f"[+] Loading: {name} -> {model_id}")
        self.root.update_idletasks()
        ok = hub.load(model_id, fp16=True)
        if ok:
            self.status.set("Model loaded")
            self.append_log("    -> Model ready.")
        else:
            self.status.set("Model load failed (fallback active)")
            self.append_log("    x Model load failed; using local illusion fallback.")

    def run_queue(self):
        base_prompt = self.prompt.get().strip()
        if not base_prompt:
            messagebox.showwarning("Warn", "Please enter a prompt / لطفاً پرومپت را وارد کنید")
            return
        count = self.count.get()
        vary = self.vary.get()
        steps = self.steps.get()
        guidance = self.guidance.get()
        seed = self.seed.get()
        size = self.size.get()
        pp = self.post_mode.get()

        orch.enqueue(base_prompt, count, vary, steps, guidance, seed, size, pp)
        self.status.set("Running...")
        self.append_log(f"[+] {datetime.now().strftime('%H:%M:%S')} Enqueued: {base_prompt} x {count} vary={vary} steps={steps} guidance={guidance} seed={seed} size={size} pp={pp}")

        def worker():
            try:
                orch.run(self.append_log)
            finally:
                self.status.set("Done")

        threading.Thread(target=worker, daemon=True).start()

    def make_offline_illusion(self):
        img = local_illusion(size=self.size.get())
        fname = OUTPUT_DIR / f"{datetime.utcnow().strftime('%Y%m%d_%H%M%S_%f')}_offline_illusion.png"
        img.save(fname, "PNG", optimize=True)
        write_meta({"prompt": "offline_illusion", "model": "local_illusion"}, fname)
        self.append_log(f"    -> Saved offline illusion: {fname}")

# ---------------------------
# Entry
# ---------------------------
def main():
    root = tk.Tk()
    style = ttk.Style()
    try:
        style.theme_use("clam")
    except Exception:
        pass
    app = App(root)
    root.mainloop()

if __name__ == "__main__":
    main()