"""Compare screen for the ALiCE TUI — side-by-side model comparison."""

from __future__ import annotations

import difflib
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from textual.binding import Binding
from textual.containers import Horizontal, Vertical
from textual.widgets import Button, Checkbox, Footer, Label, TextArea

from alice_cli.constants import DEFAULT_MAX_TOKENS
from alice_cli.models import (
    MODEL_ALIASES,
    resolve_model_id,
)
from alice_cli.tui.screens.base import AliceScreen
from alice_cli.tui.widgets.breadcrumb import BreadcrumbBar
from alice_cli.tui.widgets.header_bar import AliceHeaderBar
from alice_cli.tui.widgets.output import OutputPanel

if TYPE_CHECKING:
    from textual.app import ComposeResult


@dataclass
class _ModelResult:
    """Result from a single model invocation."""

    alias: str
    resolved_id: str
    text: str
    input_tokens: int
    output_tokens: int
    error: str | None = None


def _get_comparable_aliases() -> list[str]:
    """Return model aliases excluding convenience shorthands (those starting with →)."""
    return [m.alias for m in MODEL_ALIASES if not m.description.startswith("→")]


def _compute_diff(text_a: str, text_b: str, label_a: str, label_b: str) -> str:
    """Compute a unified diff between two response texts."""
    a_lines = text_a.splitlines(keepends=True)
    b_lines = text_b.splitlines(keepends=True)
    diff = difflib.unified_diff(a_lines, b_lines, fromfile=label_a, tofile=label_b)
    diff_text = "".join(diff)
    return diff_text if diff_text else "(no differences)"


def _invoke_single_model(
    client: Any,
    resolved_id: str,
    alias: str,
    prompt_text: str,
) -> _ModelResult:
    """Invoke a single model via the Converse API. Returns a _ModelResult."""
    try:
        response = client.converse(
            modelId=resolved_id,
            messages=[{"role": "user", "content": [{"text": prompt_text}]}],
            inferenceConfig={"maxTokens": DEFAULT_MAX_TOKENS},
        )
        output_message = response.get("output", {}).get("message", {})
        content_blocks = output_message.get("content", [])
        text = "\n".join(b["text"] for b in content_blocks if "text" in b)
        usage = response.get("usage", {})
        return _ModelResult(
            alias=alias,
            resolved_id=resolved_id,
            text=text,
            input_tokens=usage.get("inputTokens", 0),
            output_tokens=usage.get("outputTokens", 0),
        )
    except Exception as exc:
        return _ModelResult(
            alias=alias,
            resolved_id=resolved_id,
            text="",
            input_tokens=0,
            output_tokens=0,
            error=str(exc),
        )


class CompareScreen(AliceScreen):
    """Side-by-side model comparison."""

    BINDINGS = [
        Binding("escape", "pop_screen", "Back"),
    ]

    DEFAULT_CSS = """
    CompareScreen {
        layout: vertical;
    }

    CompareScreen #compare-prompt {
        height: 4;
        min-height: 4;
    }

    CompareScreen #model-checklist {
        height: auto;
        max-height: 10;
        padding: 0 1;
    }

    CompareScreen #compare-actions {
        height: 3;
        padding: 0 1;
        align: left middle;
    }

    CompareScreen #compare-btn {
        min-width: 12;
    }

    CompareScreen #diff-btn {
        min-width: 12;
    }

    CompareScreen #compare-panels {
        height: 1fr;
    }

    CompareScreen #compare-panels OutputPanel {
        width: 1fr;
        border: tall #FFB000;
    }

    CompareScreen #diff-panel {
        height: auto;
        max-height: 12;
        border: tall #00AACC;
        display: none;
    }

    CompareScreen #summary-label {
        width: 1fr;
        color: #805800;
        padding: 0 2;
    }
    """

    def __init__(self) -> None:
        super().__init__()
        self._results: list[_ModelResult] = []
        self._running: bool = False
        self._diff_visible: bool = False
        self._total_input_tokens: int = 0
        self._total_output_tokens: int = 0

    def compose(self) -> ComposeResult:
        yield AliceHeaderBar()
        yield BreadcrumbBar(id="breadcrumb")
        yield TextArea(id="compare-prompt")

        # Model checklist — one checkbox per non-shorthand alias
        with Vertical(id="model-checklist"):
            for alias in _get_comparable_aliases():
                yield Checkbox(alias, id=f"chk-{alias}")

        # Action buttons
        with Horizontal(id="compare-actions"):
            yield Button("Compare", id="compare-btn", variant="primary")
            yield Button("Diff", id="diff-btn", disabled=True)
            yield Label("", id="summary-label")

        # Dynamic result panels container
        with Horizontal(id="compare-panels"):
            pass  # Panels are added dynamically

        # Diff output panel (hidden by default)
        yield OutputPanel(id="diff-panel")
        yield Footer()

    def key_j(self) -> None:
        """Move focus to the next focusable child (vim-style down).

        No-ops when a text input (Input / TextArea) is focused so
        ``j`` types normally.
        """
        if self._text_input_has_focus():
            return
        self.focus_next()

    def key_k(self) -> None:
        """Move focus to the previous focusable child (vim-style up).

        No-ops when a text input (Input / TextArea) is focused so
        ``k`` types normally.
        """
        if self._text_input_has_focus():
            return
        self.focus_previous()

    def _get_selected_models(self) -> list[str]:
        """Return the list of checked model aliases."""
        selected: list[str] = []
        for alias in _get_comparable_aliases():
            try:
                cb = self.query_one(f"#chk-{alias}", Checkbox)
                if cb.value:
                    selected.append(alias)
            except Exception:
                continue
        return selected

    def on_button_pressed(self, event: Button.Pressed) -> None:
        """Handle Compare and Diff button presses."""
        if event.button.id == "compare-btn":
            self._run_comparison()
        elif event.button.id == "diff-btn":
            self._toggle_diff()

    def _run_comparison(self) -> None:
        """Validate inputs and kick off concurrent model invocations."""
        if self._running:
            return

        prompt_text = self.query_one("#compare-prompt", TextArea).text.strip()
        if not prompt_text:
            return

        selected = self._get_selected_models()
        if len(selected) < 2:
            self.query_one("#summary-label", Label).update("[#FF6600]Select at least 2 models[/]")
            return

        # Clear previous results
        self._results.clear()
        self._total_input_tokens = 0
        self._total_output_tokens = 0
        self._diff_visible = False

        # Remove old panels
        panels_container = self.query_one("#compare-panels", Horizontal)
        for child in list(panels_container.children):
            child.remove()

        # Create spinner panels for each selected model
        for alias in selected:
            panel = OutputPanel(id=f"panel-{alias}")
            panels_container.mount(panel)
            panel.write(f"[dim]{alias}: waiting…[/dim]")

        # Hide diff panel
        diff_panel = self.query_one("#diff-panel", OutputPanel)
        diff_panel.styles.display = "none"
        diff_panel.clear()

        self.query_one("#diff-btn", Button).disabled = True
        self.query_one("#summary-label", Label).update("")

        # Save prompt to history
        self._save_prompt("compare", prompt_text)

        self._running = True
        self.run_worker(
            lambda: self._invoke_models(prompt_text, selected),
            thread=True,
        )

    def _invoke_models(self, prompt_text: str, selected: list[str]) -> None:
        """Worker thread: invoke all selected models concurrently."""
        try:
            session = self._create_session()
            client = session.client("bedrock-runtime")

            # Resolve aliases to model IDs
            resolved: list[tuple[str, str]] = []
            for alias in selected:
                model_id = resolve_model_id(alias)
                resolved.append((alias, model_id))

            results: list[_ModelResult] = []

            with ThreadPoolExecutor(max_workers=len(resolved)) as executor:
                futures = {
                    executor.submit(
                        _invoke_single_model, client, model_id, alias, prompt_text
                    ): alias
                    for alias, model_id in resolved
                }
                for future in as_completed(futures):
                    result = future.result()
                    results.append(result)
                    # Update the panel as each result arrives
                    self.app.call_from_thread(self._update_panel, result)

            # Preserve original selection order
            order = {alias: i for i, alias in enumerate(selected)}
            results.sort(key=lambda r: order.get(r.alias, 0))
            self._results = results

            # Tally tokens
            total_in = sum(r.input_tokens for r in results)
            total_out = sum(r.output_tokens for r in results)
            self._total_input_tokens = total_in
            self._total_output_tokens = total_out

            self.app.call_from_thread(self._on_comparison_complete)

        except Exception as exc:
            self.app.call_from_thread(
                self.query_one("#summary-label", Label).update,
                f"[#FF6600]Error: {exc}[/]",
            )
        finally:
            self._running = False

    def _update_panel(self, result: _ModelResult) -> None:
        """Update a single model's output panel with its result."""
        try:
            panel = self.query_one(f"#panel-{result.alias}", OutputPanel)
            panel.clear()
            if result.error:
                panel.write(f"[#FF6600]Error: {result.error}[/]")
            else:
                panel.write(result.text or "[dim](empty response)[/dim]")
        except Exception:
            pass

    def _on_comparison_complete(self) -> None:
        """Show summary and enable Diff button after all models respond."""
        total = self._total_input_tokens + self._total_output_tokens
        n = len(self._results)
        errors = sum(1 for r in self._results if r.error)
        summary = (
            f"✓ {n} models compared  •  "
            f"↑{self._total_input_tokens:,}  ↓{self._total_output_tokens:,}  "
            f"Σ{total:,}"
        )
        if errors:
            summary += f"  •  {errors} error(s)"
        self.query_one("#summary-label", Label).update(summary)

        # Enable diff if at least 2 successful results
        successful = [r for r in self._results if r.error is None]
        if len(successful) >= 2:
            self.query_one("#diff-btn", Button).disabled = False

        # Update global TokenTracker
        self._update_token_tracker(self._total_input_tokens, self._total_output_tokens)

    # _update_token_tracker is inherited from AliceScreen

    def _toggle_diff(self) -> None:
        """Toggle the diff panel visibility and compute diff if showing."""
        diff_panel = self.query_one("#diff-panel", OutputPanel)
        self._diff_visible = not self._diff_visible

        if self._diff_visible:
            successful = [r for r in self._results if r.error is None]
            if len(successful) >= 2:
                diff_text = _compute_diff(
                    successful[0].text,
                    successful[1].text,
                    successful[0].alias,
                    successful[1].alias,
                )
                diff_panel.clear()
                diff_panel.write(diff_text)
                diff_panel.styles.display = "block"
        else:
            diff_panel.styles.display = "none"
