"""Terminal UI for streaming agent activity in real-time."""

from __future__ import annotations

import json
import re
import time

from rich.console import Console
from rich.markdown import Markdown
from rich.padding import Padding
from rich.theme import Theme

_theme = Theme({
    "markdown.code": "bold cyan",
    "markdown.code_block": "dim",
})

console = Console(stderr=True, theme=_theme)

TOOL_STYLE: dict[str, str] = {
    "Read": "cyan",
    "Write": "green",
    "Edit": "green",
    "Bash": "yellow",
    "Grep": "magenta",
    "Glob": "magenta",
    "Agent": "blue",
    "Dispatch": "blue",
    "Task": "blue",
}

_AGENT_TOOLS = {"Agent", "Dispatch", "Task"}

DEFAULT_TOOL_STYLE = "dim"

TOOL_ARG_KEYS: dict[str, list[str]] = {
    "Read": ["file_path", "path"],
    "Write": ["file_path", "path"],
    "Edit": ["file_path", "path"],
    "Bash": ["command"],
    "Grep": ["pattern", "regex"],
    "Glob": ["glob_pattern", "pattern"],
    "Agent": ["prompt", "task", "message", "description"],
    "Dispatch": ["prompt", "task", "message", "description"],
    "Task": ["prompt", "task", "message", "description"],
}

_MD_PATTERN = re.compile(
    r"(^#{1,4}\s)|(\*\*.*\*\*)|(\n\d+\.\s)|(\n[-*]\s)|(`[^`]+`)",
    re.MULTILINE,
)


class PhaseUI:
    """Renders streaming output for a single agent phase.

    Each phase gets its own instance.  For parallel reviews, each
    reviewer gets an instance with a ``prefix`` like ``"[Linus] "``.
    """

    def __init__(
        self,
        *,
        verbose: bool = False,
        prefix: str = "",
    ) -> None:
        self._verbose = verbose
        self._prefix = prefix
        self._tool_name: str | None = None
        self._tool_json: str = ""
        self._tool_displayed: bool = False
        self._text_buf: str = ""
        self._in_tool: bool = False
        self._start_time: float = time.monotonic()
        self._turn_count: int = 0

    # -- phase-level markers ------------------------------------------------

    def phase_header(
        self,
        phase_name: str,
        budget: float,
        model: str,
        extra: str = "",
    ) -> None:
        label = f" Phase: {phase_name}  ${budget:.2f} budget · {model}"
        if extra:
            label += f" · {extra}"
        console.rule(label, style="bold")

    def phase_complete(
        self,
        cost: float,
        turns: int,
        duration_ms: int,
    ) -> None:
        self._flush_text()
        elapsed = _format_duration(duration_ms)
        console.print(
            f"\n  [green]✓[/green] Phase completed  "
            f"${cost:.2f} · {turns} turns · {elapsed}\n",
            highlight=False,
        )

    def phase_error(self, error: str) -> None:
        self._flush_text()
        console.print(
            f"\n  [red]✗[/red] Phase failed: {error}\n",
            highlight=False,
        )

    # -- streaming callbacks ------------------------------------------------

    def on_tool_start(self, tool_name: str) -> None:
        self._flush_text()
        self._in_tool = True
        self._tool_name = tool_name
        self._tool_json = ""
        self._tool_displayed = False

    def on_tool_input_delta(self, partial_json: str) -> None:
        self._tool_json += partial_json
        if not self._tool_displayed:
            arg = self._try_extract_arg()
            if arg is not None:
                self._print_tool_line(arg)
                self._tool_displayed = True

    def on_tool_done(self) -> None:
        if self._tool_name and not self._tool_displayed:
            arg = self._try_extract_arg() or ""
            self._print_tool_line(arg)
        self._tool_name = None
        self._tool_json = ""
        self._tool_displayed = False
        self._in_tool = False

    def on_text_delta(self, text: str) -> None:
        if self._in_tool:
            return
        self._text_buf += text

    def on_turn_complete(self) -> None:
        self._flush_text()
        self._turn_count += 1

    # -- internals ----------------------------------------------------------

    def _flush_text(self) -> None:
        """Render buffered agent text, then clear the buffer."""
        raw = self._text_buf.strip()
        self._text_buf = ""
        if not raw:
            return
        if _looks_like_markdown(raw):
            if self._prefix:
                console.print(f"\n  {self._prefix}─────", highlight=False)
            else:
                console.print()
            console.print(Padding(Markdown(raw), (0, 4)))
        else:
            for line in raw.splitlines():
                stripped = line.strip()
                if stripped:
                    console.print(
                        f"  {self._prefix}[dim]{_truncate(stripped, 120)}[/dim]",
                        highlight=False,
                    )

    def _print_tool_line(self, arg: str) -> None:
        name = self._tool_name or "?"
        style = TOOL_STYLE.get(name, DEFAULT_TOOL_STYLE)
        label = f"{name} {arg}".rstrip() if arg else name
        console.print(
            f"  {self._prefix}[{style}]●[/{style}] {label}",
            highlight=False,
        )

    def _try_extract_arg(self) -> str | None:
        if not self._tool_name:
            return None
        keys = TOOL_ARG_KEYS.get(self._tool_name)
        if not keys:
            return None
        try:
            data = json.loads(self._tool_json)
            for key in keys:
                val = data.get(key)
                if val:
                    text = str(val)
                    if self._tool_name in _AGENT_TOOLS:
                        text = _first_meaningful_line(text)
                    return _truncate(text, 80)
            return None
        except (json.JSONDecodeError, TypeError):
            return None


class NullUI:
    """Drop-in replacement that discards all output.  Used in tests and --quiet."""

    def phase_header(self, *a: object, **kw: object) -> None: ...
    def phase_complete(self, *a: object, **kw: object) -> None: ...
    def phase_error(self, *a: object, **kw: object) -> None: ...
    def on_tool_start(self, *a: object) -> None: ...
    def on_tool_input_delta(self, *a: object) -> None: ...
    def on_tool_done(self) -> None: ...
    def on_text_delta(self, *a: object) -> None: ...
    def on_turn_complete(self) -> None: ...


# -- reviewer tag helpers --------------------------------------------------

REVIEWER_COLORS = [
    "bright_cyan", "bright_magenta", "bright_yellow", "bright_green",
    "bright_blue", "bright_red", "bright_white",
]


def _reviewer_color(index: int) -> str:
    return REVIEWER_COLORS[index % len(REVIEWER_COLORS)]


def make_reviewer_prefix(index: int) -> str:
    """Build a short numbered prefix like '[cyan]R1[/cyan] '."""
    color = _reviewer_color(index)
    return f"[{color}]R{index + 1}[/{color}] "


def print_reviewer_legend(reviewers: list[tuple[int, str]]) -> None:
    """Print legend mapping R1..RN -> full role name before review starts."""
    console.print()
    for i, role in reviewers:
        color = _reviewer_color(i)
        console.print(
            f"  [{color}]R{i + 1}[/{color}] {role}",
            highlight=False,
        )
    console.print()



def _looks_like_markdown(text: str) -> bool:
    """Return True if text contains markdown formatting worth rendering."""
    return bool(_MD_PATTERN.search(text))


def _first_meaningful_line(text: str) -> str:
    """Extract the first non-blank, non-heading line from an agent prompt."""
    for line in text.splitlines():
        stripped = line.strip().lstrip("#").strip()
        if stripped and len(stripped) > 5:
            return stripped
    return text.split("\n", 1)[0].strip()


def _truncate(s: str, maxlen: int) -> str:
    return s if len(s) <= maxlen else s[: maxlen - 1] + "…"


def _format_duration(ms: int) -> str:
    secs = ms // 1000
    if secs < 60:
        return f"{secs}s"
    mins, secs = divmod(secs, 60)
    return f"{mins}m {secs}s"
