"""probegpt — Textual TUI."""

from __future__ import annotations

import os
from pathlib import Path

import httpx
from textual import work
from textual.app import App, ComposeResult
from textual.binding import Binding
from textual.containers import Container, Horizontal, VerticalScroll
from textual.screen import Screen
from textual.widget import Widget
from textual.widgets import (
    Button,
    Checkbox,
    Footer,
    Header,
    Input,
    Label,
    ProgressBar,
    RadioButton,
    RadioSet,
    RichLog,
    Select,
    Static,
    TabbedContent,
    TabPane,
)

from probegpt.core.config import (
    ALL_STRATEGY_NAMES,
    CONFIG_PATH,
    Config,
    ProviderConfig,
    StrategyConfig,
)
from probegpt.core.discovery import DiscoveryLoop
from probegpt.core.models import DiscoveryResult
from probegpt.output.base import AbstractOutput
from probegpt.output.json_exporter import JsonOutput
from probegpt.providers.factory import create_model

# ---------------------------------------------------------------------------
# OpenRouter helpers
# ---------------------------------------------------------------------------


_VERTEX_FALLBACK: list[tuple[str, str]] = [
    ("gemini-2.5-pro", "gemini-2.5-pro"),
    ("gemini-2.5-flash", "gemini-2.5-flash"),
    ("gemini-2.0-flash", "gemini-2.0-flash"),
    ("gemini-2.0-flash-lite", "gemini-2.0-flash-lite"),
]

_BEDROCK_FALLBACK: list[tuple[str, str]] = [
    ("Claude 3.7 Sonnet", "us.anthropic.claude-3-7-sonnet-20250219-v1:0"),
    ("Claude 3.5 Sonnet v2", "us.anthropic.claude-3-5-sonnet-20241022-v2:0"),
    ("Claude 3.5 Haiku", "us.anthropic.claude-3-5-haiku-20241022-v1:0"),
    ("Nova Pro", "us.amazon.nova-pro-v1:0"),
    ("Nova Lite", "us.amazon.nova-lite-v1:0"),
    ("Nova Micro", "us.amazon.nova-micro-v1:0"),
    ("Llama 3.3 70B", "us.meta.llama3-3-70b-instruct-v1:0"),
    ("Mistral Large 2407", "mistral.mistral-large-2407-v1:0"),
]

_CEREBRAS_FALLBACK: list[tuple[str, str]] = [
    ("Llama 3.3 70B", "llama-3.3-70b"),
    ("Llama 3.1 8B", "llama3.1-8b"),
    ("GPT-OSS 120B", "gpt-oss-120b"),
    ("Qwen 3 235B", "qwen-3-235b-a22b-instruct-2507"),
    ("Qwen 3 32B", "qwen-3-32b"),
]

_MISTRAL_FALLBACK: list[tuple[str, str]] = [
    ("Mistral Large", "mistral-large-latest"),
    ("Mistral Small", "mistral-small-latest"),
    ("Codestral", "codestral-latest"),
]


def _fetch_openrouter_models() -> dict[str, list[str]]:
    """Return {org: [model_id, ...]} from the OpenRouter API. Empty on failure."""
    api_key = os.getenv("OPENROUTER_API_KEY", "")
    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
    try:
        resp = httpx.get(
            "https://openrouter.ai/api/v1/models",
            headers=headers,
            timeout=10,
        )
        resp.raise_for_status()
        by_org: dict[str, list[str]] = {}
        for m in resp.json().get("data", []):
            mid: str = m["id"]
            org = mid.split("/")[0] if "/" in mid else mid
            by_org.setdefault(org, []).append(mid)
        return {org: sorted(ids) for org, ids in sorted(by_org.items())}
    except Exception:
        return {}


def _fetch_vertex_models() -> list[tuple[str, str]]:
    """Return [(label, model_id), ...] from the Vertex AI / Gemini API. Falls back to static list."""
    try:
        import google.auth
        import google.auth.transport.requests

        creds, _ = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
        req = google.auth.transport.requests.Request()
        creds.refresh(req)
        token = creds.token
        resp = httpx.get(
            "https://generativelanguage.googleapis.com/v1beta/models",
            headers={"Authorization": f"Bearer {token}"},
            timeout=10,
        )
        resp.raise_for_status()
        models = []
        for m in resp.json().get("models", []):
            name: str = m.get("name", "")
            # name is like "models/gemini-2.0-flash"
            model_id = name.removeprefix("models/")
            if model_id and "generateContent" in m.get("supportedGenerationMethods", []):
                display = m.get("displayName", model_id)
                models.append((display, model_id))
        return sorted(models, key=lambda x: x[1]) if models else _VERTEX_FALLBACK
    except Exception:
        return _VERTEX_FALLBACK


def _fetch_bedrock_models(region: str, profile: str) -> list[tuple[str, str]]:
    """Return [(label, model_id), ...] from AWS Bedrock. Falls back to static list."""
    try:
        import boto3

        session = boto3.Session(
            region_name=region or "us-east-1",
            profile_name=profile or None,
        )
        client = session.client("bedrock")
        resp = client.list_foundation_models(byOutputModality="TEXT")
        models = []
        for m in resp.get("modelSummaries", []):
            mid: str = m.get("modelId", "")
            name: str = m.get("modelName", mid)
            if mid:
                models.append((name, mid))
        return sorted(models, key=lambda x: x[1]) if models else _BEDROCK_FALLBACK
    except Exception:
        return _BEDROCK_FALLBACK


def _fetch_cerebras_models() -> list[tuple[str, str]]:
    """Return [(label, model_id), ...] from Cerebras API. Falls back to static list."""
    api_key = os.getenv("CEREBRAS_API_KEY", "")
    if not api_key:
        return _CEREBRAS_FALLBACK
    try:
        resp = httpx.get(
            "https://api.cerebras.ai/v1/models",
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=10,
        )
        resp.raise_for_status()
        models = []
        for m in resp.json().get("data", []):
            mid: str = m.get("id", "")
            if mid:
                models.append((mid, mid))
        return sorted(models) if models else _CEREBRAS_FALLBACK
    except Exception:
        return _CEREBRAS_FALLBACK


def _fetch_mistral_models() -> list[tuple[str, str]]:
    """Return [(label, model_id), ...] from Mistral API. Falls back to static list."""
    api_key = os.getenv("MISTRAL_API_KEY", "")
    if not api_key:
        return _MISTRAL_FALLBACK
    try:
        resp = httpx.get(
            "https://api.mistral.ai/v1/models",
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=10,
        )
        resp.raise_for_status()
        models = []
        for m in resp.json().get("data", []):
            mid: str = m.get("id", "")
            if mid:
                models.append((mid, mid))
        return sorted(models) if models else _MISTRAL_FALLBACK
    except Exception:
        return _MISTRAL_FALLBACK


_STRATEGY_DESCRIPTIONS: dict[str, str] = {
    "objective": "Generate probes directly targeting the objective",
    "mutation": "Mutate high-scoring seeds from the pool",
    "encoding": "Obfuscate via base64, ROT13, leetspeak, fragmentation",
    "roleplay": "Wrap in persona, simulation, or game-context frames",
    "persuasion": "Use emotional appeals, urgency, flattery, logical fallacies",
    "linguistic": "Typos, phonetic spelling, double negatives, logic tricks",
    "structural": "Payload splitting, CoT hijack, prefix/suffix injection",
    "distraction": "Benign prefix, topic pivot, attention splitting",
}

# ---------------------------------------------------------------------------
# TextualOutput — streams results into a RichLog + updates progress
# ---------------------------------------------------------------------------


class TextualOutput(AbstractOutput):
    def __init__(
        self,
        app: App,
        log: RichLog,
        progress_bar: ProgressBar,
        status_label: Label,
        total: int,
    ) -> None:
        self._app = app
        self._log = log
        self._progress_bar = progress_bar
        self._status_label = status_label
        self._total = total
        self._count = 0
        self._best_score = 0.0

    def emit(self, result: DiscoveryResult) -> None:
        self._count += 1
        if result.score > self._best_score:
            self._best_score = result.score
        status = (
            f"Iter {result.iteration} · "
            f"Probe {self._count}/{self._total} · "
            f"Best: {self._best_score:.2f}"
        )
        self._app.call_from_thread(self._progress_bar.advance, 1)
        self._app.call_from_thread(self._status_label.update, status)

    def log(self, msg: str) -> None:
        self._app.call_from_thread(self._log.write, msg)

    def update_status(self, msg: str) -> None:
        self._app.call_from_thread(self._status_label.update, msg)

    def finalize(self, results: list[DiscoveryResult]) -> None:
        n = sum(1 for r in results if r.succeeded)
        self._app.call_from_thread(
            self._log.write,
            f"\n[bold]Done — {n} / {len(results)} succeeded[/bold]",
        )
        self._app.call_from_thread(
            self._status_label.update, f"Done · {n}/{len(results)} succeeded"
        )


# ---------------------------------------------------------------------------
# ProviderForm widget — reused for Target / Generator / Judge tabs
# ---------------------------------------------------------------------------


_PROVIDER_LABELS: dict[str, str] = {
    "azure": "azure_openai",
    "openrouter": "openrouter",
    "vertex": "vertex_ai",
    "bedrock": "bedrock",
    "cerebras": "cerebras",
    "mistral": "mistral",
}


class ProviderForm(Widget):
    DEFAULT_CSS = """
    ProviderForm {
        height: auto;
        padding: 1 2;
    }
    ProviderForm Label {
        margin-top: 1;
        color: $text-muted;
    }
    ProviderForm .hidden {
        display: none;
    }
    ProviderForm .hint {
        color: $text-muted;
        text-style: italic;
    }
    """

    def __init__(self, role: str, config: ProviderConfig) -> None:
        super().__init__(id=f"form-{role}")
        self._role = role
        self._config = config

    def _active_key(self) -> str:
        """Return the short provider key for the current config (azure, openrouter, …)."""
        mapping = {v: k for k, v in _PROVIDER_LABELS.items()}
        return mapping.get(self._config.provider_type, "azure")

    def _show(self, key: str) -> str:
        return "" if self._active_key() == key else "hidden"

    def compose(self) -> ComposeResult:
        cfg = self._config
        r = self._role

        with RadioSet(id=f"{r}-provider"):
            yield RadioButton(
                "Azure OpenAI", value=cfg.provider_type == "azure_openai", id=f"{r}-azure"
            )
            yield RadioButton(
                "OpenRouter", value=cfg.provider_type == "openrouter", id=f"{r}-openrouter"
            )
            yield RadioButton(
                "GCP Vertex AI", value=cfg.provider_type == "vertex_ai", id=f"{r}-vertex"
            )
            yield RadioButton(
                "AWS Bedrock", value=cfg.provider_type == "bedrock", id=f"{r}-bedrock"
            )
            yield RadioButton("Cerebras", value=cfg.provider_type == "cerebras", id=f"{r}-cerebras")
            yield RadioButton("Mistral", value=cfg.provider_type == "mistral", id=f"{r}-mistral")

        # ── Azure ──────────────────────────────────────────────────────────────
        with Container(id=f"{r}-azure-fields", classes=self._show("azure")):
            yield Label("Endpoint")
            yield Input(
                value=cfg.azure_endpoint,
                placeholder="https://<resource>.openai.azure.com",
                id=f"{r}-endpoint",
            )
            yield Label("Deployment")
            yield Input(value=cfg.azure_deployment, placeholder="gpt-4o", id=f"{r}-deployment")
            yield Label("API version")
            yield Input(value=cfg.azure_api_version or "2024-02-15-preview", id=f"{r}-api-version")
            yield Static("[dim]Credentials: az login[/dim]", markup=True, classes="hint")

        # ── OpenRouter ─────────────────────────────────────────────────────────
        with Container(id=f"{r}-openrouter-fields", classes=self._show("openrouter")):
            yield Label("Model", id=f"{r}-or-model-label")
            current_or = cfg.openrouter_model
            yield Select(
                [(current_or, current_or)] if current_or else [],
                value=current_or if current_or else Select.NULL,
                allow_blank=True,
                id=f"{r}-or-model",
            )
            yield Static(
                "[dim]Credentials: OPENROUTER_API_KEY in .env[/dim]", markup=True, classes="hint"
            )

        # ── GCP Vertex AI ──────────────────────────────────────────────────────
        with Container(id=f"{r}-vertex-fields", classes=self._show("vertex")):
            yield Label("Model", id=f"{r}-vertex-model-label")
            vertex_cur = cfg.vertex_model or _VERTEX_FALLBACK[0][1]
            yield Select(
                _VERTEX_FALLBACK,
                value=vertex_cur,
                id=f"{r}-vertex-model",
            )
            yield Label("GCP Project ID")
            yield Input(
                value=cfg.vertex_project, placeholder="my-gcp-project", id=f"{r}-vertex-project"
            )
            yield Label("Region")
            yield Input(value=cfg.vertex_region or "us-central1", id=f"{r}-vertex-region")
            yield Static(
                "[dim]Credentials: gcloud auth application-default login[/dim]",
                markup=True,
                classes="hint",
            )

        # ── AWS Bedrock ────────────────────────────────────────────────────────
        with Container(id=f"{r}-bedrock-fields", classes=self._show("bedrock")):
            yield Label("Model", id=f"{r}-bedrock-model-label")
            bedrock_cur = cfg.bedrock_model or _BEDROCK_FALLBACK[0][1]
            yield Select(_BEDROCK_FALLBACK, value=bedrock_cur, id=f"{r}-bedrock-model")
            yield Label("Region")
            yield Input(value=cfg.bedrock_region or "us-east-1", id=f"{r}-bedrock-region")
            yield Label("AWS Profile (optional)")
            yield Input(value=cfg.bedrock_profile, placeholder="default", id=f"{r}-bedrock-profile")
            yield Static("[dim]Credentials: aws configure[/dim]", markup=True, classes="hint")

        # ── Cerebras ───────────────────────────────────────────────────────────
        with Container(id=f"{r}-cerebras-fields", classes=self._show("cerebras")):
            yield Label("Model", id=f"{r}-cerebras-model-label")
            cerebras_cur = cfg.cerebras_model or _CEREBRAS_FALLBACK[0][1]
            yield Select(_CEREBRAS_FALLBACK, value=cerebras_cur, id=f"{r}-cerebras-model")
            yield Static(
                "[dim]Credentials: CEREBRAS_API_KEY in .env[/dim]", markup=True, classes="hint"
            )

        # ── Mistral ────────────────────────────────────────────────────────────
        with Container(id=f"{r}-mistral-fields", classes=self._show("mistral")):
            yield Label("Model", id=f"{r}-mistral-model-label")
            mistral_cur = cfg.mistral_model or _MISTRAL_FALLBACK[0][1]
            yield Select(_MISTRAL_FALLBACK, value=mistral_cur, id=f"{r}-mistral-model")
            yield Static(
                "[dim]Credentials: MISTRAL_API_KEY in .env[/dim]", markup=True, classes="hint"
            )

    def on_mount(self) -> None:
        key = self._active_key()
        self._load_models(key)

    def on_radio_set_changed(self, event: RadioSet.Changed) -> None:
        if event.pressed is None:
            return
        btn_id = event.pressed.id or ""
        key = btn_id.removeprefix(f"{self._role}-")
        for k in _PROVIDER_LABELS:
            self.query_one(f"#{self._role}-{k}-fields").set_class(k != key, "hidden")
        self._load_models(key)

    @work(thread=True)
    def _load_models(self, key: str) -> None:  # noqa: PLR0912
        r = self._role
        if key == "openrouter":
            by_org = _fetch_openrouter_models()
            options: list[tuple[str, str]] = [(mid, mid) for ids in by_org.values() for mid in ids]
            current = self._config.openrouter_model or ""
            self.app.call_from_thread(
                self._update_select, f"{r}-or-model", options, current, f"{r}-or-model-label"
            )
        elif key == "vertex":
            options = _fetch_vertex_models()
            current = self._config.vertex_model or ""
            self.app.call_from_thread(
                self._update_select,
                f"{r}-vertex-model",
                options,
                current,
                f"{r}-vertex-model-label",
            )
        elif key == "bedrock":
            region = self._config.bedrock_region or "us-east-1"
            profile = self._config.bedrock_profile or ""
            options = _fetch_bedrock_models(region, profile)
            current = self._config.bedrock_model or ""
            self.app.call_from_thread(
                self._update_select,
                f"{r}-bedrock-model",
                options,
                current,
                f"{r}-bedrock-model-label",
            )
        elif key == "cerebras":
            options = _fetch_cerebras_models()
            current = self._config.cerebras_model or ""
            self.app.call_from_thread(
                self._update_select,
                f"{r}-cerebras-model",
                options,
                current,
                f"{r}-cerebras-model-label",
            )
        elif key == "mistral":
            options = _fetch_mistral_models()
            current = self._config.mistral_model or ""
            self.app.call_from_thread(
                self._update_select,
                f"{r}-mistral-model",
                options,
                current,
                f"{r}-mistral-model-label",
            )

    def _update_select(
        self,
        widget_id: str,
        options: list[tuple[str, str]],
        current: str,
        label_id: str,
    ) -> None:
        """Update a Select widget with new options, preserving current value if valid."""
        select = self.query_one(f"#{widget_id}", Select)
        select.set_options(options)
        values = [v for _, v in options]
        if current and current in values:
            select.value = current
        elif values:
            select.value = values[0]
        label = self.query_one(f"#{label_id}", Label)
        label.update(f"Model  ({len(options)} available)")

    def _sel(self, widget_id: str) -> str:
        sel = self.query_one(f"#{widget_id}", Select)
        return "" if sel.is_blank() else str(sel.value)

    def _inp(self, widget_id: str) -> str:
        return self.query_one(f"#{widget_id}", Input).value

    def get_config(self) -> ProviderConfig:
        r = self._role
        radio_set = self.query_one(f"#{r}-provider", RadioSet)
        pressed = radio_set.pressed_button
        btn_id = pressed.id if pressed else f"{r}-azure"
        key = (btn_id or f"{r}-azure").removeprefix(f"{r}-")
        provider_type = _PROVIDER_LABELS.get(key, "azure_openai")  # type: ignore[arg-type]

        if provider_type == "azure_openai":
            return ProviderConfig(
                provider_type="azure_openai",
                azure_endpoint=self._inp(f"{r}-endpoint"),
                azure_deployment=self._inp(f"{r}-deployment"),
                azure_api_version=self._inp(f"{r}-api-version"),
            )
        elif provider_type == "openrouter":
            return ProviderConfig(
                provider_type="openrouter", openrouter_model=self._sel(f"{r}-or-model")
            )
        elif provider_type == "vertex_ai":
            return ProviderConfig(
                provider_type="vertex_ai",
                vertex_model=self._sel(f"{r}-vertex-model"),
                vertex_project=self._inp(f"{r}-vertex-project"),
                vertex_region=self._inp(f"{r}-vertex-region"),
            )
        elif provider_type == "bedrock":
            return ProviderConfig(
                provider_type="bedrock",
                bedrock_model=self._sel(f"{r}-bedrock-model"),
                bedrock_region=self._inp(f"{r}-bedrock-region"),
                bedrock_profile=self._inp(f"{r}-bedrock-profile"),
            )
        elif provider_type == "cerebras":
            return ProviderConfig(
                provider_type="cerebras", cerebras_model=self._sel(f"{r}-cerebras-model")
            )
        elif provider_type == "mistral":
            return ProviderConfig(
                provider_type="mistral", mistral_model=self._sel(f"{r}-mistral-model")
            )
        else:
            return ProviderConfig(provider_type="azure_openai")


# ---------------------------------------------------------------------------
# MainScreen
# ---------------------------------------------------------------------------


class MainScreen(Screen):
    CSS = """
    MainScreen {
        align: center middle;
    }
    #menu {
        width: 52;
        height: auto;
        border: double $primary;
        padding: 2 4;
    }
    #menu-title {
        text-align: center;
        color: $accent;
        text-style: bold;
        padding-bottom: 0;
    }
    #menu-sub {
        text-align: center;
        color: $text-muted;
        padding-bottom: 2;
    }
    #menu Button {
        width: 100%;
        margin-bottom: 1;
    }
    """

    def compose(self) -> ComposeResult:
        with Container(id="menu"):
            yield Static("probegpt", id="menu-title")
            yield Static("Automated red-team discovery", id="menu-sub")
            yield Button("Run Discovery", variant="primary", id="btn-discover")
            yield Button("Config", id="btn-config")
            yield Button("Exit", variant="error", id="btn-exit")

    def on_button_pressed(self, event: Button.Pressed) -> None:
        match event.button.id:
            case "btn-discover":
                self.app.push_screen(DiscoveryScreen())
            case "btn-config":
                self.app.push_screen(ConfigScreen())
            case "btn-exit":
                self.app.exit()


# ---------------------------------------------------------------------------
# StrategyForm widget
# ---------------------------------------------------------------------------


class StrategyForm(Widget):
    DEFAULT_CSS = """
    StrategyForm {
        height: auto;
        padding: 1 2;
    }
    StrategyForm .strategy-row {
        height: auto;
        margin-bottom: 1;
    }
    StrategyForm Label {
        color: $text-muted;
        margin-top: 1;
    }
    StrategyForm .weight-input {
        width: 8;
        margin-left: 2;
    }
    StrategyForm .desc {
        color: $text-muted;
        text-style: italic;
        padding-left: 4;
    }
    """

    def __init__(self, strategy_config: StrategyConfig) -> None:
        super().__init__()
        self._config = strategy_config

    def compose(self) -> ComposeResult:
        enabled_set = set(self._config.enabled)
        for name in ALL_STRATEGY_NAMES:
            desc = _STRATEGY_DESCRIPTIONS.get(name, "")
            weight = self._config.weights.get(name, 1.0)
            with Horizontal(classes="strategy-row"):
                yield Checkbox(
                    f" {name}",
                    value=name in enabled_set,
                    id=f"strat-{name}",
                )
                yield Input(
                    value=str(weight),
                    placeholder="weight",
                    id=f"strat-weight-{name}",
                    classes="weight-input",
                )
            yield Static(f"[dim]{desc}[/dim]", markup=True, classes="desc")

    def get_config(self) -> StrategyConfig:
        enabled: list[str] = []
        weights: dict[str, float] = {}
        for name in ALL_STRATEGY_NAMES:
            cb = self.query_one(f"#strat-{name}", Checkbox)
            if cb.value:
                enabled.append(name)
            raw = self.query_one(f"#strat-weight-{name}", Input).value.strip()
            try:
                weights[name] = max(0.0, float(raw))
            except ValueError:
                weights[name] = 1.0
        return StrategyConfig(enabled=enabled, weights=weights)


# ---------------------------------------------------------------------------
# ConfigScreen
# ---------------------------------------------------------------------------


class ConfigScreen(Screen):
    CSS = """
    ConfigScreen TabbedContent {
        height: 1fr;
    }
    #config-actions {
        height: auto;
        padding: 1 2;
        align-horizontal: right;
    }
    #config-actions Button {
        margin-left: 1;
        width: 14;
    }
    """

    def compose(self) -> ComposeResult:
        yield Header()
        config = Config.from_file() if CONFIG_PATH.exists() else Config()
        with TabbedContent():
            with TabPane("Target", id="tab-target"):
                yield Static(
                    "[dim]The model you are testing. Receives every generated probe.[/dim]",
                    markup=True,
                )
                with VerticalScroll():
                    yield ProviderForm("target", config.target)
            with TabPane("Generator", id="tab-generator"):
                yield Static(
                    "[dim]Writes new probe variants by mutating / combining seeds.[/dim]",
                    markup=True,
                )
                with VerticalScroll():
                    yield ProviderForm("generator", config.generator)
            with TabPane("Judge", id="tab-judge"):
                yield Static(
                    "[dim]Evaluates target responses to decide if a probe succeeded.[/dim]",
                    markup=True,
                )
                with VerticalScroll():
                    yield ProviderForm("judge", config.judge)
            with TabPane("Strategies", id="tab-strategies"):
                yield Static(
                    "[dim]Enable / disable probe-generation strategies and adjust their relative weights.[/dim]",
                    markup=True,
                )
                with VerticalScroll():
                    yield StrategyForm(config.strategies)
        with Horizontal(id="config-actions"):
            yield Button("Back", id="btn-back")
            yield Button("Save", variant="primary", id="btn-save")
        yield Footer()

    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "btn-back":
            self.app.pop_screen()
        elif event.button.id == "btn-save":
            forms = list(self.query(ProviderForm))
            strategy_form = self.query_one(StrategyForm)
            config = Config(
                target=forms[0].get_config(),
                generator=forms[1].get_config(),
                judge=forms[2].get_config(),
                strategies=strategy_form.get_config(),
            )
            config.to_file()
            self.notify("Config saved.")
            self.app.pop_screen()


# ---------------------------------------------------------------------------
# DiscoveryScreen
# ---------------------------------------------------------------------------


class DiscoveryScreen(Screen):
    CSS = """
    DiscoveryScreen > Horizontal {
        height: 1fr;
    }
    #discovery-form {
        width: 38;
        min-width: 38;
        border-right: solid $primary;
        padding: 0 1;
    }
    #discovery-log {
        width: 1fr;
        padding: 0 1;
    }
    #discovery-log RichLog {
        height: 1fr;
    }
    .section-label {
        color: $accent;
        text-style: bold;
        margin-top: 1;
    }
    .hidden {
        display: none;
    }
    #form-actions {
        height: auto;
        margin-top: 1;
    }
    #form-actions Button {
        margin-right: 1;
    }
    #progress-status {
        color: $text-muted;
        margin-top: 1;
    }
    #progress-bar {
        margin-bottom: 1;
    }
    """

    def compose(self) -> ComposeResult:
        yield Header()
        with Horizontal():
            with VerticalScroll(id="discovery-form"):
                yield Label("Objective", classes="section-label")
                yield Input(
                    placeholder="What should the model do?",
                    id="objective",
                )

                yield Label("Parameters", classes="section-label")
                yield Label("Iterations")
                yield Input(value="3", id="iterations")
                yield Label("Candidates per iteration")
                yield Input(value="5", id="candidates")

                yield Label("Output", classes="section-label")
                yield Checkbox("Save results to JSON", id="save-json")
                yield Input(
                    placeholder="discovery_results.json",
                    id="json-path",
                    classes="hidden",
                )

                with Horizontal(id="form-actions"):
                    yield Button("Back", id="btn-back")
                    yield Button("Run", variant="primary", id="btn-run")

            with Container(id="discovery-log"):
                yield RichLog(highlight=True, markup=True, id="log")
                yield Label("Idle", id="progress-status")
                yield ProgressBar(total=100, id="progress-bar", show_eta=False)
        yield Footer()

    def on_checkbox_changed(self, event: Checkbox.Changed) -> None:
        if event.checkbox.id == "save-json":
            self.query_one("#json-path").set_class(not event.value, "hidden")

    def on_button_pressed(self, event: Button.Pressed) -> None:
        match event.button.id:
            case "btn-back":
                self.app.pop_screen()
            case "btn-run":
                self._start_discovery()

    def _start_discovery(self) -> None:
        log = self.query_one("#log", RichLog)
        log.clear()

        objective = self.query_one("#objective", Input).value.strip()
        if not objective:
            log.write("[yellow]Please enter an objective.[/yellow]")
            return

        try:
            config = Config.from_file()
            target_model = create_model(config.target)
            generator_model = create_model(config.generator)
            judge_model = create_model(config.judge)
        except (FileNotFoundError, ValueError) as e:
            log.write(f"[red]{e}[/red]")
            return

        try:
            iterations = int(self.query_one("#iterations", Input).value or "3")
            candidates = int(self.query_one("#candidates", Input).value or "5")
        except ValueError:
            log.write("[red]Invalid iterations or candidates value.[/red]")
            return

        total = iterations * candidates

        # Reset and configure progress bar
        pb = self.query_one("#progress-bar", ProgressBar)
        pb.update(total=total, progress=0)
        status_label = self.query_one("#progress-status", Label)
        status_label.update("Starting…")

        outputs: list[AbstractOutput] = [TextualOutput(self.app, log, pb, status_label, total)]
        if self.query_one("#save-json", Checkbox).value:
            json_path = (
                self.query_one("#json-path", Input).value.strip() or "discovery_results.json"
            )
            outputs.append(JsonOutput(Path(json_path)))

        log.write(
            f"[cyan]{iterations} iter · {candidates} candidates/iter · ~{total} probes[/cyan]\n"
            f"[dim]Objective: {objective}[/dim]\n"
        )

        loop = DiscoveryLoop(
            generator_model=generator_model,
            target_model=target_model,
            judge_model=judge_model,
            objective=objective,
            outputs=outputs,
            iterations=iterations,
            candidates_per_iter=candidates,
            strategy_config=config.strategies,
        )

        run_btn = self.query_one("#btn-run", Button)
        run_btn.disabled = True
        self._run_loop(loop, log, run_btn)

    @work(thread=True)
    def _run_loop(
        self,
        loop: DiscoveryLoop,
        log: RichLog,
        run_btn: Button,
    ) -> None:
        try:
            loop.run()
        except Exception as e:
            import traceback

            tb = traceback.format_exc()
            self.app.call_from_thread(
                log.write,
                f"[red bold]Error: {type(e).__name__}: {e}[/red bold]\n[red]{tb}[/red]",
            )
        finally:
            self.app.call_from_thread(setattr, run_btn, "disabled", False)


# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------


class ProbegptApp(App):
    TITLE = "probegpt"
    BINDINGS = [Binding("q", "quit", "Quit")]

    def on_mount(self) -> None:
        self.push_screen(MainScreen())
