"""Terminal renderer for agent_ws_v3 streaming messages."""

from __future__ import annotations

import json
import random
import re
import threading
import time
from dataclasses import dataclass
from typing import Any, Optional

from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel
from rich.pretty import Pretty
from rich.status import Status
from rich.table import Table
from rich.text import Text


@dataclass
class AskPromptRequest:
    """Represents an ask_message that requires user input."""

    message_id: str
    questions: list[dict[str, Any]]


class StreamRenderer:
    """Render websocket messages with type-aware formatting."""

    def __init__(
        self, console: Optional[Console] = None, *, collapse_thinking: bool = False, debug: bool = False
    ) -> None:
        self.console = console or Console()
        self._collapse_thinking = collapse_thinking
        self._debug = debug
        self._stream_buffers: dict[str, str] = {}
        self._stream_labels: dict[str, str] = {}
        self._stream_started: set[str] = set()
        self._open_stream_ids: set[str] = set()
        self._active_stream_id: Optional[str] = None
        self._live_thinking_stream_id: Optional[str] = None
        self._live_thinking_line_len = 0
        self._completed_thinking: dict[str, dict[str, Any]] = {}
        self._thinking_status: Optional[Status] = None
        self._thinking_lock = threading.Lock()
        self._thinking_stop_event: Optional[threading.Event] = None
        self._thinking_animation_thread: Optional[threading.Thread] = None
        self._thinking_plain_text = "thinking..."
        self._thinking_sweep_index = 0
        self._thinking_sweep_width = 4
        self._agent_thinking_active = False
        self._pending_mcp_tools: dict[str, dict] = {}  # message_id -> {tool_name, sweep_index}
        self._loading_messages = [
            "Still working on your request...",
            "Gathering context and checking details...",
            "Crunching through the next steps...",
            "Running tools and validating results...",
            "Polishing the response for you...",
        ]
        self._loading_message = random.choice(self._loading_messages)
        self._loading_next_switch_at = time.monotonic() + 1.8

    def render_banner(
        self,
        *,
        endpoint: str,
        user_id: str,
        thread_id: str,
        hide_session_details: bool = False,
    ) -> None:
        # Intentionally hidden: always ignore the session banner block.
        _ = (endpoint, user_id, thread_id, hide_session_details)
        return

    def _strip_session_block(self, text: str) -> str:
        """Remove session banner-like blocks echoed in model output."""
        if not text:
            return text

        # Matches common variants that include the session heading and fields.
        block_pattern = re.compile(
            r"(?:^|\n)"
            r"(?:[^\n]*Session[^\n]*\n)?"
            r"(?:[^\n]*\n){0,3}"
            r"[^\n]*Agent Loop CLI[^\n]*\n"
            r"[^\n]*endpoint:[^\n]*\n"
            r"[^\n]*user_id:[^\n]*\n"
            r"[^\n]*thread_id:[^\n]*"
            r"(?:\n[^\n]*){0,3}",
            re.IGNORECASE,
        )
        cleaned = block_pattern.sub("\n", text)
        cleaned = re.sub(r"\n{3,}", "\n\n", cleaned).strip()
        return cleaned

    def render_help(self) -> None:
        table = Table(title="Commands", show_header=True, header_style="bold magenta")
        table.add_column("Command", style="cyan")
        table.add_column("Description")
        table.add_row("/help", "Show this help")
        table.add_row("/thread", "Show current thread id")
        table.add_row("/new", "Reconnect with a fresh session (no thread id)")
        table.add_row("/attach <path>", "Queue local file for next query upload")
        table.add_row("/attachments", "List queued local attachments")
        table.add_row("/detach <index|all>", "Remove queued attachments")
        table.add_row("/clear", "Clear terminal output")
        table.add_row("/interrupt", "Send interrupt message")
        table.add_row("/upgrade", "Upgrade CLI from configured wheel server")
        table.add_row("/thinking", "List cached thinking blocks")
        table.add_row("/thinking <message_id|prefix>", "Show one thinking block")
        table.add_row("/exit or /quit", "Terminate and exit")
        self.console.print(table)

    def render_thinking_history(self, message_id: Optional[str] = None) -> None:
        if not self._completed_thinking:
            self.console.print("[dim]no thinking blocks cached[/dim]")
            return

        if not message_id:
            self.console.print("[bold magenta]Thinking blocks[/bold magenta]")
            for key, item in self._completed_thinking.items():
                thinking_time = item.get("thinking_time")
                duration = (
                    f"{thinking_time:.1f}s"
                    if isinstance(thinking_time, (int, float))
                    else "unknown"
                )
                preview = str(item.get("text") or "").strip().replace("\n", " ")
                if len(preview) > 70:
                    preview = f"{preview[:67]}..."
                self.console.print(f"  {key} ({duration}) {preview}")
            return

        selected_id = message_id
        if selected_id not in self._completed_thinking:
            matches = [k for k in self._completed_thinking if k.startswith(selected_id)]
            if len(matches) == 1:
                selected_id = matches[0]
            elif len(matches) > 1:
                self.console.print(
                    f"[yellow]multiple matches for prefix '{message_id}'[/yellow]"
                )
                for key in matches:
                    self.console.print(f"  {key}")
                return
            else:
                self.console.print(f"[red]thinking block not found:[/red] {message_id}")
                return

        item = self._completed_thinking[selected_id]
        thinking_time = item.get("thinking_time")
        if isinstance(thinking_time, (int, float)):
            title = f"Thinking ({thinking_time:.1f}s) {selected_id}"
        else:
            title = f"Thinking {selected_id}"
        self.console.print(
            Panel(
                self._format_thinking_panel_body(str(item.get("text") or "")),
                title=title,
                border_style="magenta",
                width=self._thinking_panel_width(),
                expand=False,
            )
        )

    def render_message(self, msg: dict[str, Any]) -> Optional[AskPromptRequest]:
        msg_type = msg.get("type", "_unknown")
        if msg_type == "complete":
            self._stop_thinking_status()
            credits = msg.get("credits")
            credits_suffix = f" credits={credits}" if credits is not None else ""
            self.console.print(f"[bold green]✓ complete[/bold green]{credits_suffix}")
            return None

        if msg_type != "metadata" and not self._open_stream_ids:
            self._ensure_thinking_status()
        if msg_type not in {"agent_thinking", "metadata"}:
            self._agent_thinking_active = False

        if msg_type == "metadata":
            return None

        if msg_type in {
            "thinking_commentary",
            "commentary",
            "child_step_commentary",
            "inter_child_text",
        }:
            self._render_streaming_commentary(msg)
            return None

        if msg_type == "agent_message":
            content = msg.get("agent_message") or msg.get("message") or ""
            if isinstance(content, str):
                content = self._strip_session_block(content)
                if content.strip():
                    self.console.print(Markdown(content))
            else:
                self.console.print(Pretty(content))
            return None

        if msg_type in {"tool_execution", "tool_progress", "tool_completed"}:
            self._render_tool_event(msg)
            return None

        if msg_type == "mcp_message":
            self._handle_mcp_message(msg)
            return None

        if msg_type == "skill_message":
            self._handle_skill_message(msg)
            return None

        if msg_type == "agent_thinking":
            step_index = msg.get("step_index")
            thinking_text = str(msg.get("message") or "thinking...")
            step_prefix = f"step {step_index} " if step_index is not None else ""
            self._thinking_plain_text = f"{step_prefix}{thinking_text}"
            self._thinking_sweep_index = 0
            self._agent_thinking_active = True
            if self._thinking_status is not None:
                self._thinking_status.update(self._format_thinking_status_text())
            return None

        if msg_type in {"step_discovered", "step_started", "step_progress", "step_completed"}:
            step_index = msg.get("step_index", "?")
            title = msg.get("title") or msg.get("description") or ""
            self.console.print(f"[cyan]step {step_index}[/cyan] {msg_type}: {title}")
            return None

        if msg_type in {"interrupted", "terminated"}:
            self.console.print(
                f"[yellow]{msg_type}:[/yellow] {msg.get('message', '')}"
            )
            return None

        if msg_type == "recommended_questions":
            questions = msg.get("questions") or []
            if questions:
                self.console.print("[bold cyan]Recommended questions[/bold cyan]")
                for i, q in enumerate(questions, 1):
                    self.console.print(f"  {i}. {q}")
            return None

        if msg_type == "file_monitor":
            self._render_file_monitor(msg)
            return None

        if msg_type == "error_update":
            err = msg.get("error_message") or msg.get("ui_message") or "Unknown error"
            code = msg.get("error_code")
            self.console.print(
                Panel(
                    f"[bold red]Error[/bold red]\ncode={code}\n{err}",
                    border_style="red",
                )
            )
            return None

        if msg_type == "resume_complete":
            self.console.print(
                f"[green]resume complete[/green] lines={msg.get('line_count', '?')}"
            )
            return None

        if msg_type == "ask_message":
            status = str(msg.get("status") or "")
            if status == "called":
                # Ensure the prompt is visible and not overwritten by live status updates.
                self._stop_thinking_status()
                return AskPromptRequest(
                    message_id=str(msg.get("message_id", "")),
                    questions=list(msg.get("questions", [])),
                )
            self._render_ask_message_status(msg)
            return None

        # Fallback for unknown/new message types.
        compact = json.dumps(msg, ensure_ascii=False)
        self.console.print(
            Panel(compact, title=f"message:{msg_type}", border_style="white")
        )
        return None

    def _render_file_monitor(self, msg: dict[str, Any]) -> None:
        file_path = str(msg.get("file_path") or "(unknown)")
        file_type = str(msg.get("file_type") or "file")
        metadata = msg.get("metadata")
        source = ""
        reason = ""
        if isinstance(metadata, dict):
            source = str(metadata.get("source") or "")
            reason = str(metadata.get("reason") or "")

        lines = [f"[blue]file monitor[/blue] ({file_type}) [cyan]{file_path}[/cyan]"]
        if source:
            lines.append(f"[dim]source:[/dim] {source}")
        if reason:
            lines.append(f"[dim]reason:[/dim] {reason}")

        content = msg.get("content")
        if isinstance(content, str) and content.strip():
            lines.append("")
            lines.append(content.rstrip())

        self.console.print("\n".join(lines))

    def _render_streaming_commentary(self, msg: dict[str, Any]) -> None:
        msg_type = msg.get("type", "commentary")
        phase = msg.get("phase")
        message_id = str(msg.get("message_id") or "")
        label = msg_type.replace("_", " ")
        if message_id:
            self._stream_labels[message_id] = label

        if phase == "start":
            if message_id:
                self._stream_buffers[message_id] = ""
                self._stream_started.discard(message_id)
                self._open_stream_ids.add(message_id)
                self._active_stream_id = message_id
                if msg_type == "thinking_commentary":
                    self._live_thinking_stream_id = message_id
                    self._live_thinking_line_len = 0
            # Rich live status redraw can overwrite inline stream deltas.
            # Pause it while commentary streams are actively rendering.
            self._stop_thinking_status()
            return

        if phase == "delta":
            delta = str(msg.get("delta") or "")
            if self._thinking_status is not None:
                self._stop_thinking_status()
            if not message_id:
                if delta:
                    self.console.print(delta, end="")
                return

            if self._active_stream_id and self._active_stream_id != message_id:
                self._finish_live_thinking_line()
                self.console.print("")
            self._active_stream_id = message_id
            self._stream_buffers[message_id] = self._stream_buffers.get(message_id, "") + delta
            if message_id not in self._stream_started:
                if msg_type != "thinking_commentary":
                    self.console.print(f"[dim]{label}...[/dim]")
                self._stream_started.add(message_id)
            if delta:
                if msg_type == "thinking_commentary":
                    self._render_live_thinking_line(self._stream_buffers[message_id])
                else:
                    self.console.print(delta, end="")
            return

        if phase == "complete":
            final_text = msg.get("final_text")
            thinking_time = msg.get("thinking_time")
            if msg_type == "thinking_commentary":
                # Use streamed deltas by default; if final_text is provided, it replaces
                # the buffered delta text for the final rendered thinking block.
                text = self._stream_buffers.get(message_id, "")
                if isinstance(final_text, str) and final_text.strip():
                    text = final_text
                self._finish_live_thinking_line()
                if message_id:
                    self._completed_thinking[message_id] = {
                        "text": text,
                        "thinking_time": thinking_time,
                    }
                if text.strip():
                    self._render_thinking_commentary_line(text, thinking_time)
            else:
                text = (
                    final_text
                    if isinstance(final_text, str) and final_text.strip()
                    else self._stream_buffers.get(message_id, "")
                )
                text = self._strip_session_block(text)
                self.console.print("")  # newline after delta stream
                if text.strip():
                    self.console.print(Panel(Markdown(text), title=label, border_style="cyan"))
            self._stream_buffers.pop(message_id, None)
            self._stream_labels.pop(message_id, None)
            self._stream_started.discard(message_id)
            self._open_stream_ids.discard(message_id)
            if self._active_stream_id == message_id:
                self._active_stream_id = None
            if self._live_thinking_stream_id == message_id:
                self._live_thinking_stream_id = None
            return

        # Unknown phase or non-streaming payload
        fallback_text = str(msg.get("message") or msg.get("text") or "")
        fallback_text = self._strip_session_block(fallback_text)
        if fallback_text:
            self.console.print(f"[dim]{label}:[/dim] {fallback_text}")

    def _thinking_panel_width(self) -> int:
        # Keep thinking blocks in a stable, readable area.
        return max(56, min(self.console.width - 2, 86))

    def _format_thinking_panel_body(self, text: str) -> str:
        # Render as single-row viewport:
        # - thinking
        #   <one line clipped to fixed width>
        inner_width = max(24, self._thinking_panel_width() - 8)
        single_line = " ".join(text.split())
        if len(single_line) > inner_width:
            # Keep a fixed-width viewport on the newest content.
            single_line = f"…{single_line[-(inner_width - 1):]}"
        return f"- thinking\n   {single_line}"

    def _format_duration(self, seconds: Any) -> str:
        if not isinstance(seconds, (int, float)):
            return "unknown"
        if seconds < 60:
            return f"{seconds:.1f}s"
        mins = int(seconds // 60)
        rem = int(seconds % 60)
        return f"{mins}m {rem:02d}s"

    def _single_row_thinking_text(self, text: str) -> str:
        viewport = max(28, min(self.console.width - 22, 92))
        single_line = " ".join(text.split())
        if len(single_line) > viewport:
            single_line = f"…{single_line[-(viewport - 1):]}"
        return single_line

    def _render_thinking_commentary_line(self, text: str, thinking_time: Any) -> None:
        duration = self._format_duration(thinking_time)
        self.console.print(f"[magenta]🧠 Thinking for {duration}[/magenta]")
        self.console.print(f"[dim]   {self._single_row_thinking_text(text)}[/dim]")

    def _render_live_thinking_line(self, text: str) -> None:
        if self._debug:
            return
        line = f"🧠 {self._single_row_thinking_text(text)}"
        pad = ""
        if self._live_thinking_line_len > len(line):
            pad = " " * (self._live_thinking_line_len - len(line))
        self.console.file.write(f"\r{line}{pad}")
        self.console.file.flush()
        self._live_thinking_line_len = len(line)

    def _finish_live_thinking_line(self) -> None:
        if self._live_thinking_line_len > 0:
            self.console.file.write("\n")
            self.console.file.flush()
            self._live_thinking_line_len = 0

    def _render_ask_message_status(self, msg: dict[str, Any]) -> None:
        status = str(msg.get("status") or "unknown")
        message_id = str(msg.get("message_id") or "")
        answer = msg.get("answer")
        error = msg.get("error")

        if status == "answered":
            title = "[green]✓ answer submitted[/green]"
            suffix = f" id={message_id}" if message_id else ""
            self.console.print(f"{title}{suffix}")
            if isinstance(answer, str) and answer.strip():
                self.console.print(f"[dim]   {self._single_row_thinking_text(answer)}[/dim]")
            return

        if status == "failed":
            detail = str(error or "ask_message failed")
            self.console.print(f"[red]✗ ask_message failed[/red] id={message_id}")
            self.console.print(f"[dim]   {self._single_row_thinking_text(detail)}[/dim]")
            return

        # Keep non-called statuses readable without dumping raw JSON blobs.
        suffix = f" id={message_id}" if message_id else ""
        self.console.print(f"[dim]ask_message status={status}[/dim]{suffix}")

    def _handle_mcp_message(self, msg: dict[str, Any]) -> None:
        message_id = str(msg.get("message_id") or "")
        tool_name = str(msg.get("tool_name") or "unknown")
        status = msg.get("status")

        if status == "called":
            with self._thinking_lock:
                self._pending_mcp_tools[message_id] = {
                    "tool_name": tool_name,
                    "sweep_index": 0,
                }
            return

        if status in ("completed", "failed"):
            with self._thinking_lock:
                self._pending_mcp_tools.pop(message_id, None)

            error = msg.get("error")
            result = msg.get("result")
            content = ""
            if error:
                content = str(error)
            elif result is not None:
                if isinstance(result, dict):
                    output = result.get("output", "")
                    content = str(output) if output else json.dumps(result, ensure_ascii=False)
                else:
                    content = str(result)
            self._render_mcp_completed_line(tool_name, status, content)

    def _handle_skill_message(self, msg: dict[str, Any]) -> None:
        message_id = str(msg.get("message_id") or "")
        skill_name = str(msg.get("skill") or "unknown")
        status = msg.get("status")

        if status == "called":
            with self._thinking_lock:
                self._pending_mcp_tools[message_id] = {
                    "tool_name": f"Skill({skill_name})",
                    "sweep_index": 0,
                }
            return

        if status in ("completed", "failed"):
            with self._thinking_lock:
                self._pending_mcp_tools.pop(message_id, None)
            metadata = msg.get("metadata") or {}
            content = str(metadata.get("args") or "")
            error = msg.get("error")
            if error:
                content = str(error)
            self._render_mcp_completed_line(f"Skill({skill_name})", status, content)

    def _render_mcp_completed_line(self, tool_name: str, status: str, content: str) -> None:
        if status == "failed":
            self.console.print(f"[red]✗ {tool_name}[/red]")
        else:
            self.console.print(f"[blue]🔧 {tool_name}[/blue]")
        if content.strip():
            self.console.print(f"[dim]   {self._single_row_thinking_text(content)}[/dim]")

    def _render_tool_event(self, msg: dict[str, Any]) -> None:
        msg_type = str(msg.get("type"))
        tool_name = str(msg.get("tool_name") or msg.get("tool") or msg.get("name") or "unknown_tool")
        step_index = msg.get("step_index")
        child_index = msg.get("child_index")
        tool_key = f"tool_{step_index}_{child_index}" if step_index is not None and child_index is not None else f"tool_{tool_name}"

        if msg_type == "tool_execution":
            with self._thinking_lock:
                self._pending_mcp_tools[tool_key] = {
                    "tool_name": tool_name,
                    "sweep_index": 0,
                }
            return

        if msg_type == "tool_progress":
            return

        if msg_type == "tool_completed":
            with self._thinking_lock:
                self._pending_mcp_tools.pop(tool_key, None)
            result = msg.get("result")
            content = ""
            if isinstance(result, dict):
                output = result.get("output") or result.get("message") or result.get("summary") or ""
                content = str(output) if output else json.dumps(result, ensure_ascii=False)
            elif result is not None:
                content = str(result)
            self._render_mcp_completed_line(tool_name, "completed", content)

    def _format_thinking_status_text(self) -> Text:
        output = Text()
        if self._agent_thinking_active:
            message = self._thinking_plain_text or "thinking..."
            line = Text(f"✨ {message}", style="bold magenta")
            line_length = len(line.plain)
            if line_length > 2:
                sweep_start = min(2 + self._thinking_sweep_index, line_length)
                sweep_end = min(sweep_start + self._thinking_sweep_width, line_length)
                if sweep_start < sweep_end:
                    line.stylize("bold white", sweep_start, sweep_end)
            output.append(line)
            output.append("\n")
        for tool_info in list(self._pending_mcp_tools.values()):
            tool_name = tool_info["tool_name"]
            sweep_idx = tool_info["sweep_index"]
            display_text = f"🔧 {tool_name}"
            line = Text(display_text, style="bold blue")
            line_length = len(line.plain)
            if line_length > 2:
                sweep_start = min(2 + sweep_idx, line_length)
                sweep_end = min(sweep_start + self._thinking_sweep_width, line_length)
                if sweep_start < sweep_end:
                    line.stylize("bold white", sweep_start, sweep_end)
            output.append(line)
            output.append("\n")
        output.append(Text("────────────────────────────", style="dim"))
        output.append("\n")
        output.append(Text(f"⏳ {self._loading_message}", style="dim cyan"))
        return output

    def _start_thinking_animation_loop(self) -> None:
        stop_event = threading.Event()
        self._thinking_stop_event = stop_event

        def _animate_loop() -> None:
            while not stop_event.wait(0.12):
                with self._thinking_lock:
                    if self._thinking_status is None:
                        break
                    now = time.monotonic()
                    if now >= self._loading_next_switch_at:
                        self._loading_message = random.choice(self._loading_messages)
                        self._loading_next_switch_at = now + 1.8

                    if self._agent_thinking_active:
                        text_length = len(self._thinking_plain_text or "thinking...")
                        if text_length > 0:
                            self._thinking_sweep_index = (self._thinking_sweep_index + 1) % text_length
                    for tool_info in self._pending_mcp_tools.values():
                        text_length = max(1, len(tool_info["tool_name"]) + 4)  # "🔧 " prefix
                        tool_info["sweep_index"] = (tool_info["sweep_index"] + 1) % text_length
                    status_text = self._format_thinking_status_text()
                    self._thinking_status.update(status_text)

        self._thinking_animation_thread = threading.Thread(
            target=_animate_loop,
            name="agent-thinking-sweep",
            daemon=True,
        )
        self._thinking_animation_thread.start()

    def _stop_thinking_status(self) -> None:
        if self._thinking_stop_event is not None:
            self._thinking_stop_event.set()
            self._thinking_stop_event = None
        if self._thinking_status is not None:
            self._thinking_status.stop()
            self._thinking_status = None
        self._agent_thinking_active = False
        self._thinking_animation_thread = None

    def _ensure_thinking_status(self) -> None:
        if self._thinking_status is not None:
            return
        self._loading_next_switch_at = time.monotonic() + 1.8
        status = self.console.status(self._format_thinking_status_text(), spinner="dots")
        status.start()
        self._thinking_status = status
        self._start_thinking_animation_loop()

