#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright (C) 2026 DiogenOS
"""b70-imagegen control panel.

A small window to set up (download model + install) and then Start / Restart /
Shut down the ComfyUI server. The three control buttons are disabled until the
setup is complete (venv + portable runtime + model present).

Runs under the SYSTEM python3 (needs python3-tk); it does not import torch.
It shells out to /usr/bin/b70-imagegen and /usr/bin/b70-imagegen-setup, which
use the per-user venv.
"""
import os
import signal
import socket
import subprocess
import sys
import threading
import time
import webbrowser

HOME = os.path.expanduser("~")
VENV_PY = os.path.join(HOME, "img-gen", "bin", "python")
ENV_SH = os.path.join(HOME, "ai-image", "intel-rt", "env.sh")
MODEL = os.path.join(
    HOME, "ai-image", "ComfyUI", "models", "checkpoints",
    "flux1-schnell-fp8.safetensors",
)
PORT = 8188
URL = "http://127.0.0.1:%d" % PORT
LAUNCHER = "/usr/bin/b70-imagegen"
SETUP = "/usr/bin/b70-imagegen-setup"
VERSION_FILE = "/usr/share/b70-imagegen/VERSION"


def version():
    try:
        with open(VERSION_FILE) as f:
            return f.read().strip()
    except OSError:
        return ""


def is_set_up():
    """True once the heavy per-user install (incl. the model) is complete."""
    return (
        os.path.exists(VENV_PY)
        and os.path.exists(ENV_SH)
        and os.path.exists(MODEL)
    )


def port_open(host="127.0.0.1", port=PORT, timeout=0.5):
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True
    except OSError:
        return False


def selftest():
    print("version   :", version() or "(none)")
    print("venv_py   :", VENV_PY, "->", os.path.exists(VENV_PY))
    print("env.sh    :", ENV_SH, "->", os.path.exists(ENV_SH))
    print("model     :", MODEL, "->", os.path.exists(MODEL))
    print("is_set_up :", is_set_up())
    print("port_open :", port_open())
    return 0


# Lazy tkinter import so the helpers above (and --selftest) work even where
# python3-tk is not installed.
try:
    import tkinter as tk
    from tkinter import ttk, scrolledtext
    _HAVE_TK = True
except Exception:
    _HAVE_TK = False


class App:
    def __init__(self, root):
        self.root = root
        self.proc = None              # Popen of the ComfyUI server, if we started it
        self.q = []                   # event queue (list used as FIFO)
        self.qlock = threading.Lock()
        self.state = "unset"          # unset|setting_up|stopped|starting|running
        self._restart_pending = False
        self._intentional_stop = False

        root.title("B70 Image Generator" + (" v" + version() if version() else ""))
        root.minsize(560, 420)

        pad = {"padx": 8, "pady": 6}

        self.status = tk.StringVar(value="Checking…")
        ttk.Label(root, textvariable=self.status, font=("", 11, "bold")).pack(
            anchor="w", **pad)

        btns = ttk.Frame(root)
        btns.pack(fill="x", **pad)
        self.b_setup = ttk.Button(btns, text="Set up (download model & install)",
                                  command=self.on_setup)
        self.b_start = ttk.Button(btns, text="▶ Start & Open", command=self.on_start)
        self.b_restart = ttk.Button(btns, text="⟳ Restart", command=self.on_restart)
        self.b_stop = ttk.Button(btns, text="■ Shut down", command=self.on_stop)
        for b in (self.b_setup, self.b_start, self.b_restart, self.b_stop):
            b.pack(side="left", padx=4)

        ttk.Label(root, text="Log:").pack(anchor="w", padx=8)
        self.log = scrolledtext.ScrolledText(root, height=16, wrap="word",
                                             state="disabled")
        self.log.pack(fill="both", expand=True, padx=8, pady=(0, 8))

        tip = (
            "Image Prompt Finder: list the prompts and settings saved inside "
            "generated images, to browse and reuse past prompts.\n"
            "Install:  sudo apt update && sudo apt install dio-img-prompts"
            "      Open:  menu entry “Image Prompt Finder”, or run: dio-img-prompts"
        )
        ttk.Label(root, text=tip, justify="left", wraplength=760,
                  foreground="#555").pack(anchor="w", padx=8, pady=(0, 8))

        root.protocol("WM_DELETE_WINDOW", self.on_close)
        self.refresh_state()
        self.root.after(200, self._drain)

    # ---- event queue plumbing -------------------------------------------
    def emit(self, kind, payload=None):
        with self.qlock:
            self.q.append((kind, payload))

    def _drain(self):
        with self.qlock:
            events, self.q = self.q, []
        for kind, payload in events:
            if kind == "log":
                self._append(payload)
            elif kind == "up":
                self.set_state("running")
                self._append("Server is up; opening browser…")
                try:
                    webbrowser.open(URL)
                except Exception as e:
                    self._append("Could not open browser: %s" % e)
            elif kind == "exited":
                if not self._intentional_stop:
                    self._append("Server exited (code %s)." % payload)
                self.proc = None
                self.set_state("stopped")
            elif kind == "stopped":
                self.proc = None
                self.set_state("stopped")
                if self._restart_pending:
                    self._restart_pending = False
                    self.on_start()
            elif kind == "setup_done":
                ok = payload
                self._append("Setup finished." if ok else "Setup FAILED — see log.")
                self.refresh_state()
        self.root.after(200, self._drain)

    def _append(self, text):
        self.log.configure(state="normal")
        self.log.insert("end", text.rstrip() + "\n")
        self.log.see("end")
        self.log.configure(state="disabled")

    # ---- state / button gating ------------------------------------------
    def refresh_state(self):
        if not is_set_up():
            self.set_state("unset")
        elif self.proc and self.proc.poll() is None:
            self.set_state("running")
        elif port_open():
            self.set_state("running_external")
        else:
            self.set_state("stopped")

    def set_state(self, state):
        self.state = state
        msgs = {
            "unset": "Not set up yet — click “Set up” to download the model and install.",
            "setting_up": "Setting up… this downloads ~17GB, please wait.",
            "stopped": "Ready. Server is stopped.",
            "starting": "Starting server…",
            "running": "Running — %s" % URL,
            "running_external": "Running (started outside this window) — %s" % URL,
        }
        self.status.set(msgs.get(state, state))

        def en(b, on):
            b.state(["!disabled"] if on else ["disabled"])

        en(self.b_setup, state in ("unset", "stopped"))
        en(self.b_start, state == "stopped")
        en(self.b_restart, state in ("running", "running_external"))
        en(self.b_stop, state in ("running", "running_external"))

    # ---- actions ---------------------------------------------------------
    def on_setup(self):
        self.set_state("setting_up")
        for b in (self.b_setup, self.b_start, self.b_restart, self.b_stop):
            b.state(["disabled"])
        self._append("Running b70-imagegen-setup …")

        def worker():
            try:
                p = subprocess.Popen(
                    [SETUP], stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                    text=True, bufsize=1)
                for line in p.stdout:
                    self.emit("log", line)
                p.wait()
                self.emit("setup_done", p.returncode == 0)
            except Exception as e:
                self.emit("log", "setup error: %s" % e)
                self.emit("setup_done", False)

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

    def on_start(self):
        if port_open():
            self._append("Port %d already in use; opening browser." % PORT)
            self.set_state("running_external")
            try:
                webbrowser.open(URL)
            except Exception:
                pass
            return
        self.set_state("starting")
        self._intentional_stop = False
        self._append("Launching ComfyUI server…")
        try:
            self.proc = subprocess.Popen(
                [LAUNCHER], stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                text=True, bufsize=1, start_new_session=True)
        except Exception as e:
            self._append("Failed to launch: %s" % e)
            self.set_state("stopped")
            return
        threading.Thread(target=self._read_server, daemon=True).start()
        threading.Thread(target=self._wait_up, daemon=True).start()

    def _read_server(self):
        p = self.proc
        if not p:
            return
        for line in p.stdout:
            self.emit("log", line)
        self.emit("exited", p.poll())

    def _wait_up(self):
        for _ in range(180):
            if port_open():
                self.emit("up")
                return
            if not self.proc or self.proc.poll() is not None:
                return
            time.sleep(1)
        self.emit("log", "Timed out waiting for the server to come up.")

    def on_stop(self, then_restart=False):
        self._restart_pending = then_restart
        self._intentional_stop = True
        p = self.proc
        if not p or p.poll() is not None:
            self.emit("stopped")
            return
        self._append("Shutting down server…")

        def worker():
            try:
                os.killpg(os.getpgid(p.pid), signal.SIGINT)
            except Exception:
                pass
            if not self._wait_exit(p, 12):
                try:
                    os.killpg(os.getpgid(p.pid), signal.SIGTERM)
                except Exception:
                    pass
                if not self._wait_exit(p, 6):
                    try:
                        os.killpg(os.getpgid(p.pid), signal.SIGKILL)
                    except Exception:
                        pass
                    self._wait_exit(p, 4)
            self.emit("stopped")

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

    @staticmethod
    def _wait_exit(p, secs):
        for _ in range(int(secs * 2)):
            if p.poll() is not None:
                return True
            time.sleep(0.5)
        return p.poll() is not None

    def on_restart(self):
        self._append("Restarting…")
        self.on_stop(then_restart=True)

    def on_close(self):
        if self.proc and self.proc.poll() is None:
            try:
                os.killpg(os.getpgid(self.proc.pid), signal.SIGINT)
            except Exception:
                pass
        self.root.destroy()


def main():
    if not _HAVE_TK:
        sys.stderr.write(
            "python3-tk is not installed; cannot open the GUI.\n"
            "Install it with: sudo apt install python3-tk\n")
        return 1
    root = tk.Tk()
    App(root)
    root.mainloop()
    return 0


if __name__ == "__main__":
    if "--selftest" in sys.argv:
        sys.exit(selftest())
    sys.exit(main())
