"""Chat screen for the ALiCE TUI — multi-turn streaming conversation."""

from __future__ import annotations

import uuid
from typing import TYPE_CHECKING

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

from alice_cli.constants import DEFAULT_MAX_TOKENS
from alice_cli.models import (
    DEFAULT_MODEL_ALIAS,
    MODEL_ALIASES,
    ChatMessage,
    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.chat_log import ChatLog
from alice_cli.tui.widgets.header_bar import AliceHeaderBar

if TYPE_CHECKING:
    from textual.app import ComposeResult


class ChatScreen(AliceScreen):
    """Multi-turn chat with streaming Bedrock responses."""

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

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

    ChatScreen #chat-model-bar {
        height: 3;
        padding: 0 1;
        align: left middle;
    }

    ChatScreen #model-input {
        width: 40;
    }

    ChatScreen #token-label {
        width: auto;
        min-width: 30;
        padding: 0 2;
        color: #805800;
    }

    ChatScreen #message-input {
        height: 4;
        min-height: 4;
        width: 1fr;
    }

    ChatScreen #system-prompt-input {
        height: 3;
    }

    ChatScreen #input-row {
        height: auto;
        padding: 0 0;
        align: left bottom;
    }

    ChatScreen #btn-send {
        min-width: 10;
        margin: 0 0 0 1;
    }
    """

    def __init__(self) -> None:
        super().__init__()
        self._messages: list[ChatMessage] = []
        self._session_id: str = str(uuid.uuid4())
        self._total_input_tokens: int = 0
        self._total_output_tokens: int = 0
        self._system_prompt: str | None = None
        self._streaming: bool = False
        self._prompt_history_index: int = -1

    def compose(self) -> ComposeResult:
        alias_hint = "  ".join(f"{m.alias}={m.description}" for m in MODEL_ALIASES[:6])
        yield AliceHeaderBar()
        yield BreadcrumbBar(id="breadcrumb")
        yield ChatLog(id="chat-log")
        with Horizontal(id="chat-model-bar"):
            yield Input(
                value=DEFAULT_MODEL_ALIAS,
                placeholder=f"Model alias: {alias_hint}",
                id="model-input",
            )
            yield Label("↑0  ↓0  Σ0", id="token-label")
        with Collapsible(title="System Prompt", collapsed=True, id="settings-panel"):
            yield TextArea(id="system-prompt-input")
        with Horizontal(id="input-row"):
            yield TextArea(id="message-input")
            yield Button("Send", id="btn-send", variant="primary")
        yield Footer()

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

    def on_text_area_changed(self, event: TextArea.Changed) -> None:
        """Capture system prompt changes."""
        if event.text_area.id == "system-prompt-input":
            text = event.text_area.text.strip()
            self._system_prompt = text if text else None

    def on_button_pressed(self, event: Button.Pressed) -> None:
        """Handle the Send button press."""
        if event.button.id == "btn-send":
            self._send_message()

    def key_enter(self) -> None:
        """Send message on Enter when the message input is focused."""
        focused = self.app.focused
        if focused is not None and getattr(focused, "id", None) == "message-input":
            self._send_message()

    def key_up(self) -> None:
        """Cycle backward through prompt history."""
        focused = self.app.focused
        if focused is None or getattr(focused, "id", None) != "message-input":
            return
        history = self._get_prompt_history("chat")
        if not history:
            return
        if self._prompt_history_index == -1:
            self._prompt_history_index = len(history) - 1
        elif self._prompt_history_index > 0:
            self._prompt_history_index -= 1
        ta = self.query_one("#message-input", TextArea)
        ta.clear()
        ta.insert(history[self._prompt_history_index])

    def key_down(self) -> None:
        """Cycle forward through prompt history."""
        focused = self.app.focused
        if focused is None or getattr(focused, "id", None) != "message-input":
            return
        history = self._get_prompt_history("chat")
        if not history or self._prompt_history_index == -1:
            return
        if self._prompt_history_index < len(history) - 1:
            self._prompt_history_index += 1
            ta = self.query_one("#message-input", TextArea)
            ta.clear()
            ta.insert(history[self._prompt_history_index])
        else:
            self._prompt_history_index = -1
            ta = self.query_one("#message-input", TextArea)
            ta.clear()

    def _send_message(self) -> None:
        """Send user message and stream the response."""
        if self._streaming:
            return

        message_area = self.query_one("#message-input", TextArea)
        prompt_text = message_area.text.strip()
        if not prompt_text:
            return

        model_input = self.query_one("#model-input", Input)
        model_id = resolve_model_id(model_input.value.strip())
        if not model_id:
            return

        # Clear input and add user message to chat log
        message_area.clear()
        self._prompt_history_index = -1
        chat_log = self.query_one("#chat-log", ChatLog)
        chat_log.add_user_message(prompt_text)

        # Track message in conversation history
        self._messages.append(ChatMessage(role="user", content=prompt_text))

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

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

        # Start streaming in worker thread
        self._streaming = True
        self.run_worker(
            lambda: self._stream_response(model_id),
            thread=True,
        )

    def _stream_response(self, model_id: str) -> None:
        """Worker thread: call converse_stream and push chunks to the UI."""
        chat_log = self.query_one("#chat-log", ChatLog)

        # Create an empty assistant bubble for streaming
        self.app.call_from_thread(chat_log.add_assistant_message, "")

        api_messages = [{"role": m.role, "content": [{"text": m.content}]} for m in self._messages]

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

            kwargs = {
                "modelId": model_id,
                "messages": api_messages,
                "inferenceConfig": {"maxTokens": DEFAULT_MAX_TOKENS},
            }
            if self._system_prompt:
                kwargs["system"] = [{"text": self._system_prompt}]

            response = client.converse_stream(**kwargs)

            assistant_text = ""
            for event in response.get("stream", []):
                if "contentBlockDelta" in event:
                    delta = event["contentBlockDelta"].get("delta", {})
                    if "text" in delta:
                        chunk = delta["text"]
                        assistant_text += chunk
                        self.app.call_from_thread(chat_log.append_to_last, chunk)

                # Extract token usage from metadata event
                if "metadata" in event:
                    usage = event["metadata"].get("usage", {})
                    input_tokens = usage.get("inputTokens", 0)
                    output_tokens = usage.get("outputTokens", 0)
                    self._total_input_tokens += input_tokens
                    self._total_output_tokens += output_tokens
                    self.app.call_from_thread(self._update_token_label)
                    self.app.call_from_thread(
                        self._update_token_tracker, input_tokens, output_tokens
                    )

            # Append assistant message to conversation history
            self._messages.append(ChatMessage(role="assistant", content=assistant_text))

        except Exception as exc:
            error_msg = f"Error: {exc}"
            self.app.call_from_thread(chat_log.append_to_last, error_msg)

        finally:
            self._streaming = False

    def _update_token_label(self) -> None:
        """Update the token count label."""
        total = self._total_input_tokens + self._total_output_tokens
        label = self.query_one("#token-label", Label)
        label.update(f"↑{self._total_input_tokens:,}  ↓{self._total_output_tokens:,}  Σ{total:,}")

    # _update_token_tracker is inherited from AliceScreen

    def _persist_session(self) -> None:
        """Save the conversation to Cloud Locker on screen exit."""
        if not self._messages:
            return

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

            first_prompt = ""
            for m in self._messages:
                if m.role == "user":
                    first_prompt = m.content
                    break

            record = build_session_record(
                record_type="chat",
                model=self.query_one("#model-input", Input).value.strip(),
                prompt=first_prompt,
                response=None,
                tokens=TokenUsage(
                    input=self._total_input_tokens,
                    output=self._total_output_tokens,
                    total=self._total_input_tokens + self._total_output_tokens,
                ),
                metadata={
                    "messages": [{"role": m.role, "content": m.content} for m in self._messages],
                    "system": self._system_prompt,
                    "session_id": self._session_id,
                },
            )
            # Override the auto-generated id with our session_id
            record = type(record)(
                id=self._session_id,
                type=record.type,
                timestamp=record.timestamp,
                model=record.model,
                prompt=record.prompt,
                response=record.response,
                tokens=record.tokens,
                metadata=record.metadata,
            )

            from alice_cli.auth import detect_username

            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, "chat")
        except Exception:  # noqa: BLE001
            pass  # Best-effort — don't crash on exit

    def action_pop_screen(self) -> None:
        """Persist session before leaving the screen."""
        self._persist_session()
        self.app.pop_screen()
