"""Rich-based progress display for MCP Agent."""

import time
from contextlib import contextmanager
from typing import Any

from rich.console import Console
from rich.progress import Progress, ProgressColumn, Task, TaskID, TextColumn
from rich.spinner import Spinner
from rich.table import Column
from rich.text import Text

from fast_agent.event_progress import ProgressAction, ProgressEvent
from fast_agent.ui.console import console as default_console
from fast_agent.ui.console import ensure_blocking_console


class SpinnerDescriptionColumn(ProgressColumn):
    """Render the task description with an inline spinner (no column padding gap)."""

    def __init__(
        self,
        *,
        spinner_name: str = "dots",
        spinner_style: str | None = "progress.spinner",
        speed: float = 1.0,
        finished_text: str = " ",
        description_style: str = "progress.description",
        markup: bool = True,
        table_column: Column | None = None,
    ) -> None:
        self.spinner = Spinner(spinner_name, style=spinner_style, speed=speed)
        self.finished_text = Text.from_markup(finished_text)
        self.description_style = description_style
        self.markup = markup
        super().__init__(table_column=table_column or Column(no_wrap=True))

    def render(self, task: "Task") -> Text:
        description_markup = f"[{self.description_style}]{task.description}▎"
        if self.markup:
            description_text = Text.from_markup(description_markup)
        else:
            description_text = Text(description_markup, style=self.description_style)

        if task.finished:
            spinner_text = self.finished_text
        else:
            rendered = self.spinner.render(task.get_time())
            spinner_text = rendered if isinstance(rendered, Text) else Text(str(rendered))

        return Text.assemble(description_text, spinner_text)


class RichProgressDisplay:
    """Rich-based display for progress events."""

    def __init__(self, console: Console | None = None) -> None:
        """Initialize the progress display."""
        self.console = console or default_console
        self._taskmap: dict[str, TaskID] = {}
        self._description_spinner = SpinnerDescriptionColumn(spinner_name="dots3")
        self._progress = Progress(
            self._description_spinner,
            TextColumn(text_format="{task.fields[target]:<16}", style="Bold Blue"),
            TextColumn(text_format="{task.fields[details]}", style="dim white"),
            console=self.console,
            transient=False,
        )
        self._paused = False

    def start(self) -> None:
        """start"""
        ensure_blocking_console()
        self._progress.start()

    def stop(self) -> None:
        """Stop and clear the progress display."""
        ensure_blocking_console()
        # Set paused first to prevent race with incoming updates
        self._paused = True
        # Hide all tasks before stopping (like pause does)
        for task in self._progress.tasks:
            task.visible = False
        self._progress.stop()

    def pause(self) -> None:
        """Pause the progress display."""
        if not self._paused:
            ensure_blocking_console()
            self._paused = True
            for task in self._progress.tasks:
                task.visible = False
            self._progress.stop()

    def resume(self) -> None:
        """Resume the progress display."""
        if self._paused:
            try:
                live_stack = getattr(self.console, "_live_stack", [])
            except Exception:
                live_stack = []

            if live_stack and any(live is not self._progress.live for live in live_stack):
                return
            ensure_blocking_console()
            if getattr(self._progress.live, "_nested", False):
                self._progress.live._nested = False
            for task in self._progress.tasks:
                task.visible = True
            self._paused = False
            self._progress.start()

    def hide_task(self, task_name: str) -> None:
        """Hide an existing task from the progress display by name."""
        task_id = self._taskmap.get(task_name)
        if task_id is None:
            return
        for task in self._progress.tasks:
            if task.id == task_id:
                task.visible = False
                break

    @contextmanager
    def paused(self):
        """Context manager for temporarily pausing the display."""
        was_paused = self._paused
        self.pause()
        try:
            yield
        finally:
            if not was_paused:
                self.resume()

    def _get_action_style(self, action: ProgressAction) -> str:
        """Map actions to appropriate styles."""
        return {
            ProgressAction.STARTING: "bold yellow",
            ProgressAction.LOADED: "dim green",
            ProgressAction.INITIALIZED: "dim green",
            ProgressAction.CHATTING: "bold blue",
            ProgressAction.STREAMING: "bold green",  # Assistant Colour
            ProgressAction.THINKING: "bold yellow",  # Assistant Colour
            ProgressAction.ROUTING: "bold blue",
            ProgressAction.PLANNING: "bold blue",
            ProgressAction.READY: "dim green",
            ProgressAction.CALLING_TOOL: "bold magenta",
            ProgressAction.TOOL_PROGRESS: "bold magenta",
            ProgressAction.FINISHED: "black on green",
            ProgressAction.SHUTDOWN: "black on red",
            ProgressAction.AGGREGATOR_INITIALIZED: "bold green",
            ProgressAction.FATAL_ERROR: "black on red",
        }.get(action, "white")

    def update(self, event: ProgressEvent) -> None:
        """Update the progress display with a new event."""
        # Skip updates when display is paused (e.g., during streaming)
        if self._paused:
            return

        task_name = event.agent_name or "default"

        # Create new task if needed
        if task_name not in self._taskmap:
            task_id = self._progress.add_task(
                "",
                total=None,
                target=event.target or task_name,
                details=event.details or "",
                task_name=task_name,
            )
            self._taskmap[task_name] = task_id
        else:
            task_id = self._taskmap[task_name]

        # Ensure no None values in the update
        # For streaming, use custom description immediately to avoid flashing
        if (
            event.action == ProgressAction.STREAMING or event.action == ProgressAction.THINKING
        ) and event.streaming_tokens:
            # Account for [dim][/dim] tags (11 characters) in padding calculation
            formatted_tokens = f"▎[dim]◀[/dim] {event.streaming_tokens.strip()}".ljust(17 + 11)
            description = f"[{self._get_action_style(event.action)}]{formatted_tokens}"
        elif event.action == ProgressAction.CHATTING:
            # Add special formatting for chatting with dimmed arrow
            formatted_text = f"▎[dim]▶[/dim] {event.action.value.strip()}".ljust(17 + 11)
            description = f"[{self._get_action_style(event.action)}]{formatted_text}"
        elif event.action == ProgressAction.CALLING_TOOL:
            # Add special formatting for calling tool with dimmed arrow
            formatted_text = f"▎[dim]◀[/dim] {event.action.value}".ljust(17 + 11)
            description = f"[{self._get_action_style(event.action)}]{formatted_text}"
        elif event.action == ProgressAction.TOOL_PROGRESS:
            # Format similar to streaming - show progress numbers
            if event.progress is not None:
                if event.total is not None:
                    progress_display = f"{int(event.progress)}/{int(event.total)}"
                else:
                    progress_display = str(int(event.progress))
            else:
                progress_display = "Processing"
            formatted_text = f"▎[dim]▶[/dim] {progress_display}".ljust(17 + 11)
            description = f"[{self._get_action_style(event.action)}]{formatted_text}"
        else:
            formatted_text = f"▎[dim]•[/dim] {event.action.value}".ljust(17 + 11)
            description = f"[{self._get_action_style(event.action)}]{formatted_text}"

        # Update basic task information
        update_kwargs: dict[str, Any] = {
            "description": description,
            "target": event.target or task_name,  # Use task_name as fallback for target
            "details": event.details or "",
            "task_name": task_name,
        }

        # For TOOL_PROGRESS events, update progress if available
        if event.action == ProgressAction.TOOL_PROGRESS and event.progress is not None:
            if event.total is not None:
                update_kwargs["completed"] = event.progress
                update_kwargs["total"] = event.total
            else:
                # If no total, reset to indeterminate but keep other fields
                self._progress.reset(task_id)
                # Still need to update after reset to apply the fields

        self._progress.update(task_id, **update_kwargs)

        if (
            event.action == ProgressAction.INITIALIZED
            or event.action == ProgressAction.READY
            or event.action == ProgressAction.LOADED
        ):
            self._progress.update(task_id, completed=100, total=100)
        elif event.action == ProgressAction.FINISHED:
            self._progress.update(
                task_id,
                completed=100,
                total=100,
                target=event.target or task_name,
                details=f" / Elapsed Time {time.strftime('%H:%M:%S', time.gmtime(self._progress.tasks[task_id].elapsed))}",
                task_name=task_name,
            )
        elif event.action == ProgressAction.FATAL_ERROR:
            self._progress.update(
                task_id,
                completed=100,
                total=100,
                target=event.target or task_name,
                details=f" / {event.details}",
                task_name=task_name,
            )
        else:
            self._progress.reset(task_id)
