"""Dialog screen for the ALiCE TUI — watch two models converse in split panes."""

from __future__ import annotations

import uuid
from typing import TYPE_CHECKING, Any

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

from alice_cli.constants import DEFAULT_MAX_TOKENS
from alice_cli.models import (
    DEFAULT_MODEL_ALIAS,
    TokenUsage,
    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


def _default_dialog_system(my_alias: str, other_alias: str, user_system: str | None) -> str:
    """Build a system prompt that frames the model as a dialog participant."""
    base = (
        f"You are {my_alias} in a dialog with {other_alias}. "
        "A human has posed a topic for discussion. "
        "Build on what has been said so far — add new ideas, "
        "challenge assumptions, or deepen the analysis. "
        "Do not simply repeat or rephrase the previous message. "
        "Keep your response focused and concise."
    )
    if user_system:
        return f"{user_system}\n\n{base}"
    return base


def _build_messages_for(
    transcript: list[tuple[str, str]],
    my_label: str,
    opening_prompt: str,
) -> list[dict[str, Any]]:
    """Build a Converse-API message list from the shared transcript.

    Each transcript entry is ``(speaker_label, text)``.  Entries where
    ``speaker_label == my_label`` become ``"assistant"`` messages; all
    others become ``"user"`` messages.  The opening prompt is always
    prepended as the first ``"user"`` message so the model has the
    original context.

    Adjacent messages with the same role are merged (the Converse API
    requires strict user/assistant alternation).
    """
    raw: list[dict[str, Any]] = [{"role": "user", "content": [{"text": opening_prompt}]}]
    for speaker, text in transcript:
        role = "assistant" if speaker == my_label else "user"
        raw.append({"role": role, "content": [{"text": text}]})

    # Merge consecutive same-role messages to satisfy the alternation rule.
    merged: list[dict[str, Any]] = []
    for msg in raw:
        if merged and merged[-1]["role"] == msg["role"]:
            merged[-1]["content"].extend(msg["content"])
        else:
            merged.append(msg)

    # The list must start with "user" and end with "user" (the model replies).
    if merged and merged[-1]["role"] == "assistant":
        merged.append({"role": "user", "content": [{"text": "Please continue."}]})

    return merged


class DialogScreen(AliceScreen):
    """Watch two models converse in split panes."""

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

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

    DialogScreen #dialog-config-row {
        height: 3;
        padding: 0 1;
        align: left middle;
    }

    DialogScreen #prompt-input {
        width: 1fr;
    }

    DialogScreen #model-a-input {
        width: 16;
    }

    DialogScreen #model-b-input {
        width: 16;
    }

    DialogScreen #rounds-input {
        width: 8;
    }

    DialogScreen #start-btn {
        min-width: 10;
    }

    DialogScreen #dialog-panels {
        height: 1fr;
    }

    DialogScreen #dialog-panel-a {
        width: 1fr;
        border: tall #00AACC;
    }

    DialogScreen #dialog-panel-b {
        width: 1fr;
        border: tall #CC44AA;
    }

    DialogScreen #round-counter {
        width: auto;
        min-width: 12;
        text-align: center;
        color: #FFD866;
        padding: 1 1;
        content-align: center middle;
    }

    DialogScreen #save-btn {
        min-width: 10;
        dock: bottom;
    }

    DialogScreen #dialog-footer-row {
        height: 3;
        padding: 0 1;
        align: center middle;
    }

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

    def __init__(self) -> None:
        super().__init__()
        self._transcript: list[tuple[str, str]] = []
        self._session_id: str = str(uuid.uuid4())
        self._total_input_tokens: int = 0
        self._total_output_tokens: int = 0
        self._running: bool = False
        self._completed: bool = False

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

        # Config row: prompt, model-a, model-b, rounds, Start
        with Horizontal(id="dialog-config-row"):
            yield Input(
                placeholder="Opening prompt for the dialog…",
                id="prompt-input",
            )
            yield Input(
                value=DEFAULT_MODEL_ALIAS,
                placeholder="Model A",
                id="model-a-input",
            )
            yield Input(
                value="haiku",
                placeholder="Model B",
                id="model-b-input",
            )
            yield Input(
                value="3",
                placeholder="Rounds",
                id="rounds-input",
            )
            yield Button("Start", id="start-btn", variant="primary")

        # Split panes with round counter between them
        with Horizontal(id="dialog-panels"):
            yield OutputPanel(id="dialog-panel-a")
            with Vertical(id="round-counter-container"):
                yield Label("Round 0/0", id="round-counter")
            yield OutputPanel(id="dialog-panel-b")

        # Footer row with summary and save button
        with Horizontal(id="dialog-footer-row"):
            yield Label("", id="summary-label")
            yield Button("Save", id="save-btn", disabled=True)

        yield Footer()

    def on_mount(self) -> None:
        """Set model input defaults from persisted TUI state."""
        state = getattr(self.app, "tui_state", None)
        if state and state.last_model:
            self.query_one("#model-a-input", Input).value = state.last_model

    def on_button_pressed(self, event: Button.Pressed) -> None:
        """Handle Start and Save button presses."""
        if event.button.id == "start-btn":
            self._start_dialog()
        elif event.button.id == "save-btn":
            self._save_transcript()

    def _start_dialog(self) -> None:
        """Validate inputs and kick off the dialog loop in a worker thread."""
        if self._running:
            return

        prompt_text = self.query_one("#prompt-input", Input).value.strip()
        if not prompt_text:
            return

        model_a_alias = self.query_one("#model-a-input", Input).value.strip()
        model_b_alias = self.query_one("#model-b-input", Input).value.strip()
        if not model_a_alias or not model_b_alias:
            return

        rounds_text = self.query_one("#rounds-input", Input).value.strip()
        try:
            rounds = int(rounds_text)
        except ValueError:
            rounds = 3
        if rounds < 1:
            rounds = 1

        model_a_id = resolve_model_id(model_a_alias)
        model_b_id = resolve_model_id(model_b_alias)

        # Clear previous state
        self._transcript.clear()
        self._total_input_tokens = 0
        self._total_output_tokens = 0
        self._completed = False
        self.query_one("#dialog-panel-a", OutputPanel).clear()
        self.query_one("#dialog-panel-b", OutputPanel).clear()
        self.query_one("#summary-label", Label).update("")
        self.query_one("#save-btn", Button).disabled = True

        # Persist model alias
        state = getattr(self.app, "tui_state", None)
        if state:
            state.last_model = model_a_alias

        self._running = True
        self.run_worker(
            lambda: self._dialog_loop(
                prompt_text,
                model_a_alias,
                model_b_alias,
                model_a_id,
                model_b_id,
                rounds,
            ),
            thread=True,
        )

    def _dialog_loop(
        self,
        prompt_text: str,
        model_a_alias: str,
        model_b_alias: str,
        model_a_id: str,
        model_b_id: str,
        rounds: int,
    ) -> None:
        """Worker thread: run the dialog loop, updating panels via call_from_thread."""
        panel_a = self.query_one("#dialog-panel-a", OutputPanel)
        panel_b = self.query_one("#dialog-panel-b", OutputPanel)

        sys_a = [{"text": _default_dialog_system(model_a_alias, model_b_alias, None)}]
        sys_b = [{"text": _default_dialog_system(model_b_alias, model_a_alias, None)}]

        try:
            session = self._create_session()
            client = session.client("bedrock-runtime")

            for round_num in range(1, rounds + 1):
                self.app.call_from_thread(self._update_round_counter, round_num, rounds)

                # Model A's turn
                messages_a = _build_messages_for(self._transcript, model_a_alias, prompt_text)
                self.app.call_from_thread(
                    panel_a.write,
                    f"[dim]── Round {round_num}/{rounds} ──[/dim]",
                )

                resp_a = client.converse(
                    modelId=model_a_id,
                    messages=messages_a,
                    system=sys_a,
                    inferenceConfig={"maxTokens": DEFAULT_MAX_TOKENS},
                )

                text_a = "\n".join(
                    b["text"]
                    for b in resp_a.get("output", {}).get("message", {}).get("content", [])
                    if "text" in b
                )
                self._transcript.append((model_a_alias, text_a))
                self.app.call_from_thread(panel_a.write, text_a)

                # Track tokens from model A
                usage_a = resp_a.get("usage", {})
                in_a = usage_a.get("inputTokens", 0)
                out_a = usage_a.get("outputTokens", 0)
                self._total_input_tokens += in_a
                self._total_output_tokens += out_a
                self.app.call_from_thread(self._update_token_tracker, in_a, out_a)

                # Model B's turn
                messages_b = _build_messages_for(self._transcript, model_b_alias, prompt_text)
                self.app.call_from_thread(
                    panel_b.write,
                    f"[dim]── Round {round_num}/{rounds} ──[/dim]",
                )

                resp_b = client.converse(
                    modelId=model_b_id,
                    messages=messages_b,
                    system=sys_b,
                    inferenceConfig={"maxTokens": DEFAULT_MAX_TOKENS},
                )

                text_b = "\n".join(
                    b["text"]
                    for b in resp_b.get("output", {}).get("message", {}).get("content", [])
                    if "text" in b
                )
                self._transcript.append((model_b_alias, text_b))
                self.app.call_from_thread(panel_b.write, text_b)

                # Track tokens from model B
                usage_b = resp_b.get("usage", {})
                in_b = usage_b.get("inputTokens", 0)
                out_b = usage_b.get("outputTokens", 0)
                self._total_input_tokens += in_b
                self._total_output_tokens += out_b
                self.app.call_from_thread(self._update_token_tracker, in_b, out_b)

            # Dialog complete
            self._completed = True
            self.app.call_from_thread(self._on_dialog_complete, rounds)

        except Exception as exc:
            self.app.call_from_thread(panel_a.write, f"[#FF6600]Error: {exc}[/]")

        finally:
            self._running = False

    def _update_round_counter(self, current: int, total: int) -> None:
        """Update the round counter label."""
        self.query_one("#round-counter", Label).update(f"Round {current}/{total}")

    def _on_dialog_complete(self, rounds: int) -> None:
        """Show summary and enable Save button after dialog finishes."""
        total_tokens = self._total_input_tokens + self._total_output_tokens
        summary = (
            f"✓ Dialog complete — {rounds} rounds  •  "
            f"↑{self._total_input_tokens:,}  ↓{self._total_output_tokens:,}  "
            f"Σ{total_tokens:,}"
        )
        self.query_one("#summary-label", Label).update(summary)
        self.query_one("#save-btn", Button).disabled = False

    # _update_token_tracker is inherited from AliceScreen

    def _save_transcript(self) -> None:
        """Persist the dialog transcript to Cloud Locker."""
        if not self._transcript:
            return

        try:
            from alice_cli.auth import detect_username
            from alice_cli.constants import get_memory_bucket
            from alice_cli.locker import CloudLocker
            from alice_cli.session_record import build_session_record

            prompt_text = self.query_one("#prompt-input", Input).value.strip()
            model_a_alias = self.query_one("#model-a-input", Input).value.strip()
            model_b_alias = self.query_one("#model-b-input", Input).value.strip()

            record = build_session_record(
                record_type="dialog",
                model=f"{model_a_alias} vs {model_b_alias}",
                prompt=prompt_text,
                response=None,
                tokens=TokenUsage(
                    input=self._total_input_tokens,
                    output=self._total_output_tokens,
                    total=self._total_input_tokens + self._total_output_tokens,
                ),
                metadata={
                    "model_a": model_a_alias,
                    "model_b": model_b_alias,
                    "rounds": len(self._transcript) // 2,
                    "transcript": [
                        {"speaker": speaker, "text": text} for speaker, text in self._transcript
                    ],
                    "session_id": self._session_id,
                },
            )

            result = detect_username(self.config.profile, self.config.region)
            locker = CloudLocker(
                config=self.config,
                username=result.username,
                bucket=get_memory_bucket(self.config.namespace, self.config.environment),
            )
            locker.persist(record, "dialog")
            self.query_one("#summary-label", Label).update("✓ Transcript saved to Cloud Locker")
        except Exception:  # noqa: BLE001
            self.query_one("#summary-label", Label).update(
                "[#FF6600]Save failed — check credentials[/]"
            )
