#!/usr/bin/env python3
# -*- coding: utf-8 -*-
r"""
ramble.py — Lean "Ramble → Generate" dialog (PlantUML + optional images)

Changes (per latest request)
- Fields are always editable (no Edit button)
- Lock via checkbox; locked fields are highlighted and not overwritten by Generate
- Hide UML code; only show rendered PlantUML image
- Scrollable main content so long generations fit smaller screens
- "Generate / Update" moved under the Ramble input
- Bottom-right "Submit" button returns JSON (summary + fields) via the public API and prints in CLI demo
"""

from __future__ import annotations
from dataclasses import dataclass
from typing import Dict, List, Optional, Any, Tuple, Protocol, runtime_checkable, Mapping, cast
import os, sys, json, textwrap, base64, re, time, shutil, subprocess, argparse, threading

try:
    import requests  # type: ignore
except ImportError:
    requests = None  # runtime-guarded

# --- Qt Import (PySide6 preferred; fallback to PyQt5) -------------------------
QT_LIB = None
try:
    from PySide6.QtCore import (Qt, QTimer, QPropertyAnimation, QEasingCurve, QObject,
                                QThreadPool, QRunnable, Signal, Slot)
    from PySide6.QtGui import (QFont, QLinearGradient, QPainter, QPalette, QColor, QPixmap, QTextCursor)
    from PySide6.QtWidgets import (
        QApplication, QDialog, QGridLayout, QVBoxLayout, QHBoxLayout, QLabel, QPushButton,
        QPlainTextEdit, QTextEdit, QWidget, QGroupBox, QFormLayout, QSizePolicy,
        QFrame, QMessageBox, QScrollArea, QCheckBox, QSpacerItem
    )
    QT_LIB = "PySide6"
except ImportError:
    from PyQt5.QtCore import (Qt, QTimer, QPropertyAnimation, QEasingCurve, QObject,
                              QThreadPool, QRunnable, pyqtSignal as Signal, pyqtSlot as Slot)
    from PyQt5.QtGui import (QFont, QLinearGradient, QPainter, QPalette, QColor, QPixmap, QTextCursor)
    from PyQt5.QtWidgets import (
        QApplication, QDialog, QGridLayout, QVBoxLayout, QHBoxLayout, QLabel, QPushButton,
        QPlainTextEdit, QTextEdit, QWidget, QGroupBox, QFormLayout, QSizePolicy,
        QFrame, QMessageBox, QScrollArea, QCheckBox, QSpacerItem
    )
    QT_LIB = "PyQt5"


# --- Provider Protocol ---------------------------------------------------------
@runtime_checkable
class RambleProvider(Protocol):
    """Return a dict with: summary, fields, uml_blocks, image_descriptions."""
    def generate(self, *, prompt: str, ramble_text: str, fields: List[str]) -> Dict[str, Any]: ...


# --- Mock Provider (offline/testing) ------------------------------------------
class MockProvider:
    def generate(self, *, prompt: str, ramble_text: str, fields: List[str]) -> Dict[str, Any]:
        words = ramble_text.strip().split()
        cap = min(25, len(words))
        summary = " ".join(words[-cap:]) if words else "(no content yet)"
        summary = (summary[:1].upper() + summary[1:]).rstrip(".") + "."
        field_map = {f: f"{f}: Derived from ramble ({len(words)} words)." for f in fields}
        uml_blocks = [
            ("@startuml\nactor User\nUser -> System: Ramble\nSystem -> LLM: Generate\nLLM --> System: Summary/Fields/UML\nSystem --> User: Updates\n@enduml", None)
        ]
        image_descriptions = [
            "A conceptual sketch representing the core entities mentioned in the ramble.",
            "An abstract icon related to the domain of the idea."
        ]
        return {
            "summary": summary,
            "fields": field_map,
            "uml_blocks": uml_blocks,
            "image_descriptions": image_descriptions,
        }


# --- Claude CLI Provider -------------------------------------------------------
class ClaudeCLIProvider:
    def __init__(
        self,
        cmd: str = "claude",
        extra_args: Optional[List[str]] = None,
        timeout_s: int = 120,
        tail_chars: int = 8000,
        use_arg_p: bool = True,
        debug: bool = False,
        log_path: str = "/tmp/ramble_claude.log",
    ):
        self.cmd = shutil.which(cmd) or cmd
        self.extra_args = extra_args or []
        self.timeout_s = timeout_s
        self.tail_chars = tail_chars
        self.use_arg_p = use_arg_p
        self.debug = debug
        self.log_path = log_path

    def _log(self, msg: str):
        if not self.debug: return
        with open(self.log_path, "a", encoding="utf-8") as f:
            print(f"[{time.strftime('%H:%M:%S')}] {msg}", file=f)

    def _build_prompt(self, *, user_prompt: str, ramble_text: str, fields: List[str]) -> str:
        fields_list = "\n".join(f'- "{f}"' for f in fields)
        ramble_tail = ramble_text[-self.tail_chars:] if self.tail_chars and len(ramble_text) > self.tail_chars else ramble_text
        return textwrap.dedent(f"""\
        You are an assistant that returns ONLY compact JSON. No preamble, no markdown, no code fences.

        Task:
        1) Read the user's 'ramble' text.
        2) Produce a JSON object with exactly these keys:
           "summary": string (≤ 2 sentences)
           "fields": object with these keys: {fields_list}
           "uml_blocks": array of objects: {{"uml_text": string, "png_base64": null}}
           "image_descriptions": array of short strings
        Constraints:
        - Output MUST be valid JSON. No markdown fences. No extra keys.

        Context:
        - Prompt: {user_prompt}
        - Ramble (tail, possibly truncated):
        {ramble_tail}
        """).strip() + "\n"

    @staticmethod
    def _strip_fences(s: str) -> str:
        s = s.strip()
        if s.startswith("```"):
            s = re.sub(r"^```(?:json)?\s*", "", s, flags=re.IGNORECASE)
            s = re.sub(r"\s*```$", "", s)
        return s.strip()

    def _run_once(self, prompt_text: str, timeout: int) -> str:
        if self.use_arg_p:
            argv = [self.cmd, "-p", prompt_text, *self.extra_args]
            stdin = None
        else:
            argv = [self.cmd, *self.extra_args]
            stdin = prompt_text

        self._log(f"argv: {argv}")
        t0 = time.time()
        try:
            proc = subprocess.run(
                argv, input=stdin, capture_output=True, text=True,
                timeout=timeout, check=False,
            )
        except FileNotFoundError as e:
            self._log(f"FileNotFoundError: {e}")
            raise RuntimeError(f"Claude CLI not found at {self.cmd}.") from e
        except subprocess.TimeoutExpired:
            self._log("TimeoutExpired")
            raise

        out = (proc.stdout or "").strip()
        err = (proc.stderr or "").strip()
        self._log(f"rc={proc.returncode} elapsed={time.time()-t0:.2f}s out={len(out)}B err={len(err)}B")
        if proc.returncode != 0:
            raise RuntimeError(f"Claude CLI exited {proc.returncode}:\n{err or out or '(no output)'}")
        return out

    def _normalize(self, raw: Dict[str, Any], fields: List[str]) -> Dict[str, Any]:
        fields_map: Mapping[str, Any] = raw.get("fields", {}) or {}
        uml_objs = raw.get("uml_blocks", []) or []
        image_desc = raw.get("image_descriptions", []) or []
        uml_blocks: List[tuple[str, Optional[bytes]]] = []
        for obj in uml_objs:
            uml_text = (obj or {}).get("uml_text") or ""
            png_b64 = (obj or {}).get("png_base64")
            png_bytes = None
            if isinstance(png_b64, str) and png_b64:
                try:
                    png_bytes = base64.b64decode(png_b64)
                except Exception:
                    png_bytes = None
            uml_blocks.append((uml_text, png_bytes))
        normalized_fields = {name: str(fields_map.get(name, "")) for name in fields}
        return {
            "summary": str(raw.get("summary", "")),
            "fields": normalized_fields,
            "uml_blocks": uml_blocks,
            "image_descriptions": [str(s) for s in image_desc],
        }

    def generate(self, *, prompt: str, ramble_text: str, fields: List[str]) -> Dict[str, Any]:
        p = self._build_prompt(user_prompt=prompt, ramble_text=ramble_text, fields=fields)
        try:
            out = self._run_once(p, timeout=self.timeout_s)
        except subprocess.TimeoutExpired:
            shorter = self._build_prompt(user_prompt=prompt, ramble_text=ramble_text[-3000:], fields=fields)
            out = self._run_once(shorter, timeout=max(45, self.timeout_s // 2))

        txt = self._strip_fences(out)
        try:
            data = json.loads(txt)
        except json.JSONDecodeError:
            words = txt.split()
            summary = " ".join(words[:30]) + ("..." if len(words) > 30 else "")
            data = {
                "summary": summary or "(no content)",
                "fields": {f: f"{f}: {summary}" for f in fields},
                "uml_blocks": [{"uml_text": "@startuml\nactor User\nUser -> System: Ramble\n@enduml", "png_base64": None}],
                "image_descriptions": ["Abstract illustration related to the idea."],
            }
        return self._normalize(data, fields)


# --- Helpers: Fade label + fading ramble edit ---------------------------------
class FadeLabel(QLabel):
    def __init__(self, parent=None):
        super().__init__(parent)
        self._anim = QPropertyAnimation(self, b"windowOpacity")
        self._anim.setDuration(300)
        self._anim.setEasingCurve(QEasingCurve.InOutQuad)
        self.setWindowOpacity(1.0)

    def fade_to_text(self, text: str):
        def set_text():
            self.setText(text)
            try:
                self._anim.finished.disconnect(set_text)  # type: ignore
            except Exception:
                pass
            self._anim.setStartValue(0.0)
            self._anim.setEndValue(1.0)
            self._anim.start()
        self._anim.stop()
        self._anim.setStartValue(1.0)
        self._anim.setEndValue(0.0)
        self._anim.finished.connect(set_text)  # type: ignore
        self._anim.start()


class FadingRambleEdit(QPlainTextEdit):
    def __init__(self, max_blocks: int = 500, parent=None):
        super().__init__(parent)
        self.setPlaceholderText("Type freely… just keep rambling.")
        self.setVerticalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff)
        self.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff)
        self.setFrameStyle(QFrame.StyledPanel | QFrame.Raised)
        self.document().setMaximumBlockCount(max_blocks)
        font = QFont("DejaVu Sans Mono"); font.setPointSize(10)
        self.setFont(font)
        self.setFixedHeight(180)

    def wheelEvent(self, event):  # prevent manual scrolling
        event.ignore()

    def keyPressEvent(self, event):
        super().keyPressEvent(event)
        self.moveCursor(QTextCursor.End)

    def paintEvent(self, e):
        super().paintEvent(e)
        painter = QPainter(self.viewport())
        grad = QLinearGradient(0, 0, 0, 30)
        bg = self.palette().color(QPalette.ColorRole.Base)
        grad.setColorAt(0.0, QColor(bg.red(), bg.green(), bg.blue(), 255))
        grad.setColorAt(1.0, QColor(bg.red(), bg.green(), bg.blue(), 0))
        painter.fillRect(0, 0, self.viewport().width(), 30, grad)
        painter.end()


# --- Image backends ------------------------------------------------------------
def have_requests_or_raise():
    if requests is None:
        raise RuntimeError("The 'requests' library is required for image fetching. Install with: pip install requests")

def fetch_image_stability(prompt: str) -> Optional[bytes]:
    have_requests_or_raise(); assert requests is not None
    api_key = os.getenv("STABILITY_API_KEY", "").strip()
    if not api_key:
        raise RuntimeError("STABILITY_API_KEY is not set")
    url = "https://api.stability.ai/v2beta/stable-image/generate/core"
    headers = {"Authorization": f"Bearer {api_key}", "Accept": "image/png"}
    data = {"prompt": prompt[:1000], "output_format": "png"}
    resp = requests.post(url, headers=headers, data=data, timeout=60)
    if resp.status_code == 200 and resp.content:
        return resp.content
    raise RuntimeError(f"Stability API error: {resp.status_code} {resp.text[:200]}")

def fetch_image_pexels(query: str) -> Optional[bytes]:
    have_requests_or_raise(); assert requests is not None
    api_key = os.getenv("PEXELS_API_KEY", "").strip()
    if not api_key:
        raise RuntimeError("PEXELS_API_KEY is not set")
    s_url = "https://api.pexels.com/v1/search"
    headers = {"Authorization": api_key}
    params = {"query": query, "per_page": 1}
    r = requests.get(s_url, headers=headers, params=params, timeout=30)
    r.raise_for_status()
    data = r.json()
    photos = data.get("photos") or []
    if not photos:
        return None
    src = photos[0].get("src") or {}
    raw_url = src.get("original") or src.get("large") or src.get("medium")
    if not raw_url:
        return None
    img = requests.get(raw_url, timeout=60)
    img.raise_for_status()
    return img.content


# --- PlantUML rendering --------------------------------------------------------
def ensure_plantuml_present():
    exe = shutil.which("plantuml")
    if not exe:
        raise RuntimeError("plantuml not found in PATH. Install it (e.g., sudo apt install plantuml).")
    return exe

def render_plantuml_to_png_bytes(uml_text: str) -> bytes:
    exe = ensure_plantuml_present()
    proc = subprocess.run(
        [exe, "-tpng", "-pipe"],
        input=uml_text.encode("utf-8"),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=False,
    )
    if proc.returncode != 0 or not proc.stdout:
        raise RuntimeError(f"plantuml failed: {proc.stderr.decode('utf-8', 'ignore')[:200]}")
    return proc.stdout


# --- Background worker ---------------------------------------------------------
class GenWorker(QRunnable):
    class Signals(QObject):
        finished = Signal(dict)
        error = Signal(str)
    def __init__(self, provider: RambleProvider, payload: Dict[str, Any]):
        super().__init__()
        self.provider = provider
        self.payload = payload
        self.signals = GenWorker.Signals()
    def run(self):
        try:
            data = self.provider.generate(**self.payload)
            self.signals.finished.emit(data)  # type: ignore
        except Exception as e:
            self.signals.error.emit(str(e))   # type: ignore


# --- Main Dialog ---------------------------------------------------------------
@dataclass
class RambleResult:
    summary: str
    fields: Dict[str, str]
    uml_blocks: List[Tuple[str, Optional[bytes]]]
    image_descriptions: List[str]

class RambleDialog(QDialog):
    def __init__(
        self,
        *,
        prompt: str,
        fields: List[str],
        hints: Optional[List[str]] = None,
        provider: Optional[RambleProvider] = None,
        enable_stability: bool = False,
        enable_pexels: bool = False,
        parent: Optional[QWidget] = None
    ):
        super().__init__(parent)
        self.setWindowTitle("Ramble → Generate")
        self.resize(1100, 750)  # sensible default, scroll keeps it usable

        self.provider = provider or cast(RambleProvider, MockProvider())
        self._prompt_text = prompt
        self._fields = fields[:]
        self._hints = hints or [
            "What is it called?",
            "What is it for?",
            "What type of idea is it?",
            "What problem does it solve?",
            "Who benefits most?",
            "What are success criteria?",
        ]
        self.enable_stability = enable_stability
        self.enable_pexels = enable_pexels and not enable_stability  # stability takes precedence

        self.thread_pool = QThreadPool.globalInstance()
        self.result: Optional[RambleResult] = None

        # Track per-field lock states
        self.field_locked: Dict[str, bool] = {f: False for f in self._fields}

        # --- Outer layout: scrollable content + fixed footer row
        outer = QVBoxLayout(self)

        # Scroll area wraps the main content
        self.scroll = QScrollArea()
        self.scroll.setWidgetResizable(True)
        content = QWidget()
        self.scroll.setWidget(content)
        outer.addWidget(self.scroll, 1)

        # Footer with Submit button (fixed)
        footer = QHBoxLayout()
        footer.addStretch(1)
        self.submit_btn = QPushButton("Submit")
        self.submit_btn.clicked.connect(self.on_submit)  # type: ignore
        footer.addWidget(self.submit_btn)
        outer.addLayout(footer)

        # === Inside the scrollable content ===
        grid = QGridLayout(content)
        grid.setHorizontalSpacing(16)
        grid.setVerticalSpacing(12)

        # Top prompt + hint
        prompt_label = QLabel(f"<b>{self._prompt_text}</b>"); prompt_label.setWordWrap(True)
        hint_label = FadeLabel(); hint_label.setText(self._hints[0])
        hint_label.setStyleSheet("color:#666; font-style:italic;")
        grid.addWidget(prompt_label, 0, 0, 1, 2)
        grid.addWidget(hint_label, 1, 0, 1, 2)
        self.hint_label = hint_label

        # Left column -----------------------------------------------------------
        left_col = QVBoxLayout()

        # Ramble input
        self.ramble_edit = FadingRambleEdit(max_blocks=600)
        left_col.addWidget(self.ramble_edit)

        # Generate button right under input
        gen_row = QHBoxLayout()
        self.generate_btn = QPushButton("Generate / Update")
        self.generate_btn.clicked.connect(self.on_generate)  # type: ignore
        self.provider_label = QLabel(f"Provider: {type(self.provider).__name__}")
        self.provider_label.setStyleSheet("color:#555;")
        gen_row.addWidget(self.generate_btn)
        gen_row.addStretch(1)
        gen_row.addWidget(self.provider_label)
        left_col.addLayout(gen_row)

        # UML group (image only)
        uml_group = QGroupBox("Generated UML (PlantUML)")
        uml_v = QVBoxLayout(uml_group)
        self.uml_image_label = QLabel()
        self.uml_image_label.setAlignment(Qt.AlignCenter)
        if QT_LIB == "PyQt5":
            self.uml_image_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        else:
            self.uml_image_label.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
        uml_v.addWidget(self.uml_image_label)
        left_col.addWidget(uml_group)

        # Images / descriptions
        img_group = QGroupBox("Generated Images / Descriptions")
        img_v = QVBoxLayout(img_group)
        self.img_desc_text = QTextEdit(); self.img_desc_text.setReadOnly(True); self.img_desc_text.setMinimumHeight(90)
        self.img_desc_text.setPlaceholderText("Descriptions of suggested/related images…")
        img_v.addWidget(self.img_desc_text)

        self.img_scroll_inner = QScrollArea()
        self.img_scroll_inner.setWidgetResizable(True)
        self.img_panel = QWidget()
        self.img_panel_layout = QVBoxLayout(self.img_panel)
        self.img_scroll_inner.setWidget(self.img_panel)
        img_v.addWidget(self.img_scroll_inner)
        left_col.addWidget(img_group)

        # Right column ----------------------------------------------------------
        right_col = QVBoxLayout()

        # Summary (editable)
        summary_group = QGroupBox("Generated Summary")
        sg_v = QVBoxLayout(summary_group)
        self.summary_edit = QTextEdit()
        self.summary_edit.setReadOnly(False)
        self.summary_edit.setMinimumHeight(110)
        self.summary_edit.setPlaceholderText("Summary will appear here…")
        sg_v.addWidget(self.summary_edit)
        right_col.addWidget(summary_group)

        # Fields (editable) with Lock checkbox
        fields_group = QGroupBox("Generated Fields")
        fg_form = QFormLayout(fields_group)
        self.field_outputs: Dict[str, QTextEdit] = {}
        self.field_lock_boxes: Dict[str, QCheckBox] = {}

        def apply_lock_style(widget: QTextEdit, locked: bool):
            if locked:
                widget.setStyleSheet("background-color: #fff7cc;")  # light yellow
            else:
                widget.setStyleSheet("")

        for f in self._fields:
            row = QHBoxLayout()
            te = QTextEdit()
            te.setReadOnly(False)            # always editable
            te.setMinimumHeight(70)
            te.setPlaceholderText(f"{f}…")
            self.field_outputs[f] = te

            lock = QCheckBox("Lock")
            lock.setChecked(False)
            self.field_lock_boxes[f] = lock
            lock.stateChanged.connect(lambda _=0, name=f: apply_lock_style(self.field_outputs[name], self.field_lock_boxes[name].isChecked()))  # type: ignore

            row.addWidget(te, 1)
            row.addWidget(lock)
            fg_form.addRow(f"{f}:", row)
            apply_lock_style(te, False)

        right_col.addWidget(fields_group, 1)

        # Status label at the very bottom of content (inside scroll)
        self.status_label = QLabel(""); self.status_label.setStyleSheet("color:#777;")
        right_col.addWidget(self.status_label)

        # Place columns into the grid
        left_w = QWidget(); left_w.setLayout(left_col)
        right_w = QWidget(); right_w.setLayout(right_col)
        grid.addWidget(left_w, 2, 0)
        grid.addWidget(right_w, 2, 1)

        self._setup_hint_timer()

    def _setup_hint_timer(self):
        self._hint_idx = 0
        self._hint_timer = QTimer(self)
        self._hint_timer.setInterval(3500)
        self._hint_timer.timeout.connect(self._advance_hint)  # type: ignore
        self._hint_timer.start()

    def _advance_hint(self):
        if not self._hints: return
        self._hint_idx = (self._hint_idx + 1) % len(self._hints)
        self.hint_label.fade_to_text(self._hints[self._hint_idx])

    # Actions -------------------------------------------------------------------
    @Slot()
    def on_generate(self):
        ramble_text = self.ramble_edit.toPlainText()
        if not ramble_text.strip():
            QMessageBox.information(self, "Nothing to generate", "Type your thoughts first, then click Generate / Update.")
            return

        self._set_busy(True, "Generating…")
        payload = {"prompt": self._prompt_text, "ramble_text": ramble_text, "fields": self._fields}
        worker = GenWorker(self.provider, payload)
        worker.signals.finished.connect(self._on_generated)  # type: ignore
        worker.signals.error.connect(self._on_gen_error)     # type: ignore
        self.thread_pool.start(worker)

    def _on_generated(self, data: Dict[str, Any]):
        # Summary: do not auto-overwrite if locked? (No lock UI for summary by request; always overwrite but is editable)
        self.summary_edit.setPlainText(data.get("summary", ""))

        # Fields: respect locks
        new_fields: Dict[str, str] = data.get("fields", {})
        for name, widget in self.field_outputs.items():
            if self.field_lock_boxes.get(name, QCheckBox()).isChecked():
                continue  # keep user's text
            widget.setPlainText(new_fields.get(name, ""))

        # UML: render PlantUML if provider didn't supply PNG bytes (show image only)
        uml_blocks = data.get("uml_blocks", [])
        self.uml_image_label.clear()
        for (uml_text, maybe_png) in uml_blocks:
            png_bytes = maybe_png
            if not png_bytes and uml_text:
                try:
                    png_bytes = render_plantuml_to_png_bytes(uml_text)
                except Exception as e:
                    print(f"[UML render] {e}", file=sys.stderr)
                    png_bytes = None
            if png_bytes:
                pix = QPixmap()
                if pix.loadFromData(png_bytes):
                    self.uml_image_label.setPixmap(pix)
                    break  # show first image only

        # Images / descriptions
        img_desc = [str(s) for s in data.get("image_descriptions", [])]
        self.img_desc_text.setPlainText("\n\n".join(f"• {d}" for d in img_desc))
        self._clear_img_panel()

        if img_desc and (self.enable_stability or self.enable_pexels):
            threading.Thread(target=self._fetch_and_show_images, args=(img_desc,), daemon=True).start()

        # Store result (fields reflect current UI, including locks/edits)
        self.result = RambleResult(
            summary=self.summary_edit.toPlainText(),
            fields={k: self.field_outputs[k].toPlainText() for k in self.field_outputs.keys()},
            uml_blocks=uml_blocks,
            image_descriptions=img_desc,
        )
        self._set_busy(False, "Ready.")

    def _fetch_and_show_images(self, descriptions: List[str]):
        images: List[Tuple[str, Optional[bytes], Optional[str]]] = []
        for d in descriptions:
            try:
                if self.enable_stability:
                    img_bytes = fetch_image_stability(d)
                elif self.enable_pexels:
                    img_bytes = fetch_image_pexels(d)
                else:
                    img_bytes = None
                images.append((d, img_bytes, None))
            except Exception as e:
                images.append((d, None, str(e)))

        def apply():
            for (label, img_bytes, err) in images:
                row = QHBoxLayout()
                lbl = QLabel(label); lbl.setWordWrap(True)
                row.addWidget(lbl, 1)
                if img_bytes:
                    pix = QPixmap()
                    if pix.loadFromData(img_bytes):
                        thumb = QLabel()
                        thumb.setPixmap(pix)
                        thumb.setAlignment(Qt.AlignCenter)
                        if QT_LIB == "PyQt5":
                            thumb.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
                        else:
                            thumb.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
                        row.addWidget(thumb, 1)
                elif err:
                    err_lbl = QLabel(f"(image error: {err})"); err_lbl.setStyleSheet("color:#b33;")
                    row.addWidget(err_lbl)
                container = QWidget(); container.setLayout(row)
                self.img_panel_layout.addWidget(container)
            self.img_panel_layout.addStretch(1)
        QTimer.singleShot(0, apply)

    def _clear_img_panel(self):
        lay = self.img_panel_layout
        while lay.count():
            item = lay.takeAt(0)
            w = item.widget()
            if w:
                w.setParent(None)

    def _on_gen_error(self, msg: str):
        self._set_busy(False, "Error.")
        QMessageBox.critical(self, "Generation Error", msg)

    def _set_busy(self, busy: bool, msg: str):
        self.generate_btn.setEnabled(not busy)
        self.submit_btn.setEnabled(not busy)
        self.status_label.setText(msg)

    @Slot()
    def on_submit(self):
        """Collect current summary + fields and return them to caller."""
        self.result = RambleResult(
            summary=self.summary_edit.toPlainText(),
            fields={k: self.field_outputs[k].toPlainText() for k in self.field_outputs.keys()},
            uml_blocks=[],  # you can include PNG bytes if you want to pipe them out too
            image_descriptions=self.img_desc_text.toPlainText().splitlines(),
        )
        self.accept()


# --- Public API ---------------------------------------------------------------
def open_ramble_dialog(
    *,
    prompt: str,
    fields: List[str],
    hints: Optional[List[str]] = None,
    provider: Optional[RambleProvider] = None,
    enable_stability: bool = False,
    enable_pexels: bool = False,
    parent: Optional[QWidget] = None
) -> Optional[Dict[str, Any]]:
    app_created = False
    app = QApplication.instance()
    if app is None:
        app_created = True
        app = QApplication(sys.argv)

    dlg = RambleDialog(
        prompt=prompt,
        fields=fields,
        hints=hints,
        provider=provider,
        enable_stability=enable_stability,
        enable_pexels=enable_pexels,
        parent=parent
    )
    rc = dlg.exec_() if QT_LIB == "PyQt5" else dlg.exec()

    out = None
    if rc == QDialog.Accepted and dlg.result:
        out = {
            "summary": dlg.result.summary,
            "fields": dlg.result.fields,
            # keeping UML/image lists out of the returned JSON by default; add if you need them:
            # "uml_blocks": dlg.result.uml_blocks,
            # "image_descriptions": dlg.result.image_descriptions,
        }

    if app_created:
        app.quit()
    return out


# --- CLI demo ------------------------------------------------------------------
def parse_args():
    p = argparse.ArgumentParser(description="Ramble → Generate (PlantUML + optional images)")
    p.add_argument("--provider", choices=["mock", "claude"], default="mock")
    p.add_argument("--claude-cmd", default="claude", help="Path to claude CLI")
    p.add_argument("--stability", action="store_true", help="Enable Stability AI images (needs STABILITY_API_KEY)")
    p.add_argument("--pexels", action="store_true", help="Enable Pexels images (needs PEXELS_API_KEY); ignored if --stability set")
    p.add_argument("--prompt", default="Explain your new feature idea")
    p.add_argument("--fields", nargs="+", default=["Title", "Intent", "Problem it solves", "Brief overview of logical flow"])
    p.add_argument("--timeout", type=int, default=90)
    p.add_argument("--tail", type=int, default=6000)
    p.add_argument("--debug", action="store_true")
    return p.parse_args()

if __name__ == "__main__":
    args = parse_args()

    if args.provider == "claude":
        provider = cast(RambleProvider, ClaudeCLIProvider(
            cmd=args.claude_cmd,
            use_arg_p=True,
            timeout_s=args.timeout,
            tail_chars=args.tail,
            debug=args.debug,
        ))
    else:
        provider = cast(RambleProvider, MockProvider())

    # Ensure PlantUML exists early (fail fast)
    try:
        ensure_plantuml_present()
    except Exception as e:
        print(f"[FATAL] {e}", file=sys.stderr); sys.exit(2)

    if (args.stability or args.pexels) and requests is None:
        print("[FATAL] 'requests' is required for image backends. Install with: pip install requests", file=sys.stderr)
        sys.exit(3)

    demo = open_ramble_dialog(
        prompt=args.prompt,
        fields=args.fields,
        hints=None,
        provider=provider,
        enable_stability=args.stability,
        enable_pexels=args.pexels,
    )
    if demo:
        # Emit JSON to stdout (standard practice for tool-style dialogs)
        print(json.dumps(demo, ensure_ascii=False, indent=2))