"""Textual UI application for ieum."""
# ruff: noqa: BLE001, PLR0912, PLR2004, S110

from __future__ import annotations

import asyncio
import uuid
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar

from ieum.config import get_package_path
from ieum.approval_policy import (
    APPROVAL_POLICY_FULL_TRUSTED,
    APPROVAL_POLICY_MANUAL,
    APPROVAL_POLICY_PLAN_ONLY,
    normalize_approval_policy,
)

from textual.app import App
from textual.binding import Binding, BindingType
from textual.containers import Container, VerticalScroll
from textual.css.query import NoMatches
from textual.events import Click, MouseUp, Resize
from textual.widgets import Static

from ieum.clipboard import copy_selection_to_clipboard
from ieum.textual_adapter import TextualUIAdapter
from ieum.widgets.chat_input import ChatInput
from ieum.widgets.loading import LoadingWidget
from ieum.widgets.messages import (
    SystemMessage,
    ToolCallMessage,
)
from ieum.widgets.status import StatusBar
from ieum.widgets.welcome import WelcomeBanner

from ieum.handlers.approval import ApprovalHandlerMixin
from ieum.handlers.chat import ChatHandlerMixin
from ieum.handlers.commands import CommandHandlerMixin
from ieum.handlers.env import EnvHandlerMixin
from ieum.handlers.interview import InterviewHandlerMixin
from ieum.handlers.mcp import McpHandlerMixin

if TYPE_CHECKING:
    from langgraph.pregel import Pregel
    from textual.app import ComposeResult
    from textual.worker import Worker


class TextualTokenTracker:
    """Token tracker that updates the status bar with tokens and cost."""

    def __init__(
        self,
        update_callback: callable,
        hide_callback: callable | None = None,
        cost_callback: callable | None = None,
    ) -> None:
        """Initialize with callbacks to update the display."""
        self._update_callback = update_callback
        self._hide_callback = hide_callback
        self._cost_callback = cost_callback
        self.current_context = 0
        self.session_cost: float = 0.0
        self._model_name: str | None = None

    def set_model_name(self, name: str) -> None:
        """Set current model name for cost calculation."""
        self._model_name = name

    def add(
        self,
        total_tokens: int,
        _output_tokens: int = 0,
        input_tokens: int = 0,
        output_tokens: int = 0,
    ) -> None:
        """Update token count and accumulate cost.

        Args:
            total_tokens: Total context tokens (input + output from usage_metadata)
            _output_tokens: Unused, kept for backwards compatibility
            input_tokens: Input/prompt tokens for cost calculation
            output_tokens: Output/completion tokens for cost calculation
        """
        self.current_context = total_tokens
        self._update_callback(self.current_context)

        # Calculate and accumulate cost
        if self._model_name and (input_tokens or output_tokens):
            from ieum.pricing import calculate_cost
            cost = calculate_cost(self._model_name, input_tokens, output_tokens)
            if cost > 0:
                self.session_cost += cost
                if self._cost_callback:
                    self._cost_callback(self.session_cost)

    def reset(self) -> None:
        """Reset token count (cost persists across resets)."""
        self.current_context = 0
        self._update_callback(0)

    def reset_cost(self) -> None:
        """Reset session cost."""
        self.session_cost = 0.0
        if self._cost_callback:
            self._cost_callback(0.0)

    def hide(self) -> None:
        """Hide the token display (e.g., during streaming)."""
        if self._hide_callback:
            self._hide_callback()

    def show(self) -> None:
        """Show the token display with current value (e.g., after interrupt)."""
        self._update_callback(self.current_context)


class TextualSessionState:
    """Session state for the Textual app."""

    def __init__(
        self,
        *,
        auto_approve: bool = False,
        approval_policy: str = "manual",
        thread_id: str | None = None,
    ) -> None:
        """Initialize session state.

        Args:
            auto_approve: Whether to auto-approve tool calls
            thread_id: Optional thread ID (generates 8-char hex if not provided)
        """
        self.auto_approve = auto_approve
        self.approval_policy = normalize_approval_policy(approval_policy)
        self.thread_id = thread_id if thread_id else uuid.uuid4().hex[:8]
        # Auto-commit state
        self.auto_commit: bool = False
        # Verify loop state
        self.verify_enabled: bool = False
        self.verify_config: Any = None  # VerifyConfig from verify_loop
        # Checkpoint restore state
        self._restore_checkpoint_id: str | None = None
        # Message counter for auto-memory triggers
        self._user_message_count: int = 0
        # Plan mode state
        self.plan_mode: bool = False
        self.current_plan: str | None = None
        self._plan_original_task: str | None = None
        self._saved_approval_policy: str | None = None

    def reset_thread(self) -> str:
        """Reset to a new thread. Returns the new thread_id."""
        self.thread_id = uuid.uuid4().hex[:8]
        return self.thread_id

    def set_approval_policy(self, policy: str) -> str:
        """Set session approval policy and return the normalized value."""
        self.approval_policy = normalize_approval_policy(policy)
        return self.approval_policy

    # --- Three-mode cycling: manual → auto → plan → manual ---

    def get_current_mode(self) -> str:
        """Return current mode label: 'manual', 'auto', or 'plan'."""
        if self.plan_mode:
            return "plan"
        if self.auto_approve:
            return "auto"
        return "manual"

    def cycle_mode(self) -> str:
        """Cycle to next mode. Returns the new mode label."""
        current = self.get_current_mode()
        if current == "manual":
            # manual → auto
            self.auto_approve = True
            return "auto"
        if current == "auto":
            # auto → plan
            self.auto_approve = False
            self.enter_plan_mode()
            return "plan"
        # plan → manual
        self.exit_plan_mode()
        return "manual"

    def enter_plan_mode(self) -> None:
        """Enter plan mode: backup policy, switch to plan_only."""
        self._saved_approval_policy = self.approval_policy
        self.approval_policy = APPROVAL_POLICY_PLAN_ONLY
        self.plan_mode = True
        self.current_plan = None

    def exit_plan_mode(self) -> None:
        """Exit plan mode: restore previous approval policy."""
        self.plan_mode = False
        if self._saved_approval_policy is not None:
            self.approval_policy = self._saved_approval_policy
            self._saved_approval_policy = None
        else:
            self.approval_policy = APPROVAL_POLICY_MANUAL


class IeumApp(
    ChatHandlerMixin,
    CommandHandlerMixin,
    ApprovalHandlerMixin,
    InterviewHandlerMixin,
    McpHandlerMixin,
    EnvHandlerMixin,
    App,
):
    """Main Textual application for ieum."""

    TITLE = "IEUM"
    CSS_PATH = get_package_path("app.tcss")
    ENABLE_COMMAND_PALETTE = False

    # Scroll speed (default is 3 lines per scroll event)
    SCROLL_SENSITIVITY_Y = 1.0

    # Timeout (seconds) for bash commands executed via ! prefix
    BASH_COMMAND_TIMEOUT: ClassVar[int] = 60

    BINDINGS: ClassVar[list[BindingType]] = [
        Binding("escape", "interrupt", "Interrupt", show=True, priority=True),
        Binding("ctrl+c", "quit_or_interrupt", "Quit/Interrupt", show=False),
        Binding("ctrl+d", "quit_app", "Quit", show=True, priority=True),
        Binding("ctrl+l,cmd+l", "refresh_screen", "Refresh", show=False, priority=True),
        Binding("ctrl+t,cmd+t", "toggle_auto_approve", "Toggle Auto-Approve", show=False),
        Binding(
            "shift+tab", "toggle_auto_approve", "Auto-Approve", show=True, priority=True
        ),
        Binding("ctrl+o,cmd+o", "toggle_tool_output", "Toggle Tool Output", show=False),
        # Approval menu keys (handled at App level for reliability)
        Binding("up", "approval_up", "Up", show=False),
        Binding("k", "approval_up", "Up", show=False),
        Binding("down", "approval_down", "Down", show=False),
        Binding("j", "approval_down", "Down", show=False),
        Binding("enter", "approval_select", "Select", show=False),
        Binding("y", "approval_yes", "Yes", show=False),
        Binding("1", "approval_yes", "Yes", show=False),
        Binding("n", "approval_no", "No", show=False),
        Binding("2", "approval_no", "No", show=False),
        Binding("a", "approval_auto", "Auto", show=False),
        Binding("3", "approval_auto", "Auto", show=False),
    ]

    def __init__(
        self,
        *,
        agent: Pregel | None = None,
        assistant_id: str | None = None,
        backend: Any = None,  # noqa: ANN401  # CompositeBackend
        auto_approve: bool = False,
        approval_policy: str = "manual",
        cwd: str | Path | None = None,
        thread_id: str | None = None,
        initial_prompt: str | None = None,
        max_context_tokens: int | None = None,
        setup_mode: bool = False,
        **kwargs: Any,
    ) -> None:
        """Initialize the IEUM application.

        Args:
            agent: Pre-configured LangGraph agent (optional for standalone mode)
            assistant_id: Agent identifier for memory storage
            backend: Backend for file operations
            auto_approve: Whether to start with auto-approve enabled
            cwd: Current working directory to display
            thread_id: Optional thread ID for session persistence
            initial_prompt: Optional prompt to auto-submit when session starts
            max_context_tokens: Maximum context window size from model profile
            setup_mode: Whether app is in setup mode (no LLM configured)
            **kwargs: Additional arguments passed to parent
        """
        super().__init__(**kwargs)
        self._agent = agent
        self._assistant_id = assistant_id
        self._backend = backend
        self._auto_approve = auto_approve
        self._approval_policy = normalize_approval_policy(approval_policy)
        self._cwd = str(cwd) if cwd else str(Path.cwd())
        # Avoid collision with App._thread_id
        self._lc_thread_id = thread_id
        self._initial_prompt = initial_prompt
        self._max_context_tokens = max_context_tokens  # Context window size
        self._setup_mode = setup_mode  # No LLM configured - show setup help
        self._current_model_name: str | None = None  # Track current model for hot-swap
        self._checkpointer: Any = None  # Store checkpointer for agent recreation
        self._store: Any = None  # Store for cross-thread persistent memories
        self._status_bar: StatusBar | None = None
        self._chat_input: ChatInput | None = None
        self._quit_pending = False
        self._session_state: TextualSessionState | None = None
        self._ui_adapter: TextualUIAdapter | None = None
        self._pending_approval_widget: Any = None
        self._pending_interview_widget: Any = None
        # Agent task tracking for interruption
        self._agent_worker: Worker[None] | None = None
        self._agent_running = False
        self._loading_widget: LoadingWidget | None = None
        self._token_tracker: TextualTokenTracker | None = None

    def compose(self) -> ComposeResult:
        """Compose the application layout."""
        # Main chat area with scrollable messages
        # VerticalScroll tracks user scroll intent for better auto-scroll behavior
        with VerticalScroll(id="chat"):
            yield WelcomeBanner(id="welcome-banner")
            yield Container(id="messages")
            with Container(id="bottom-app-container"):
                yield ChatInput(cwd=self._cwd, id="input-area")
            yield Static(id="chat-spacer")  # Fills remaining space below input

        # Status bar at bottom
        yield StatusBar(cwd=self._cwd, id="status-bar")

    async def on_mount(self) -> None:
        """Initialize components after mount."""
        from ieum.config import settings

        self._status_bar = self.query_one("#status-bar", StatusBar)
        self._chat_input = self.query_one("#input-area", ChatInput)

        # Set initial model name from settings
        if settings.model_name:
            self._status_bar.set_model_name(settings.model_name)
        elif settings.default_model:
            self._status_bar.set_model_name(settings.default_model)

        # Set initial max context tokens (for usage percentage display)
        if self._max_context_tokens:
            self._status_bar.set_max_tokens(self._max_context_tokens)

        # Set initial auto-approve state
        if self._auto_approve:
            self._status_bar.set_auto_approve(enabled=True)

        # Create session state
        self._session_state = TextualSessionState(
            auto_approve=self._auto_approve,
            approval_policy=self._approval_policy,
            thread_id=self._lc_thread_id,
        )

        # Create token tracker that updates status bar (with cost tracking)
        self._token_tracker = TextualTokenTracker(
            self._update_tokens, self._hide_tokens, self._update_cost,
        )

        # Set model name on tracker for cost calculation
        if settings.model_name:
            self._token_tracker.set_model_name(settings.model_name)
        elif settings.default_model:
            self._token_tracker.set_model_name(settings.default_model)

        # Create UI adapter if agent is provided
        if self._agent:
            self._ui_adapter = TextualUIAdapter(
                mount_message=self._mount_message,
                update_status=self._update_status,
                request_approval_fn=self._request_approval,
                on_auto_approve_enabled=self._on_auto_approve_enabled,
                scroll_to_bottom=self._scroll_chat_to_bottom,
                show_thinking=self._show_thinking,
                hide_thinking=self._hide_thinking,
                request_interview_fn=self._request_interview,
            )
            self._ui_adapter.set_token_tracker(self._token_tracker)

        # Focus the input (autocomplete is now built into ChatInput)
        self._chat_input.focus_input()

        # Size the spacer to fill remaining viewport below input
        self.call_after_refresh(self._size_initial_spacer)

        # Show setup message if in setup mode (no LLM configured)
        if self._setup_mode:
            self.call_after_refresh(lambda: asyncio.create_task(self._show_setup_message()))
        # Load thread history if resuming a session
        elif self._lc_thread_id and self._agent:
            self.call_after_refresh(lambda: asyncio.create_task(self._load_thread_history()))
        # Auto-submit initial prompt if provided (but not when resuming - let user see history first)
        elif self._initial_prompt and self._initial_prompt.strip():
            # Use call_after_refresh to ensure UI is fully mounted before submitting
            self.call_after_refresh(
                lambda: asyncio.create_task(self._handle_user_message(self._initial_prompt))
            )

    def on_resize(self, _event: Resize) -> None:
        """Handle terminal resize to recalculate layout."""
        # Force full UI refresh on resize
        self.refresh(repaint=True, layout=True)

        try:
            self.query_one("#chat-spacer", Static)
            # Spacer exists, recalculate its height
            self.call_after_refresh(self._size_initial_spacer)
        except NoMatches:
            pass  # Spacer already removed, no action needed

    # =========================================================================
    # UI Helper Methods
    # =========================================================================

    def _update_status(self, message: str) -> None:
        """Update the status bar with a message."""
        if self._status_bar:
            self._status_bar.set_status_message(message)

    def _update_tokens(self, count: int) -> None:
        """Update the token count in status bar."""
        if self._status_bar:
            self._status_bar.set_tokens(count)

    def _hide_tokens(self) -> None:
        """Hide the token display during streaming."""
        if self._status_bar:
            self._status_bar.hide_tokens()

    def _update_cost(self, cost: float) -> None:
        """Update the cost display in status bar."""
        if self._status_bar:
            self._status_bar.set_cost(cost)

    def _scroll_chat_to_bottom(self) -> None:
        """Scroll chat to bottom using sticky scroll pattern.

        Only scrolls if user is already at/near the bottom.
        This prevents dragging the user back if they've scrolled up to read.
        """
        chat = self.query_one("#chat", VerticalScroll)

        # Nothing to scroll if content fits in viewport
        if chat.max_scroll_y <= 0:
            return

        # Sticky scroll: only scroll to bottom if user is near the bottom
        # "Near" means within 100 pixels of the bottom (about 6-7 lines)
        distance_from_bottom = chat.max_scroll_y - chat.scroll_y
        if distance_from_bottom < 100:
            chat.scroll_end(animate=False)

    async def _show_thinking(self) -> None:
        """Show or reposition the thinking spinner at the bottom of messages."""
        if self._loading_widget:
            await self._loading_widget.remove()
            self._loading_widget = None

        self._loading_widget = LoadingWidget("Thinking")
        messages = self.query_one("#messages", Container)
        await messages.mount(self._loading_widget)
        # NOTE: Don't call _scroll_chat_to_bottom() here - it would re-anchor
        # and drag user back to bottom if they've scrolled away during streaming

    async def _hide_thinking(self) -> None:
        """Hide the thinking spinner."""
        if self._loading_widget:
            await self._loading_widget.remove()
            self._loading_widget = None

    def _size_initial_spacer(self) -> None:
        """Size the spacer to fill remaining viewport below input."""
        try:
            chat = self.query_one("#chat", VerticalScroll)
            welcome = self.query_one("#welcome-banner", WelcomeBanner)
            input_container = self.query_one("#bottom-app-container", Container)
            spacer = self.query_one("#chat-spacer", Static)
            content_height = welcome.size.height + input_container.size.height + 4
            spacer_height = chat.size.height - content_height
            spacer.styles.height = max(0, spacer_height)
        except NoMatches:
            # Spacer may have been removed already (e.g., when resuming a session)
            pass

    async def _remove_spacer(self) -> None:
        """Remove the initial spacer when first message is sent."""
        try:
            spacer = self.query_one("#chat-spacer", Static)
            await spacer.remove()
        except NoMatches:
            pass

    async def _mount_message(self, widget: Static) -> None:
        """Mount a message widget to the messages area.

        Args:
            widget: The message widget to mount
        """
        await self._remove_spacer()
        messages = self.query_one("#messages", Container)
        await messages.mount(widget)
        # Scroll to keep input bar visible
        input_container = self.query_one("#bottom-app-container", Container)
        input_container.scroll_visible()

    async def _clear_messages(self) -> None:
        """Clear the messages area."""
        try:
            messages = self.query_one("#messages", Container)
            await messages.remove_children()
        except NoMatches:
            # Widget not found - can happen during shutdown
            pass

    # =========================================================================
    # Event Handlers
    # =========================================================================

    async def on_chat_input_submitted(self, event: ChatInput.Submitted) -> None:
        """Handle submitted input from ChatInput widget."""
        value = event.value
        mode = event.mode

        # Reset quit pending state on any input
        self._quit_pending = False

        # Handle different modes
        if mode == "bash":
            # Bash command - strip the ! prefix
            await self._handle_bash_command(value.removeprefix("!"))
        elif mode == "command":
            # Slash command
            await self._handle_command(value)
        else:
            # Normal message - will be sent to agent
            await self._handle_user_message(value)

    def on_chat_input_mode_changed(self, event: ChatInput.ModeChanged) -> None:
        """Update status bar when input mode changes."""
        if self._status_bar:
            self._status_bar.set_mode(event.mode)

    # =========================================================================
    # Action Methods (Keybindings)
    # =========================================================================

    def action_quit_or_interrupt(self) -> None:
        """Handle Ctrl+C - interrupt agent, reject approval, or quit on double press.

        Priority order:
        1. If agent is running, interrupt it (preserve input)
        2. If approval menu is active, reject it
        3. If double press (quit_pending), quit
        4. Otherwise show quit hint
        """
        # If agent is running, interrupt it
        if self._agent_running and self._agent_worker:
            self._agent_worker.cancel()
            self._quit_pending = False
            return

        # If approval menu is active, reject it
        if self._pending_approval_widget:
            self._pending_approval_widget.action_select_reject()
            self._quit_pending = False
            return

        # If interview widget is active, cancel it
        if self._pending_interview_widget:
            self._pending_interview_widget.action_cancel()
            self._quit_pending = False
            return

        # Double Ctrl+C to quit
        if self._quit_pending:
            self.exit()
        else:
            self._quit_pending = True
            self.notify("Press Ctrl+C again to quit", timeout=3)

    def action_interrupt(self) -> None:
        """Handle escape key - interrupt agent, reject approval, or cancel interview.

        This is the primary way to stop a running agent.
        """
        # If agent is running, interrupt it
        if self._agent_running and self._agent_worker:
            self._agent_worker.cancel()
            return

        # If approval menu is active, reject it
        if self._pending_approval_widget:
            self._pending_approval_widget.action_select_reject()
            return

        # If interview widget is active, cancel it
        if self._pending_interview_widget:
            self._pending_interview_widget.action_cancel()

    def action_quit_app(self) -> None:
        """Handle quit action (Ctrl+D)."""
        self.exit()

    def action_refresh_screen(self) -> None:
        """Handle screen refresh (Ctrl+L).

        Forces a full screen redraw to fix display corruption
        that can occur on some terminals (e.g., PowerShell after Alt+Tab).
        """
        self.refresh(repaint=True, layout=True)

    def action_toggle_auto_approve(self) -> None:
        """Cycle through three modes: manual → auto → plan.

        Skips toggle if an InterviewWidget is currently focused, allowing
        shift+tab to propagate to the interview widget for tab navigation.
        """
        from ieum.widgets.interview import InterviewWidget

        # If an interview widget is active, don't toggle - let shift+tab
        # propagate to the interview widget for tab navigation
        if self._pending_interview_widget is not None:
            focused = self.focused
            if focused is not None and isinstance(focused, InterviewWidget):
                return

        if not self._session_state:
            return

        new_mode = self._session_state.cycle_mode()
        self._auto_approve = self._session_state.auto_approve

        if self._status_bar:
            self._status_bar.set_approval_mode(new_mode)

    def action_toggle_tool_output(self) -> None:
        """Toggle expand/collapse of the most recent tool output."""
        # Find all tool messages with output, get the most recent one
        try:
            tool_messages = list(self.query(ToolCallMessage))
            # Find ones with output, toggle the most recent
            for tool_msg in reversed(tool_messages):
                if tool_msg.has_output:
                    tool_msg.toggle_output()
                    return
        except Exception:
            pass

    def _is_input_focused(self) -> bool:
        """Check if the chat input (or its text area) has focus."""
        if not self._chat_input:
            return False
        focused = self.focused
        if focused is None:
            return False
        # Check if focused widget is the text area inside chat input
        return focused.id == "chat-input" or focused in self._chat_input.walk_children()

    def on_click(self, _event: Click) -> None:
        """Handle clicks anywhere in the terminal to focus on the command line."""
        if not self._chat_input:
            return
        # Don't steal focus from approval or interview widgets
        if self._pending_approval_widget or self._pending_interview_widget:
            return
        self.call_after_refresh(self._chat_input.focus_input)

    def on_mouse_up(self, event: MouseUp) -> None:  # noqa: ARG002
        """Copy selection to clipboard on mouse release."""
        copy_selection_to_clipboard(self)


async def run_textual_app(
    *,
    agent: Pregel | None = None,
    assistant_id: str | None = None,
    backend: Any = None,  # noqa: ANN401  # CompositeBackend
    auto_approve: bool = False,
    approval_policy: str = "manual",
    cwd: str | Path | None = None,
    thread_id: str | None = None,
    initial_prompt: str | None = None,
    checkpointer: Any = None,  # noqa: ANN401  # BaseCheckpointSaver
    store: Any = None,  # noqa: ANN401  # BaseStore
    max_context_tokens: int | None = None,
    setup_mode: bool = False,
) -> None:
    """Run the Textual application.

    Args:
        agent: Pre-configured LangGraph agent (optional)
        assistant_id: Agent identifier for memory storage
        backend: Backend for file operations
        auto_approve: Whether to start with auto-approve enabled
        approval_policy: Approval preset policy for HITL decisions
        cwd: Current working directory to display
        thread_id: Optional thread ID for session persistence
        initial_prompt: Optional prompt to auto-submit when session starts
        checkpointer: Checkpointer for session persistence (for model hot-swap)
        store: Store for cross-thread persistent memories (for model hot-swap)
        max_context_tokens: Maximum context window size from model profile
        setup_mode: Whether to start in setup mode (no LLM configured)
    """
    app = IeumApp(
        agent=agent,
        assistant_id=assistant_id,
        backend=backend,
        auto_approve=auto_approve,
        approval_policy=approval_policy,
        cwd=cwd,
        thread_id=thread_id,
        initial_prompt=initial_prompt,
        max_context_tokens=max_context_tokens,
        setup_mode=setup_mode,
    )
    app._checkpointer = checkpointer  # Store for model hot-swap
    app._store = store  # Store for model hot-swap
    await app.run_async()


if __name__ == "__main__":
    import asyncio

    asyncio.run(run_textual_app())
