"""outofthebox — 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 core.config import DEFAULT_CONFIG_PATH, Config, ProviderConfig
from core.discovery import DiscoveryLoop
from core.models import DiscoveryResult
from output.base import AbstractOutput
from output.json_exporter import JsonOutput
from providers.factory import create_model

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


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 {}


# ---------------------------------------------------------------------------
# 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
# ---------------------------------------------------------------------------


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

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

    def compose(self) -> ComposeResult:
        cfg = self._config
        is_azure = cfg.provider_type == "azure_openai"

        with RadioSet(id=f"{self._role}-provider"):
            yield RadioButton("Azure OpenAI", value=is_azure, id=f"{self._role}-azure")
            yield RadioButton("OpenRouter", value=not is_azure, id=f"{self._role}-openrouter")

        # Azure fields
        with Container(id=f"{self._role}-azure-fields", classes="" if is_azure else "hidden"):
            yield Label("Endpoint")
            yield Input(
                value=cfg.azure_endpoint,
                placeholder="https://<resource>.openai.azure.com",
                id=f"{self._role}-endpoint",
            )
            yield Label("Deployment")
            yield Input(
                value=cfg.azure_deployment,
                placeholder="gpt-4",
                id=f"{self._role}-deployment",
            )
            yield Label("API version")
            yield Input(
                value=cfg.azure_api_version or "2024-02-15-preview",
                id=f"{self._role}-api-version",
            )

        # OpenRouter fields
        with Container(
            id=f"{self._role}-openrouter-fields", classes="" if not is_azure else "hidden"
        ):
            yield Label("Model", id=f"{self._role}-model-label")
            current = cfg.openrouter_model
            initial: list[tuple[str, str]] = [(current, current)] if current else []
            yield Select(
                initial,
                value=current if current else Select.NULL,
                allow_blank=True,
                id=f"{self._role}-model",
            )

    def on_mount(self) -> None:
        if self._config.provider_type == "openrouter":
            self._load_models()

    def on_radio_set_changed(self, event: RadioSet.Changed) -> None:
        is_azure = event.pressed is not None and "azure" in (event.pressed.id or "")
        self.query_one(f"#{self._role}-azure-fields").set_class(not is_azure, "hidden")
        self.query_one(f"#{self._role}-openrouter-fields").set_class(is_azure, "hidden")
        if not is_azure:
            self._load_models()

    @work(thread=True)
    def _load_models(self) -> None:
        by_org = _fetch_openrouter_models()
        options = [(mid, mid) for ids in by_org.values() for mid in ids]
        self.app.call_from_thread(self._set_model_options, options)

    def _set_model_options(self, options: list[tuple[str, str]]) -> None:
        select = self.query_one(f"#{self._role}-model", Select)
        select.set_options(options)
        current = self._config.openrouter_model
        if current:
            select.value = current
        self.query_one(f"#{self._role}-model-label", Label).update(
            f"Model  ({len(options)} available, type to filter)"
        )

    def get_config(self) -> ProviderConfig:
        radio_set = self.query_one(f"#{self._role}-provider", RadioSet)
        pressed = radio_set.pressed_button
        is_azure = pressed is not None and "azure" in (pressed.id or "")

        if is_azure:
            return ProviderConfig(
                provider_type="azure_openai",
                azure_endpoint=self.query_one(f"#{self._role}-endpoint", Input).value,
                azure_deployment=self.query_one(f"#{self._role}-deployment", Input).value,
                azure_api_version=self.query_one(f"#{self._role}-api-version", Input).value,
            )
        else:
            sel = self.query_one(f"#{self._role}-model", Select)
            model = "" if sel.is_blank() else str(sel.value)
            return ProviderConfig(provider_type="openrouter", openrouter_model=model)


# ---------------------------------------------------------------------------
# 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()


# ---------------------------------------------------------------------------
# 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 DEFAULT_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 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))
            config = Config(
                target=forms[0].get_config(),
                generator=forms[1].get_config(),
                judge=forms[2].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,
        )

        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:
            self.app.call_from_thread(log.write, f"[red]Error: {e}[/red]")
        finally:
            self.app.call_from_thread(setattr, run_btn, "disabled", False)


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


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

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