cli_entrypoint.py

"""
cli_entrypoint.py — CLI front-end for the SDXL batch queue.

Usage
─────
    python cli_entrypoint.py --workload workload.json
    python cli_entrypoint.py --workload workload.json \\
                             --models-cfg models_cfg.json \\
                             --nodes-cfg  nodes_cfg.json

JSON config schemas
───────────────────
  --models-cfg  →  SDXLModelsCfg fields  (see models_cfg_template.json)
  --nodes-cfg   →  { "latent": {...}, "sampler": {...},
                     "ipadapter": {...}, "upscale": {...}, "rembg": {...} }
                   (see nodes_cfg_template.json)
                   All keys are optional; missing keys use built-in defaults.
"""

from __future__ import annotations

import argparse
import json
import os
import signal
import sys
import threading
import time
from pathlib import Path

from comfy_script.runtime import load

load()

src_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "src")
if src_path not in sys.path:
    sys.path.insert(0, src_path)

from src.configs import (
    ConfigRegistry,
    IPAdapterConfig,
    KSamplerConfig,
    LatentConfig,
    RembgConfig,
    UpscaleConfig,
)
from src.lora_models import SomeLORA, union_lora_list_adapter
from src.models_cfg import SDXLModelsCfg
from src.queue import GroupStats, QueueReporter, SilentReporter, process_queue


# ── ANSI helpers ──────────────────────────────────────────────────────────────

_RESET = "\033[0m"
_BOLD = "\033[1m"
_GREEN = "\033[32m"
_RED = "\033[31m"
_YELLOW = "\033[33m"
_CYAN = "\033[36m"
_DIM = "\033[2m"


def _g(s: str) -> str:
    return f"{_GREEN}{_BOLD}{s}{_RESET}"  # bold green


def _r(s: str) -> str:
    return f"{_RED}{_BOLD}{s}{_RESET}"  # bold red


def _y(s: str) -> str:
    return f"{_YELLOW}{s}{_RESET}"  # yellow


def _c(s: str) -> str:
    return f"{_CYAN}{s}{_RESET}"  # cyan


def _b(s: str) -> str:
    return f"{_BOLD}{s}{_RESET}"  # bold


def _d(s: str) -> str:
    return f"{_DIM}{s}{_RESET}"  # dim


def _stats_line(stats: GroupStats) -> str:
    return (
        f"Всього {_b(str(stats.total))}, "
        f"згенеровано {_g(str(stats.done))}, "
        f"залишилось {_c(str(stats.remaining))}, "
        f"помилка {_r(str(stats.errors))}"
    )


def _is_jupyter() -> bool:
    try:
        shell = get_ipython().__class__.__name__  # type: ignore[name-defined]
        return shell in ("ZMQInteractiveShell", "TerminalInteractiveShell")
    except NameError:
        return False


# ── Spinner / pipeline display ────────────────────────────────────────────────

_SPINNER_CHARS = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏"


class PipelineDisplay:
    """
    Shows an in-place animated spinner with the pipeline stage list
    while a workflow is running.

    Terminal mode  — uses \\r to update a single line in place.
    Jupyter mode   — prints static lines (no \\r tricks needed).
    """

    def __init__(self, img_name: str, stages: list[str], jupyter: bool = False):
        self._name = img_name
        self._stages = stages
        self._jupyter = jupyter
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None

    # ── Stage formatting ──────────────────────────────────────────────────────

    def _pending_pipeline(self) -> str:
        return _d(" → ".join(self._stages))

    def _done_pipeline(self, ok: bool) -> str:
        icon = "✓" if ok else "✗"
        color = _g if ok else _r
        parts = [color(f"{icon} {s}") for s in self._stages]
        return f" {_g('→') if ok else _r('→')} ".join(parts)

    # ── Spinner thread ────────────────────────────────────────────────────────

    def _spin(self) -> None:
        idx = 0
        stages_str = self._pending_pipeline()
        start = time.time()

        while not self._stop.is_set():
            char = _SPINNER_CHARS[idx % len(_SPINNER_CHARS)]
            elapsed = int(time.time() - start)
            line = f"  {_c(char)} {stages_str}  {_d(f'[{elapsed:03d}s]')}"
            print(f"\r{line}        ", end="", flush=True)
            idx += 1
            time.sleep(0.1)

        # Clear the spinner line so finish() can overwrite cleanly.
        print("\r" + " " * 80 + "\r", end="", flush=True)

    # ── Public API ────────────────────────────────────────────────────────────

    def start(self) -> None:
        if self._jupyter:
            print(f"    Конвеєр: {self._pending_pipeline()}", flush=True)
            print(f"    Генерація...", flush=True)
        else:
            self._thread = threading.Thread(target=self._spin, daemon=True)
            self._thread.start()

    def finish(self, duration: float, ok: bool = True) -> None:
        self._stop.set()
        if self._thread:
            self._thread.join()

        pipeline = self._done_pipeline(ok)
        if self._jupyter:
            print(f"    Конвеєр: {pipeline}")
        else:
            print(f"  {pipeline}  {_d(f'[{duration:.1f}s]')}", flush=True)


# ── Reporter implementation ───────────────────────────────────────────────────


class CliReporter:
    """
    Translates queue lifecycle events into formatted terminal output.
    No tqdm — raw Python prints + ANSI codes only.
    """

    def __init__(self) -> None:
        self._jupyter = _is_jupyter()
        self._display: PipelineDisplay | None = None
        self._sep = _d("─" * 60)

    # ── Separator ─────────────────────────────────────────────────────────────

    def _line(self) -> None:
        print(self._sep)

    # ── QueueReporter protocol ────────────────────────────────────────────────

    def on_groups_found(self, count: int) -> None:
        print()
        print(f"{_b('Всього знайдено груп:')} {_c(str(count))}")

    def on_group_start(self, name: str, stats: GroupStats) -> None:
        print()
        self._line()
        print(f"{_b('Почато обробку групи:')} {_y(name)}")
        print(f"  {_stats_line(stats)}")

    def on_image_start(self, name: str, stages: list[str]) -> None:
        print()
        print(f"  {_b('Почато створення зображення:')} {name}")
        self._display = PipelineDisplay(name, stages, jupyter=self._jupyter)
        self._display.start()

    def on_image_done(self, name: str, duration: float, stats: GroupStats) -> None:
        if self._display:
            self._display.finish(duration, ok=True)
            self._display = None
        print(f"  {_g('✓')} Завершено: {_b(name)}  {_d(f'({duration:.1f}s)')}")
        self._line()
        print(f"  {_stats_line(stats)}")

    def on_image_error(
        self, name: str, duration: float, exc: Exception, stats: GroupStats
    ) -> None:
        if self._display:
            self._display.finish(duration, ok=False)
            self._display = None
        print(f"  {_r('✗')} Помилка: {_b(name)}  — {_r(str(exc))}")
        self._line()
        print(f"  {_stats_line(stats)}")

    def on_group_done(self, name: str, path: Path, stats: GroupStats) -> None:
        print()
        print(f"{_g('✓')} Генерація групи {_b(name)} завершена")
        print(f"  {_stats_line(stats)}")
        print(f"  Збережено у: {_c(str(path))}")
        self._line()

    def on_fatal(self, message: str) -> None:
        print(f"\n{_r('FATAL:')} {message}", file=sys.stderr)


# ── Graceful interrupt ────────────────────────────────────────────────────────


def _install_signal_handler(stop_flag: list[bool]) -> None:
    """SIGUSR1 → finish current image, then stop."""

    def _handler(*_):
        stop_flag[0] = True
        print(
            f"\n{_y('!' * 60)}\n"
            f"{_y('SIGUSR1 отримано — завершую поточне зображення, після зупинюсь...')}\n"
            f"{_y('!' * 60)}\n",
            flush=True,
        )

    try:
        signal.signal(signal.SIGUSR1, _handler)
        print(_d(f"PID {os.getpid()} | м'яка зупинка: kill -USR1 {os.getpid()}"))
    except (OSError, AttributeError):
        # Windows — SIGUSR1 not available
        pass


# ── Config loaders ────────────────────────────────────────────────────────────

_CFG_CLASSES = {
    "latent": LatentConfig,
    "sampler": KSamplerConfig,
    "ipadapter": IPAdapterConfig,
    "upscale": UpscaleConfig,
    "rembg": RembgConfig,
}


def _load_nodes_cfg(path: str) -> ConfigRegistry:
    """
    Parse nodes_cfg.json into a ConfigRegistry.

    Expected shape (all keys optional):
        {
            "latent":    { "width": 1024, "height": 1024 },
            "sampler":   { "steps": 30, "cfg": 6.5, ... },
            "ipadapter": { "weight": 0.55, ... },
            "upscale":   { "upscale_by": 4.0, ... },
            "rembg":     { "threshold": 0.5, ... }
        }
    """
    raw = json.loads(Path(path).read_text(encoding="utf-8"))
    registry = ConfigRegistry()
    for key, cls in _CFG_CLASSES.items():
        if key in raw:
            registry = registry.with_cfg(key, cls.model_validate(raw[key]))
    return registry


# ── CLI entrypoint ────────────────────────────────────────────────────────────


def main() -> None:
    ap = argparse.ArgumentParser(
        description="Comfyui automation batch generation — CLI",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
examples:
  python cli_entrypoint.py --workload workload.json
  python cli_entrypoint.py --workload workload.json --models-cfg models_cfg.json
  python cli_entrypoint.py --workload workload.json \\
                            --models-cfg models_cfg.json \\
                            --nodes-cfg  nodes_cfg.json \\
                            -o /path/to/export_folder/ \\
                            --lora lora.json 0 1 \\
                            --lora lora2.json \\
                            --lora lora3.json 5 0 
        """,
    )
    ap.add_argument(
        "--workload",
        required=True,
        metavar="FILE",
        help="Path to workload.json",
    )
    ap.add_argument(
        "--models-cfg",
        metavar="FILE",
        help="Path to models_cfg.json  (SDXLModelsCfg fields)",
    )
    ap.add_argument(
        "--nodes-cfg",
        metavar="FILE",
        help=(
            "Path to nodes_cfg.json  "
            "(latent/sampler/ipadapter/upscale/rembg keys, all optional)"
        ),
    )
    ap.add_argument(
        "--output",
        "-o",
        type=Path,
        metavar="DIR",
        help=(
            "Path to export folder, default in DATA_DIRECORY/export or root_save_folder from workfload.json if exist"
        ),
    )
    ap.add_argument(
        "--lora",
        action="append",
        nargs="+",
        metavar=("FILE", "IDXS"),
        help=(
            "Path to lora.json file(s)"
            "file must contain list with LocalLORA or CivitaiLORA model(s)"
            "IDXS is indices of loras in that list"
            "Local LORA (file_name\strength_model\strength_clip\clip_skip\triger_words, triger_words and clip_skip - optional)"
            "Civitai LORA (lora_air\lora_name\strength_model\strength_clip\clip_skip\triger_words; triger_words and clip_skip - optional)"
        ),
    )
    args = ap.parse_args()

    # ── Load configs ──────────────────────────────────────────────────────────

    models_cfg = SDXLModelsCfg.model_validate_json(Path(args.models_cfg).read_text())
    configs = _load_nodes_cfg(args.nodes_cfg) if args.nodes_cfg else ConfigRegistry()

    loras: list[SomeLORA] = []
    if args.lora:
        for entry in args.lora:
            # entry[0] - lora.json file path
            # entry[1:] - lora list indices to apply
            file_path = Path(entry[0])
            try:
                indices = [int(i) for i in entry[1:]]
                if not file_path.exists:
                    print(f"Warning: File {file_path} not found.")
                    continue
                file_content = file_path.read_text()
                loras_from_file = union_lora_list_adapter.validate_json(file_content)
                for i in indices:
                    loras.append(loras_from_file[i])

            except ValueError:
                print(f"Error: indices for {file_path} must be integers.")
            except Exception as e:
                print(f"Error parsing {file_path}: {e}")

    # ── Graceful stop ─────────────────────────────────────────────────────────
    stop_flag: list[bool] = [False]
    _install_signal_handler(stop_flag)

    # ── Run ───────────────────────────────────────────────────────────────────
    reporter = CliReporter()

    process_queue(
        config_path=args.workload,
        models_cfg=models_cfg,
        loras=loras,
        output_folder=args.output,
        configs=configs,
        reporter=reporter,
        stop_flag=stop_flag,
    )

    status = "ПЕРЕРВАНО" if stop_flag[0] else "ЗАВЕРШЕНО"
    print(f"\n{_b(status)}: {Path(args.workload).stem}\n")


if __name__ == "__main__":
    main()