"""Message widgets for ieum."""

from __future__ import annotations

from datetime import datetime
from time import time
from typing import TYPE_CHECKING, Any, ClassVar

from rich.text import Text
from textual.containers import Vertical
from textual.events import Click
from textual.timer import Timer
from textual.widgets import Markdown, Static
from textual.widgets._markdown import MarkdownStream

from ieum.i18n import t
from ieum.ui import format_tool_display
from ieum.widgets.diff import format_diff_textual

if TYPE_CHECKING:
    from textual.app import ComposeResult

# Maximum number of tool arguments to display inline
_MAX_INLINE_ARGS = 3

# Tools that have their key info already in the header (no need for args line)
_TOOLS_WITH_HEADER_INFO: set[str] = {
    # Filesystem tools
    "ls",
    "read_file",
    "write_file",
    "edit_file",
    "glob",
    "grep",
    # Shell tools
    "shell",
    # Web tools
    "web_search",
    "fetch_url",
    "http_request",
    # Agent tools
    "task",
    "write_todos",
}


class UserMessage(Vertical):
    """Widget displaying a user message."""

    DEFAULT_CSS = """
    UserMessage {
        height: auto;
        padding: 0 1;
        margin: 1 0 0 0;
        background: transparent;
        border-left: wide #10b981;
    }

    UserMessage .message-timestamp {
        text-align: right;
        color: $text-muted;
        width: 100%;
    }
    """

    def __init__(self, content: str, timestamp: datetime | None = None, **kwargs: Any) -> None:
        """Initialize a user message.

        Args:
            content: The message content
            timestamp: Message creation time (defaults to now)
            **kwargs: Additional arguments passed to parent
        """
        super().__init__(**kwargs)
        self._content = content
        self._timestamp = timestamp or datetime.now()

    def compose(self) -> ComposeResult:
        """Compose the user message layout."""
        text = Text()
        text.append("> ", style="bold #10b981")
        text.append(self._content)
        yield Static(text)
        # Timestamp at bottom right
        time_str = self._timestamp.strftime("%m/%d %H:%M")
        yield Static(time_str, classes="message-timestamp")


class AssistantMessage(Vertical):
    """Widget displaying an assistant message with markdown support.

    Uses MarkdownStream for smoother streaming instead of re-rendering
    the full content on each update.
    """

    DEFAULT_CSS = """
    AssistantMessage {
        height: auto;
        padding: 0 1;
        margin: 1 0 0 0;
    }

    AssistantMessage .message-timestamp {
        text-align: right;
        color: $text-muted;
        width: 100%;
    }

    AssistantMessage .message-meta {
        text-align: right;
        color: $text-muted;
        width: 100%;
    }

    AssistantMessage .token-count {
        color: #6b7280;
    }

    AssistantMessage Markdown {
        padding: 0;
        margin: 0;
    }
    """

    def __init__(self, content: str = "", timestamp: datetime | None = None, **kwargs: Any) -> None:
        """Initialize an assistant message.

        Args:
            content: Initial markdown content
            timestamp: Message creation time (defaults to now)
            **kwargs: Additional arguments passed to parent
        """
        super().__init__(**kwargs)
        self._content = content
        self._timestamp = timestamp or datetime.now()
        self._markdown: Markdown | None = None
        self._stream: MarkdownStream | None = None
        self._tokens: int = 0
        self._meta_widget: Static | None = None

    def compose(self) -> ComposeResult:
        """Compose the assistant message layout."""
        yield Markdown("", id="assistant-content")
        # Timestamp and tokens at bottom right
        time_str = self._timestamp.strftime("%m/%d %H:%M")
        yield Static(time_str, classes="message-meta", id="message-meta")

    def on_mount(self) -> None:
        """Store reference to markdown widget."""
        self._markdown = self.query_one("#assistant-content", Markdown)
        self._meta_widget = self.query_one("#message-meta", Static)

    def set_tokens(self, tokens: int) -> None:
        """Set the token count and update display.

        Args:
            tokens: Number of tokens used for this message
        """
        self._tokens = tokens
        self._update_meta_display()

    def _update_meta_display(self) -> None:
        """Update the meta display with timestamp and tokens."""
        if self._meta_widget is None:
            return

        time_str = self._timestamp.strftime("%m/%d %H:%M")
        meta = Text()
        meta.append(time_str, style="#6b7280")

        if self._tokens > 0:
            meta.append(" • ", style="#4b5563")
            if self._tokens >= 1000:
                token_str = f"{self._tokens / 1000:.1f}k"
            else:
                token_str = str(self._tokens)
            meta.append(f"{token_str} tokens", style="#6b7280")

        self._meta_widget.update(meta)

    def _get_markdown(self) -> Markdown:
        """Get the markdown widget, querying if not cached."""
        if self._markdown is None:
            self._markdown = self.query_one("#assistant-content", Markdown)
        return self._markdown

    def _ensure_stream(self) -> MarkdownStream:
        """Ensure the markdown stream is initialized."""
        if self._stream is None:
            self._stream = Markdown.get_stream(self._get_markdown())
        return self._stream

    async def append_content(self, text: str) -> None:
        """Append content to the message (for streaming).

        Uses MarkdownStream for smoother rendering instead of re-rendering
        the full content on each chunk.

        Args:
            text: Text to append
        """
        if not text:
            return
        self._content += text
        stream = self._ensure_stream()
        await stream.write(text)

    async def write_initial_content(self) -> None:
        """Write initial content if provided at construction time."""
        if self._content:
            stream = self._ensure_stream()
            await stream.write(self._content)

    async def stop_stream(self) -> None:
        """Stop the streaming and finalize the content."""
        if self._stream is not None:
            await self._stream.stop()
            self._stream = None

    async def set_content(self, content: str) -> None:
        """Set the full message content.

        This stops any active stream and sets content directly.

        Args:
            content: The markdown content to display
        """
        await self.stop_stream()
        self._content = content
        if self._markdown:
            await self._markdown.update(content)


class ToolCallMessage(Vertical):
    """Widget displaying a tool call with collapsible output.

    Tool outputs are shown as a 3-line preview by default.
    Press Ctrl+O to expand/collapse the full output.
    Shows an animated "Running..." indicator while the tool is executing.
    """

    DEFAULT_CSS = """
    ToolCallMessage {
        height: auto;
        padding: 0 1;
        margin: 0 0 1 0;
        background: transparent;
        border-left: wide #3b3b3b;
    }

    ToolCallMessage .tool-header {
        height: auto;
    }

    ToolCallMessage .tool-args {
        color: #6b7280;
        margin-left: 3;
    }

    ToolCallMessage .tool-status {
        margin-left: 3;
    }

    ToolCallMessage .tool-status.pending {
        color: #f59e0b;
    }

    ToolCallMessage .tool-status.success {
        color: #10b981;
    }

    ToolCallMessage .tool-status.error {
        color: #ef4444;
    }

    ToolCallMessage .tool-status.rejected {
        color: #f59e0b;
    }

    ToolCallMessage .tool-output {
        margin-left: 3;
        margin-top: 0;
        padding: 0;
        height: auto;
    }

    ToolCallMessage .tool-output-preview {
        margin-left: 3;
        margin-top: 0;
    }

    ToolCallMessage .tool-output-hint {
        margin-left: 3;
        color: #6b7280;
    }

    ToolCallMessage:hover {
        border-left: wide #525252;
    }
    """

    # Spinner frames for running animation
    _SPINNER_FRAMES: ClassVar[tuple[str, ...]] = ("⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏")

    # Max lines/chars to show in preview mode
    _PREVIEW_LINES = 6
    _PREVIEW_CHARS = 400

    def __init__(
        self,
        tool_name: str,
        args: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        """Initialize a tool call message.

        Args:
            tool_name: Name of the tool being called
            args: Tool arguments (optional)
            **kwargs: Additional arguments passed to parent
        """
        super().__init__(**kwargs)
        self._tool_name = tool_name
        self._args = args or {}
        self._status = "pending"  # Waiting for approval or auto-approve
        self._output: str = ""
        self._expanded: bool = False
        # Widget references (set in on_mount)
        self._status_widget: Static | None = None
        self._preview_widget: Static | None = None
        self._hint_widget: Static | None = None
        self._full_widget: Static | None = None
        # Animation state
        self._spinner_position = 0
        self._start_time: float | None = None
        self._animation_timer: Timer | None = None

    def compose(self) -> ComposeResult:
        """Compose the tool call message layout."""
        tool_label = format_tool_display(self._tool_name, self._args)
        yield Static(
            f"[bold #f59e0b]{tool_label}[/bold #f59e0b]",
            classes="tool-header",
        )
        # Only show args for tools where header doesn't capture the key info
        if self._tool_name not in _TOOLS_WITH_HEADER_INFO:
            args = self._filtered_args()
            if args:
                args_str = ", ".join(f"{k}={v!r}" for k, v in list(args.items())[:_MAX_INLINE_ARGS])
                if len(args) > _MAX_INLINE_ARGS:
                    args_str += ", ..."
                yield Static(f"[dim]({args_str})[/dim]", classes="tool-args")
        # Status - shows running animation while pending, then final status
        yield Static("", classes="tool-status", id="status")
        # Output area - hidden initially, shown when output is set
        yield Static("", classes="tool-output-preview", id="output-preview")
        yield Static("", classes="tool-output-hint", id="output-hint")
        yield Static("", classes="tool-output", id="output-full")

    def on_mount(self) -> None:
        """Cache widget references and hide all status/output areas initially."""
        self._status_widget = self.query_one("#status", Static)
        self._preview_widget = self.query_one("#output-preview", Static)
        self._hint_widget = self.query_one("#output-hint", Static)
        self._full_widget = self.query_one("#output-full", Static)
        # Hide everything initially - status only shown when running or on error/reject
        self._status_widget.display = False
        self._preview_widget.display = False
        self._hint_widget.display = False
        self._full_widget.display = False

    def set_running(self) -> None:
        """Mark the tool as running (approved and executing).

        Call this when approval is granted to start the running animation.
        """
        if self._status == "running":
            return  # Already running

        self._status = "running"
        self._start_time = time()
        if self._status_widget:
            self._status_widget.add_class("pending")
            self._status_widget.display = True
        self._update_running_animation()
        self._animation_timer = self.set_interval(0.1, self._update_running_animation)

    # Thresholds (seconds) for spinner timeout warnings
    _SPINNER_WARNING_SECS = 120
    _SPINNER_ERROR_SECS = 300

    def _update_running_animation(self) -> None:
        """Update the running spinner animation.

        Changes color and adds a warning after extended durations:
        - After 120s: yellow/warning with hint text
        - After 300s: red/error with hint text
        """
        if self._status != "running" or self._status_widget is None:
            return

        frame = self._SPINNER_FRAMES[self._spinner_position]
        self._spinner_position = (self._spinner_position + 1) % len(self._SPINNER_FRAMES)

        elapsed_text = ""
        color = "yellow"
        suffix = ""
        if self._start_time is not None:
            elapsed_secs = int(time() - self._start_time)
            elapsed_text = f" ({elapsed_secs}s)"
            if elapsed_secs >= self._SPINNER_ERROR_SECS:
                color = "red"
                suffix = f" {t('tool.taking_longer')}"
            elif elapsed_secs >= self._SPINNER_WARNING_SECS:
                color = "#f59e0b"
                suffix = f" {t('tool.taking_longer')}"

        self._status_widget.update(
            f"[{color}]{frame} {t('tool.running')}{elapsed_text}{suffix}[/{color}]"
        )

    def _stop_animation(self) -> None:
        """Stop the running animation."""
        if self._animation_timer is not None:
            self._animation_timer.stop()
            self._animation_timer = None

    def set_success(self, result: str = "") -> None:
        """Mark the tool call as successful.

        Args:
            result: Tool output/result to display
        """
        self._stop_animation()
        self._status = "success"
        self._output = result
        if self._status_widget:
            self._status_widget.remove_class("pending")
            # Hide status on success - output speaks for itself
            self._status_widget.display = False
        self._update_output_display()

    def set_error(self, error: str) -> None:
        """Mark the tool call as failed.

        Args:
            error: Error message
        """
        self._stop_animation()
        self._status = "error"
        self._output = error
        if self._status_widget:
            self._status_widget.remove_class("pending")
            self._status_widget.add_class("error")
            self._status_widget.update(f"[red]✗ {t('tool.error')}[/red]")
            self._status_widget.display = True
        # Always show full error - errors should be visible
        self._expanded = True
        self._update_output_display()

    def set_rejected(self) -> None:
        """Mark the tool call as rejected by user."""
        self._stop_animation()
        self._status = "rejected"
        if self._status_widget:
            self._status_widget.remove_class("pending")
            self._status_widget.add_class("rejected")
            self._status_widget.update(f"[yellow]✗ {t('tool.rejected')}[/yellow]")
            self._status_widget.display = True

    def set_skipped(self) -> None:
        """Mark the tool call as skipped (due to another rejection)."""
        self._stop_animation()
        self._status = "skipped"
        if self._status_widget:
            self._status_widget.remove_class("pending")
            self._status_widget.add_class("rejected")  # Use same styling as rejected
            self._status_widget.update(f"[dim]- {t('tool.skipped')}[/dim]")
            self._status_widget.display = True

    def toggle_output(self) -> None:
        """Toggle between preview and full output display."""
        if not self._output:
            return
        self._expanded = not self._expanded
        self._update_output_display()

    def on_click(self, event: Click) -> None:
        """Handle click to toggle output expansion."""
        event.stop()  # Prevent click from bubbling up and scrolling
        self.toggle_output()

    def _format_output(self, output: str, *, is_preview: bool = False) -> str:
        """Format tool output based on tool type for nicer display.

        Args:
            output: Raw output string
            is_preview: Whether this is for preview (truncated) display

        Returns:
            Formatted output string with Rich markup
        """
        from ieum.widgets.tool_formatters import ToolOutputFormatter

        return ToolOutputFormatter.format(self._tool_name, output, is_preview=is_preview)

    def _escape_markup(self, text: str) -> str:
        """Escape Rich markup characters."""
        from ieum.widgets.tool_formatters import escape_markup

        return escape_markup(text)

    def _update_output_display(self) -> None:
        """Update the output display based on expanded state."""
        if not self._output or not self._preview_widget:
            return

        output_stripped = self._output.strip()
        lines = output_stripped.split("\n")
        total_lines = len(lines)
        total_chars = len(output_stripped)

        # Truncate if too many lines OR too many characters
        needs_truncation = total_lines > self._PREVIEW_LINES or total_chars > self._PREVIEW_CHARS

        if self._expanded:
            # Show full output with formatting
            self._preview_widget.display = False
            formatted = self._format_output(self._output, is_preview=False)
            self._full_widget.update(formatted)
            self._full_widget.display = True
            # Show collapse hint
            self._hint_widget.update(f"[dim italic]{t('tool.click_collapse')}[/dim italic]")
            self._hint_widget.display = True
        else:
            # Show preview
            self._full_widget.display = False
            if needs_truncation:
                # Show formatted preview
                formatted_preview = self._format_output(self._output, is_preview=True)
                self._preview_widget.update(formatted_preview)
                self._preview_widget.display = True

                # Show expand hint
                self._hint_widget.update(f"[dim italic]{t('tool.click_expand')}[/dim italic]")
                self._hint_widget.display = True
            elif output_stripped:
                # Output fits in preview, show formatted
                formatted = self._format_output(output_stripped, is_preview=False)
                self._preview_widget.update(formatted)
                self._preview_widget.display = True
                self._hint_widget.display = False
            else:
                self._preview_widget.display = False
                self._hint_widget.display = False

    @property
    def has_output(self) -> bool:
        """Check if this tool message has output to display."""
        return bool(self._output)

    def _filtered_args(self) -> dict[str, Any]:
        """Filter large tool args for display."""
        if self._tool_name not in {"write_file", "edit_file"}:
            return self._args

        filtered: dict[str, Any] = {}
        for key in ("file_path", "path", "replace_all"):
            if key in self._args:
                filtered[key] = self._args[key]
        return filtered


class DiffMessage(Static):
    """Widget displaying a diff with syntax highlighting."""

    DEFAULT_CSS = """
    DiffMessage {
        height: auto;
        padding: 1;
        margin: 1 0;
        background: $surface;
        border: solid $primary;
    }

    DiffMessage .diff-header {
        text-style: bold;
        margin-bottom: 1;
    }

    DiffMessage .diff-add {
        color: #10b981;
        background: #10b98120;
    }

    DiffMessage .diff-remove {
        color: #ef4444;
        background: #ef444420;
    }

    DiffMessage .diff-context {
        color: $text-muted;
    }

    DiffMessage .diff-hunk {
        color: $secondary;
        text-style: bold;
    }
    """

    def __init__(self, diff_content: str, file_path: str = "", **kwargs: Any) -> None:
        """Initialize a diff message.

        Args:
            diff_content: The unified diff content
            file_path: Path to the file being modified
            **kwargs: Additional arguments passed to parent
        """
        super().__init__(**kwargs)
        self._diff_content = diff_content
        self._file_path = file_path

    def compose(self) -> ComposeResult:
        """Compose the diff message layout."""
        if self._file_path:
            yield Static(f"[bold]File: {self._file_path}[/bold]", classes="diff-header")

        # Render the diff with enhanced formatting
        rendered = format_diff_textual(self._diff_content, max_lines=100)
        yield Static(rendered)


class ErrorMessage(Static):
    """Widget displaying an error message."""

    DEFAULT_CSS = """
    ErrorMessage {
        height: auto;
        padding: 1;
        margin: 1 0;
        background: #7f1d1d;
        color: white;
        border-left: thick $error;
    }
    """

    def __init__(self, error: str, **kwargs: Any) -> None:
        """Initialize an error message.

        Args:
            error: The error message
            **kwargs: Additional arguments passed to parent
        """
        # Use Text object to combine styled prefix with unstyled error content
        text = Text(t("error.prefix"), style="bold red")
        text.append(error)
        super().__init__(text, **kwargs)


class SystemMessage(Vertical):
    """Widget displaying a system message."""

    DEFAULT_CSS = """
    SystemMessage {
        height: auto;
        padding: 0 1;
        margin: 1 0;
        color: $text-muted;
        text-style: italic;
    }

    SystemMessage .message-timestamp {
        text-align: right;
        color: $text-muted;
        width: 100%;
    }
    """

    def __init__(self, message: str, timestamp: datetime | None = None, **kwargs: Any) -> None:
        """Initialize a system message.

        Args:
            message: The system message
            timestamp: Message creation time (defaults to now)
            **kwargs: Additional arguments passed to parent
        """
        super().__init__(**kwargs)
        self._message = message
        self._timestamp = timestamp or datetime.now()

    def compose(self) -> ComposeResult:
        """Compose the system message layout."""
        # Use Text object to safely render message without markup parsing
        yield Static(Text(self._message, style="dim italic"))
        # Timestamp at bottom right
        time_str = self._timestamp.strftime("%m/%d %H:%M")
        yield Static(time_str, classes="message-timestamp")
