#!/usr/bin/env python3
# -*- coding: utf-8 -*-
r"""
Ramble → Generate (PlantUML + optional images, locked-field context, per-field criteria)

What’s new:
- Locked fields feed back into the prompt as authoritative context.
- Summary is just another field; all fields can have custom criteria.
- Fields are always editable; “Lock” checkbox prevents overwrite and highlights the field.
- “How to use” toggle with simple instructions.
- Soft colors; Ramble input visually stands out.
- Scrollable main content, Generate under ramble, Submit bottom-right returns JSON.

CLI examples
------------
# Mock provider (no external calls)
python3 ramble.py --fields Title Summary "Problem it solves" --criteria '{"Summary":"<=2 sentences","Title":"camelCase, no spaces, <=20 chars"}'

# Claude CLI
python3 ramble.py \
  --provider=claude \
  --claude-cmd=/home/rob/.npm-global/bin/claude \
  --fields Title Summary "Problem it solves" "Brief overview" \
  --criteria '{"Summary":"<= 2 sentences","Title":"camelCase, no spaces, <= 20 characters"}' \
  --tail 6000 --debug

Requirements
------------
- PlantUML CLI in PATH           -> sudo apt install plantuml
- PySide6 (or PyQt5)             -> pip install PySide6
- requests (only for images)     -> pip install requests
- Stability key (optional)       -> STABILITY_API_KEY
- Pexels key (optional)          -> PEXELS_API_KEY
"""

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

# ── Qt (PySide6 preferred; PyQt5 fallback) ────────────────────────────────────
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
    )
    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
    )
    QT_LIB = "PyQt5"


# ── Provider Protocol ─────────────────────────────────────────────────────────
@runtime_checkable
class RambleProvider(Protocol):
    """
    generate() must return:
      - summary: str
      - fields: Dict[str, str]
      - uml_blocks: List[Tuple[str, Optional[bytes]]]
      - image_descriptions: List[str]
    """
    def generate(
        self,
        *,
        prompt: str,
        ramble_text: str,
        fields: List[str],
        field_criteria: Dict[str, str],
        locked_context: Dict[str, str],
    ) -> Dict[str, Any]:
        ...


# ── Mock Provider ─────────────────────────────────────────────────────────────
class MockProvider:
    def generate(
        self, *, prompt: str, ramble_text: str, fields: List[str],
        field_criteria: Dict[str, str], locked_context: Dict[str, 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 = {}
        for f in fields:
            crit = field_criteria.get(f, "").strip()
            suffix = f" [criteria: {crit}]" if crit else ""
            field_map[f] = f"{f}: Derived from ramble ({len(words)} words).{suffix}"
        uml_blocks = [
            ("@startuml\nactor User\nUser -> System: Ramble\nSystem -> LLM: Generate\nLLM --> System: Summary/Fields/UML\nSystem --> User: Updates\n@enduml", None)
        ]
        image_descriptions = [
            "Illustrate the core actor interacting with the system.",
            "Abstract icon that represents the idea’s domain."
        ]
        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], field_criteria: Dict[str,str], locked_context: Dict[str,str]
    ) -> str:
        fields_yaml = "\n".join([f'- "{f}"' for f in fields])
        criteria_yaml = "\n".join([f'- {name}: {field_criteria[name]}' for name in fields if field_criteria.get(name)])
        locked_yaml = "\n".join([f'- {k}: {locked_context[k]}' for k in locked_context.keys()]) if locked_context else ""
        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.

        Goal:
        - Read the user's "ramble" and synthesize structured outputs.

        Required JSON keys:
        - "summary": string
        - "fields": object with these keys: {fields_yaml}
        - "uml_blocks": array of objects: {{"uml_text": string, "png_base64": null}}
        - "image_descriptions": array of short strings

        Per-field criteria (enforce tightly; if ambiguous, choose the most useful interpretation):
        {criteria_yaml if criteria_yaml else "- (no special criteria provided)"}

        Guidance:
        - The PlantUML diagram must reflect concrete entities/flows drawn from the ramble and any locked fields.
        - Use 3–7 nodes where possible; prefer meaningful arrow labels.
        - Image descriptions must be specific to this idea (avoid generic phrasing).

        Authoritative context from previously LOCKED fields (treat as ground truth if present):
        {locked_yaml if locked_yaml else "- (none locked yet)"}

        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],
        field_criteria: Dict[str, str], locked_context: Dict[str, str]
    ) -> Dict[str, Any]:
        p = self._build_prompt(
            user_prompt=prompt, ramble_text=ramble_text,
            fields=fields, field_criteria=field_criteria, locked_context=locked_context
        )
        try:
            out = self._run_once(p, timeout=self.timeout_s)
        except subprocess.TimeoutExpired:
            # Retry with a smaller tail but still honoring tail_chars if set
            rt = ramble_text[-min(self.tail_chars or 3000, 3000):]
            self._log("Retrying with smaller prompt…")
            shorter = self._build_prompt(
                user_prompt=prompt, ramble_text=rt,
                fields=fields, field_criteria=field_criteria, locked_context=locked_context
            )
            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)


# ── UI helpers ─────────────────────────────────────────────────────────────────
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("Start here. Ramble about your idea…")
        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(11)
        self.setFont(font)
        self.setFixedHeight(190)
        # Soft highlight
        self.setStyleSheet("""
            QPlainTextEdit {
                background:#f4fbff; border:2px solid #9ed6ff; border-radius:8px; padding:8px;
            }
        """)
    def wheelEvent(self, event): 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()


# ── Images (optional backends) ────────────────────────────────────────────────
def have_requests_or_raise():
    if requests is None:
        raise RuntimeError("The 'requests' library is required for image fetching. 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 ──────────────────────────────────────────────────────────────────
def ensure_plantuml_present():
    exe = shutil.which("plantuml")
    if not exe: raise RuntimeError("plantuml not found in PATH. 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


# ── Dialog ────────────────────────────────────────────────────────────────────
@dataclass
class RambleResult:
    summary: str
    fields: Dict[str, str]

class RambleDialog(QDialog):
    def __init__(
        self,
        *,
        prompt: str,
        fields: List[str],
        field_criteria: Optional[Dict[str,str]] = None,
        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(1120, 760)

        self.provider = provider or cast(RambleProvider, MockProvider())
        self._prompt_text = prompt
        self._fields = fields[:] if fields else ["Summary"]
        if "Summary" not in self._fields:
            self._fields.insert(0, "Summary")
        self._criteria = field_criteria or {}
        # sensible default for Summary if not provided
        self._criteria.setdefault("Summary", "<= 2 sentences")

        self._hints = hints or [
            "What is it called?",
            "Who benefits most?",
            "What problem does it solve?",
            "What would success look like?",
        ]
        self.enable_stability = enable_stability
        self.enable_pexels = enable_pexels and not enable_stability

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

        # Track per-field lock states
        self.field_lock_boxes: Dict[str, QCheckBox] = {}
        self.field_outputs: Dict[str, QTextEdit] = {}

        # Layout: scrollable content + fixed footer
        outer = QVBoxLayout(self)

        self.scroll = QScrollArea(); self.scroll.setWidgetResizable(True)
        content = QWidget(); self.scroll.setWidget(content)
        outer.addWidget(self.scroll, 1)

        footer = QHBoxLayout()
        # “How to use” toggle
        self.help_btn = QPushButton("How to use ▾"); self.help_btn.setCheckable(True); self.help_btn.setChecked(False)
        self.help_btn.clicked.connect(self._toggle_help)  # type: ignore
        footer.addWidget(self.help_btn)
        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)

        grid = QGridLayout(content)
        grid.setHorizontalSpacing(16); grid.setVerticalSpacing(12)

        # Title & hint
        title = QLabel("<b>Ramble about your idea. Fields will fill themselves.</b>")
        title.setWordWrap(True)
        grid.addWidget(title, 0, 0, 1, 2)
        self.hint_label = FadeLabel(); self.hint_label.setText(self._hints[0])
        self.hint_label.setStyleSheet("color:#666; font-style:italic;")
        grid.addWidget(self.hint_label, 1, 0, 1, 2)

        # Left column: Ramble → Generate → UML → Images
        left_col = QVBoxLayout()

        self.ramble_edit = FadingRambleEdit(max_blocks=900)
        left_col.addWidget(self.ramble_edit)

        gen_row = QHBoxLayout()
        self.generate_btn = QPushButton("Generate / Update")
        self.generate_btn.setStyleSheet("QPushButton { background:#e6f4ff; border:1px solid #9ed6ff; padding:6px 12px; }")
        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 (image only)
        uml_group = QGroupBox("Diagram (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("Images / Descriptions"); img_v = QVBoxLayout(img_group)
        self.img_desc_text = QTextEdit(); self.img_desc_text.setReadOnly(True); self.img_desc_text.setMinimumHeight(80)
        self.img_desc_text.setPlaceholderText("Descriptions of suggested images…")
        img_v.addWidget(self.img_desc_text)
        left_col.addWidget(img_group)

        # Right column: Fields
        right_col = QVBoxLayout()

        fields_group = QGroupBox("Fields")
        fields_group.setStyleSheet("QGroupBox { background:#fafafa; border:1px solid #ddd; border-radius:8px; margin-top:6px; }")
        fg_form = QFormLayout(fields_group)

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

        for f in self._fields:
            row = QHBoxLayout()
            te = QTextEdit(); te.setReadOnly(False); te.setMinimumHeight(64)
            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

            # show criteria hint (subtle)
            crit = self._criteria.get(f, "")
            if crit:
                hint = QLabel(f"<span style='color:#777; font-size:11px'>({crit})</span>")
            else:
                hint = QLabel("")
            row.addWidget(te, 1); row.addWidget(lock)
            fg_form.addRow(QLabel(f"{f}:"), row)
            fg_form.addRow(hint)
            apply_lock_style(te, False)

        right_col.addWidget(fields_group, 1)

        # Inline help panel (hidden by default)
        self.help_panel = QLabel(
            "How to use:\n"
            "• Start in the blue box above and just ramble about your idea.\n"
            "• Click “Generate / Update” to fill fields automatically.\n"
            "• Not happy? Add more detail to your ramble or edit fields directly.\n"
            "• Love a field? Tick “Lock” so it won’t change next time. Locked fields also help guide the next generation.\n"
            "• When satisfied, click Submit to output your structured JSON."
        )
        self.help_panel.setWordWrap(True)
        self.help_panel.setVisible(False)
        self.help_panel.setStyleSheet("background:#f8fbff; border:1px dashed #bcdcff; padding:8px; border-radius:6px;")
        right_col.addWidget(self.help_panel)

        # Status
        self.status_label = QLabel(""); self.status_label.setStyleSheet("color:#777;")
        right_col.addWidget(self.status_label)

        # Place columns
        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 _toggle_help(self):
        self.help_panel.setVisible(self.help_btn.isChecked())
        self.help_btn.setText("How to use ▴" if self.help_btn.isChecked() else "How to use ▾")

    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])

    @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

        # Collect locked fields to feed back as authoritative context
        locked_context: Dict[str, str] = {}
        for name, box in self.field_lock_boxes.items():
            if box.isChecked():
                locked_context[name] = self.field_outputs[name].toPlainText().strip()

        payload = {
            "prompt": self._prompt_text,
            "ramble_text": ramble_text,
            "fields": self._fields,
            "field_criteria": self._criteria,
            "locked_context": locked_context,
        }

        self._set_busy(True, "Generating…")
        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]):
        # Respect locks
        new_fields: Dict[str, str] = data.get("fields", {})
        for name, widget in self.field_outputs.items():
            if self.field_lock_boxes[name].isChecked():
                continue
            widget.setPlainText(new_fields.get(name, ""))

        # Also update Summary if present and not locked
        if "Summary" in self.field_outputs and not self.field_lock_boxes["Summary"].isChecked():
            # If provider also returns top-level "summary", prefer that
            s = data.get("summary", "") or new_fields.get("Summary", "")
            if s: self.field_outputs["Summary"].setPlainText(s)

        # UML render (first diagram)
        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

        # 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._set_busy(False, "Ready.")

    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):
        out_fields = {k: self.field_outputs[k].toPlainText() for k in self.field_outputs.keys()}
        # prefer explicit Summary field if present
        summary = out_fields.get("Summary", "").strip()
        self.result = RambleResult(summary=summary, fields=out_fields)
        self.accept()


# ── Public API ────────────────────────────────────────────────────────────────
def open_ramble_dialog(
    *,
    prompt: str,
    fields: List[str],
    field_criteria: Optional[Dict[str, str]] = None,
    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,
        field_criteria=field_criteria,
        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,
        }

    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=["Summary","Title","Intent","ProblemItSolves","BriefOverview"])
    p.add_argument("--criteria", default="", help="JSON mapping of field -> criteria")
    p.add_argument("--hints", default="", help="JSON list of hint strings")
    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()

    # Field criteria parsing
    criteria: Dict[str, str] = {}
    if args.criteria.strip():
        try:
            criteria = json.loads(args.criteria)
        except Exception as e:
            print(f"[WARN] Could not parse --criteria JSON: {e}", file=sys.stderr)

    # Provider
    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
    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. pip install requests", file=sys.stderr)
        sys.exit(3)

    # Parse JSON args (tolerate empty/invalid)
    try:
        criteria = json.loads(args.criteria) if args.criteria else {}
        if not isinstance(criteria, dict): criteria = {}
    except Exception:
        criteria = {}
    try:
        hints = json.loads(args.hints) if args.hints else None
        if hints is not None and not isinstance(hints, list): hints = None
    except Exception:
        hints = None

    demo = open_ramble_dialog(
        prompt=args.prompt,
        fields=args.fields,
        field_criteria=criteria,
        hints = hints,
        provider=provider,
        enable_stability=args.stability,
        enable_pexels=args.pexels,
    )
    if demo:
        print(json.dumps(demo, ensure_ascii=False, indent=2))