"""Chat panel widget for the TUI."""

from __future__ import annotations

import json
import time
from typing import Any

from rich.console import RenderableType
from rich.segment import Segment
from rich.style import Style as RichStyle
from rich.text import Text
from textual.app import ComposeResult
from textual.containers import ScrollableContainer
from textual.content import Content
from textual.geometry import Offset
from textual.selection import Selection
from textual.style import Style
from textual.widgets import Static

from clients.terminal.tui.chat_model import ChatItem, ChatModel
from clients.terminal.tui.renderers import default_registry
from clients.terminal.tui.settings import get_tui_settings
from clients.terminal.tui.themes import get_active_theme

# Minimum gap between expensive rich-renderable rebuilds of a streaming
# assistant bubble. Assistant answers render as rich ``Markdown``; re-parsing
# the growing text on every ``stream_delta`` is O(N²) over the answer length.
# Coalesce rebuilds to at most ~one per window (the ChatItem content still
# updates per delta; only the Markdown parse + rich render is throttled). Plain
# renderables (thinking/planning are ``Text``) bypass this — they are cheap.
_STREAMING_RENDER_THROTTLE = 0.15


def _item_msg(item: ChatItem) -> dict[str, Any]:
    """Map a ChatItem to the msg dict its renderer expects."""
    kind = item.kind
    if kind == "user_message":
        return {"type": "user_message", "content": item.content}
    if kind == "assistant_message":
        return {"type": "assistant_message", "content": item.content}
    if kind == "thinking_block":
        return {
            "type": "thinking_block",
            "content": item.content,
            "is_subagent": item.meta.get("is_subagent", False),
        }
    if kind == "tool_started":
        return {"type": "tool_started", **item.meta}
    if kind == "tool_call":
        return {"type": "tool_call", **item.meta}
    if kind == "error":
        return {"type": "error", "message": item.content}
    if kind == "status":
        return {"type": "status", "content": item.content}
    if kind == "planning_status":
        return {
            "type": "planning_status",
            "label": item.content,
            "phase": item.meta.get("phase"),
            "is_subagent": item.meta.get("is_subagent", False),
        }
    if kind == "plan_ready":
        return {
            "type": "plan_ready",
            "plan": item.content,
            "is_subagent": item.meta.get("is_subagent", False),
        }
    if kind == "turn_meta":
        return {"type": "turn_meta", "elapsed_seconds": item.meta.get("elapsed_seconds")}
    if kind == "context_summary":
        return {
            "type": "context_summary",
            "content": item.content,
            "messages_before": item.meta.get("messages_before"),
            "messages_after": item.meta.get("messages_after"),
        }
    if kind == "recall":
        return {"type": "recall_update", **item.meta}
    return {"type": "plain", "content": item.content}


def _signature(msg: dict[str, Any]) -> str:
    """Stable serialized form of a renderer msg, for cache invalidation.

    Only an item whose rendered inputs changed should re-render; this captures
    content + meta so in-place mutations (e.g. a streaming assistant bubble
    growing token by token, or a rolling planning_status line) invalidate
    exactly their own widget.
    """
    return json.dumps(msg, sort_keys=True, default=str)


class _ChatItemView(Static):
    """One chat item as its own widget.

    Making each message a distinct widget is what keeps streaming cheap: a
    per-token ``stream_delta`` only re-renders *this* widget, not the whole
    conversation — Textual refreshes just the changed region instead of
    re-laying-out every visible item.

    The item's rich renderable (Panel/Markdown/Group/…) is rendered to a
    Textual ``Content`` in ``render()`` rather than handed to ``Static`` as a
    raw rich renderable. Only ``Text``/``str``/``Content`` visuals are
    selection-aware in Textual; a raw rich ``Panel`` wraps into a ``RichVisual``
    which neither shows the drag-select highlight nor returns text from
    ``get_selection`` — so a drag over a message looked like "nothing
    selected". Building a ``Content`` makes the widget selection-aware (the
    magenta highlight renders) while ``get_selection`` strips the bubble
    chrome so the copy is clean.
    """

    DEFAULT_CSS = """
    _ChatItemView {
        height: auto;
        width: 1fr;
        padding: 0;
    }
    """

    def __init__(self, renderable: RenderableType, *, item: ChatItem, signature: str) -> None:
        super().__init__("")
        self._item: ChatItem = item
        self._signature: str = signature
        # The raw rich renderable (Panel/Text/Group/...) this widget shows. Kept
        # alongside the rendered Content so render() can rebuild on resize and
        # get_selection can strip the bubble chrome from the rendered text.
        self._rich_renderable: RenderableType = renderable
        # (id(renderable), width) -> Content cache. Rebuilding a Content renders
        # the rich renderable to segments, so we cache by content identity +
        # width: a repeated paint at the same width/ content reuses the Content,
        # and a resize or token update (new renderable) invalidates it.
        self._content_cache: tuple[tuple[int, int], Content] | None = None
        # (id(renderable), width) -> height cache for get_content_height, so the
        # layout pass does not re-render the rich renderable to measure it when
        # neither the renderable nor the width changed (e.g. between throttled
        # rebuilds, or repeated paints). Invalidated alongside _content_cache.
        self._height_cache: tuple[tuple[int, int], int] | None = None
        # Throttle state for streaming assistant rebuilds (see
        # _STREAMING_RENDER_THROTTLE). ``_last_render_at`` starts at 0 so the
        # first mutating delta after mount rebuilds immediately (snappy start),
        # then rapid deltas coalesce into ~one rebuild per window.
        self._render_timer = None
        self._render_pending: bool = False
        self._last_render_at: float = 0.0
        # Captured at maybe_update time for the deferred rebuild (the registry
        # and reading-mode flag are owned by ChatPanel).
        self._registry_ref = None
        self._reading: bool = False

    def maybe_update(self, item: ChatItem, registry, reading: bool = False) -> bool:
        """Re-render only if the item's signature changed. Returns True if updated.

        Assistant answers are rich Markdown — rebuilding the renderable on every
        stream delta is O(N²) over the answer. Throttle those rebuilds to ~one
        per ``_STREAMING_RENDER_THROTTLE`` window (coalesced); the ChatItem
        content still updates per delta, only the Markdown parse/render is
        deferred. Plain renderables (thinking/planning are Text) are cheap and
        rebuild immediately.
        """
        msg = _item_msg(item)
        sig = _signature(msg)
        if sig == self._signature:
            return False
        self._item = item
        self._signature = sig
        self._registry_ref = registry
        self._reading = reading
        if item.kind == "assistant_message":
            self._schedule_rebuild()
        else:
            self._do_rebuild()
        return True

    def _schedule_rebuild(self) -> None:
        """Rebuild the rich renderable now, or coalesce into the pending one.

        A rapid burst of stream deltas should produce at most one rebuild per
        throttle window, not one per token. If a rebuild is already scheduled,
        it will pick up the freshest ``_item`` when it fires — so just return.
        """
        if self._render_pending:
            return
        elapsed = time.monotonic() - self._last_render_at
        delay = _STREAMING_RENDER_THROTTLE - elapsed
        if delay <= 0:
            self._do_rebuild()
        else:
            self._render_pending = True
            self._render_timer = self.set_timer(delay, self._do_rebuild)

    def _do_rebuild(self) -> None:
        """Build a fresh rich renderable for the current item and refresh."""
        self._render_pending = False
        self._render_timer = None
        self._last_render_at = time.monotonic()
        registry = self._registry_ref
        if registry is None:
            return
        msg = _item_msg(self._item)
        # Always update the renderable (and invalidate the caches keyed on it):
        # if the widget is mid-mount, the fresh renderable is what gets painted
        # when mount completes. ``refresh`` only means something once mounted.
        self._rich_renderable = (
            registry.render_plain(msg) if self._reading else registry.render(msg)
        )
        self._content_cache = None
        self._height_cache = None
        if not self.is_mounted:
            return
        self.refresh(layout=True)
        # A deferred rebuild grows the widget's height AFTER ChatPanel._sync
        # already ran scroll_end against the old height. If the panel was
        # sticking to the bottom, re-anchor so the freshly rendered streaming
        # text stays in view instead of scrolling past the bottom.
        parent = self.parent
        if parent is not None and getattr(parent, "_stick_to_bottom", False):
            self.call_after_refresh(parent.scroll_end, animate=False)

    def flush(self) -> None:
        """Force a pending throttled rebuild to happen now.

        Called by ChatPanel on ``stream_end`` so the final chunk of streamed
        text appears immediately rather than after the throttle window.
        """
        if not self._render_pending:
            return
        if self._render_timer is not None:
            self._render_timer.stop()
        self._do_rebuild()

    # ── rendering + selection ────────────────────────────────────────────────

    def _build_content(self) -> Content:
        """Render the item's rich renderable to a selection-aware ``Content``,
        the same way Textual would draw it (same console, same width → identical
        line layout). Cached by (renderable identity, width) so repeated paints
        are cheap and resize/token-update invalidate it.
        """
        renderable = self._rich_renderable
        width = self.size.width
        if renderable is None or width <= 0:
            return Content("")
        key = (id(renderable), width)
        if self._content_cache is not None and self._content_cache[0] == key:
            return self._content_cache[1]
        console = self.app.console
        # Textual's selection-aware Content stores spans with Textual ``Style``
        # objects (the magenta selection_style is merged in via Style + Style at
        # paint time); rich segments carry rich ``Style`` objects, which crash
        # that merge with "unsupported operand type(s) for +". Convert each
        # segment's rich Style to a Textual Style (mirroring Content.from_rich_text)
        # so the selection highlight composes cleanly over our panel borders.
        ansi_theme = self.app.ansi_theme
        to_style = Style.from_rich_style
        options = console.options.update(width=width, highlight=False)
        segments = list(console.render(renderable, options))
        lines = list(Segment.split_lines(segments))
        parts: list[tuple[str, Any]] = []
        last = len(lines) - 1
        for i, line in enumerate(lines):
            for seg in line:
                if seg.text:
                    # rich segments may carry ``style=None`` (e.g. a bare
                    # Segment with no style); Style.from_rich_style would crash
                    # on None, so fall back to rich's empty default Style.
                    parts.append(
                        (seg.text, to_style(seg.style or RichStyle(), ansi_theme))
                    )
            if i < last:
                parts.append(("\n", None))
        content = Content.assemble(*parts, end="")
        self._content_cache = (key, content)
        return content

    def render(self):  # type: ignore[override]
        try:
            return self._build_content()
        except Exception:
            return Content("")

    def get_content_height(self, container, viewport, width: int) -> int:  # type: ignore[override]
        """Auto-height of the bubble, computed directly from the rich renderable.

        Textual measures ``height: auto`` widgets by calling ``_render()`` (which
        wraps ``render()``) and asking the resulting Visual for its height. Our
        ``render()`` builds a selection-aware ``Content`` keyed by width, so every
        width change yields a *new* Content/Visual object — and since Textual's
        ``_render`` visual cache is not width-keyed, that spreads the width-settling
        cascade (0 → panel-width → final) across multiple refreshes. A deferred
        ``scroll_end`` (``call_after_refresh``) then lands at an *intermediate*
        ``max_scroll_y`` instead of the final one, so a freshly-filled chat no
        longer sticks to the bottom.

        Computing the height straight from the rich renderable (the same thing a
        raw ``RichVisual`` does) is a pure function of ``width`` with no Content
        object churn, so Textual reaches the layout fixpoint in one refresh and
        the deferred scroll lands at the true bottom. ``render()`` still supplies
        the selection-aware Content for painting.
        """
        renderable = self._rich_renderable
        if not width or renderable is None:
            return 0
        # Height is a pure function of (renderable, width); cache it so the
        # layout pass does not re-render the rich renderable to measure it when
        # neither changed (between throttled rebuilds, or repeated paints). The
        # cache key is content identity + width, mirroring _content_cache.
        key = (id(renderable), width)
        cache = self._height_cache
        if cache is not None and cache[0] == key:
            return cache[1]
        try:
            console = self.app.console
            options = console.options.update(width=width, highlight=False)
            segments = list(console.render(renderable, options))
            height = sum(1 for _ in Segment.split_lines(segments))
        except Exception:
            return 0
        self._height_cache = (key, height)
        return height

    @staticmethod
    def _denoise(lines: list[str]) -> tuple[list[str], int, bool]:
        """Split rendered lines into clean content + the column content starts at.

        Chat items render rich ``Panel``s (rounded borders + "Navi"/"You"
        titles). The outer panel adds one top and one bottom border line and a
        ``│ `` / `` │`` gutter on every content line; stripping those gives the
        selectable text the user actually wants to copy. Non-panel renderables
        (plain ``Text``/``Group`` — status lines, filesystem tool output,
        turn_meta) have no chrome: their lines are kept verbatim with no column
        shift, so the selection maps 1:1 onto what is on screen.

        Returns ``(content_lines, content_col, is_panel)``: ``content_lines``
        is the clean text (one entry per content row, trailing panel padding
        stripped), ``content_col`` is how far in a rendered line the content
        begins (0 for plain, pad+2 for a panel), ``is_panel`` flags the layout.
        """
        n = len(lines)

        def is_border(line: str) -> bool:
            s = line.lstrip(" ")
            return bool(s) and s[0] in "╭╮╰╯┌┐└┘"

        if n >= 3 and is_border(lines[0]) and is_border(lines[-1]):
            spans: list[tuple[int, str]] | None = []
            for line in lines[1:-1]:
                left = line.find("│")
                if left < 0:
                    spans = None
                    break
                right = line.rfind("│")
                content_col = left + 2
                end = right if right > left else len(line)
                spans.append((content_col, line[content_col:end].rstrip()))
            if spans and all(s[0] == spans[0][0] for s in spans):
                return [s[1] for s in spans], spans[0][0], True
        return list(lines), 0, False

    def get_selection(self, selection: Selection) -> tuple[str, str] | None:
        """Clean, chrome-free text under the selection.

        The widget renders a selection-aware ``Content`` (so the magenta
        highlight draws), but that Content carries the panel borders verbatim;
        ``get_selection`` would otherwise copy the borders too. We take the
        Content's plain text (line layout identical to the screen), strip the
        outer panel borders and gutters, and remap the mouse offsets onto the
        clean content so a partial-line selection still lands on the right
        characters. ``Ctrl+C`` then copies the result via OSC 52.
        """
        try:
            content = self._build_content()
        except Exception:
            return None
        lines = content.plain.split("\n")
        if not lines or lines == [""]:
            return None
        content_lines, content_col, is_panel = self._denoise(lines)
        if not content_lines:
            return None
        n_render = len(lines)
        n_content = len(content_lines)

        def map_point(offset: Offset | None) -> Offset | None:
            if offset is None:
                return None
            y, x = offset.y, offset.x
            if is_panel:
                if y <= 0:
                    return Offset(x=0, y=0)
                if y >= n_render - 1:
                    return Offset(x=len(content_lines[-1]), y=n_content - 1)
                cy = y - 1
            else:
                cy = y
            cy = max(0, min(cy, n_content - 1))
            cx = max(0, x - content_col)
            cx = min(cx, len(content_lines[cy]))
            return Offset(x=cx, y=cy)

        try:
            mapped = Selection(map_point(selection.start), map_point(selection.end))
            text = mapped.extract("\n".join(content_lines))
        except Exception:
            return None
        return text, "\n"


class ChatPanel(ScrollableContainer):
    """Scrollable conversation panel.

    Each visible message is a separate ``_ChatItemView`` child widget rather
    than one big ``Static`` rebuilt on every event. On a per-token stream the
    only widget that re-renders is the streaming bubble; everything else is
    untouched. Items are mounted/unmounted incrementally as the model grows and
    as the visible-window cap truncates the oldest.
    """

    DEFAULT_CSS = """
    ChatPanel {
        border: solid $tui-border;
        background: $tui-surface;
        color: $tui-text;
        padding: 0 1;
        height: 1fr;
        width: 2fr;
    }
    ChatPanel Static {
        color: $tui-text;
    }
    """

    def __init__(self) -> None:
        super().__init__()
        self._model = ChatModel()
        self._registry = default_registry()
        # Reading mode: render items without their Panel bubble chrome so a raw
        # terminal Shift+drag + Ctrl+Shift+C copies clean text in terminals that
        # lack OSC 52 (e.g. GNOME Console). Toggled by the app; when it changes
        # every visible widget is re-rendered in place.
        self._reading_mode: bool = False
        # Truncation hint shown at the top when the history exceeds the cap.
        self._hint_view = Static("", id="chat-truncation-hint")
        self._hint_view.styles.display = "none"
        # Number of older items currently truncated out of view (0 = none).
        self._hidden_count = 0
        # item -> view widget, keyed by id(item). Only holds visible items, so
        # it stays bounded by the window cap, not by total history length.
        self._widgets: dict[int, _ChatItemView] = {}
        # Captured in _sync (before any DOM change) so a deferred throttled
        # rebuild (see _ChatItemView._do_rebuild) knows whether to re-anchor the
        # scroll to the bottom after the streaming bubble grows.
        self._stick_to_bottom: bool = False

    def _render_for(self, item: ChatItem):
        """Renderable for an item, honouring the current reading mode.

        Reading mode uses the registry's borderless ``render_plain`` (raw text /
        unwrapped bodies) instead of the Panel-based ``render``, so the bubble
        borders never reach the screen and a raw terminal selection copies clean.
        """
        msg = _item_msg(item)
        if self._reading_mode:
            return self._registry.render_plain(msg)
        return self._registry.render(msg)

    def set_reading_mode(self, reading: bool) -> None:
        """Switch between bordered (normal) and borderless (reading) rendering.

        Re-renders every visible item widget in place — no remount — so scroll
        position is preserved. Items mounted later (new stream events) pick up
        the mode via :meth:`_render_for`.
        """
        if reading == self._reading_mode:
            return
        self._reading_mode = reading
        for widget in self._widgets.values():
            widget._rich_renderable = self._render_for(widget._item)
            widget._content_cache = None
            widget.refresh(layout=True)

    def compose(self) -> ComposeResult:
        yield self._hint_view

    def add_user_message(self, text: str) -> None:
        self._model.add_user_message(text)
        self._sync()

    def load_history(self, messages: list[dict]) -> None:
        """Populate the chat panel with a session's stored messages on resume."""
        self._model.load_history(messages)
        self._rebuild_all()

    def handle_ws_event(self, msg: dict) -> None:
        self._model.handle_ws_event(msg)
        self._sync()
        # stream_end closes the turn — flush any throttled assistant rebuild so
        # the final chunk of streamed text appears immediately, not after the
        # throttle window. (No-op for widgets without a pending rebuild.)
        if msg.get("type") == "stream_end":
            self._flush_pending_renders()

    def _flush_pending_renders(self) -> None:
        for widget in self._widgets.values():
            widget.flush()

    def clear(self) -> None:
        """Reset the chat model and redraw an empty conversation."""
        self._model.items.clear()
        self._model._current_assistant = None
        self._model._current_thinking = None
        self._rebuild_all()

    def _max_visible_items(self) -> int:
        return get_tui_settings().max_visible_items

    def _truncation_hint(self, hidden: int) -> RenderableType:
        theme = get_active_theme()
        return Text(
            f"… {hidden} earlier messages not shown",
            style=theme.text.hex,
        )

    def _rebuild_all(self) -> None:
        """Unmount every item widget and re-sync from scratch.

        Used on session switch / clear, where the model is replaced wholesale:
        the surviving-vs-new diff is meaningless when nothing survives.
        """
        for widget in self._widgets.values():
            widget.remove()
        self._widgets.clear()
        self._sync()

    def _sync(self) -> None:
        """Reconcile mounted widgets with the model's visible window.

        The model only ever appends new items and pops/purges existing ones —
        it never reorders survivors — so new widgets always belong at the end
        and a simple mount-at-end keeps the DOM in the right order. Surviving
        items keep their widget and are only re-rendered if their signature
        changed (in-place mutation, e.g. a growing stream).
        """
        # Stick-to-bottom: only auto-follow new content if the user was already
        # at the bottom. Without this, scroll_end() below yanks the view back
        # down on every streamed token, so it is impossible to scroll up and
        # read earlier messages while the agent responds. Capture this BEFORE
        # any DOM change (truncation-hint toggle, mount, update) — those grow
        # max_scroll_y, so is_vertical_scroll_end checked afterwards would be
        # false for the "content just grew" case too and we'd never follow.
        # Auto-follow resumes on its own when the user scrolls back to bottom.
        stick_to_bottom = self.is_vertical_scroll_end
        # Record for _ChatItemView._do_rebuild: a deferred (throttled) rebuild
        # grows the streaming bubble's height after this _sync already ran
        # scroll_end; the rebuild re-anchors only if we were at the bottom here.
        self._stick_to_bottom = stick_to_bottom
        items = self._model.items
        limit = self._max_visible_items()
        hidden = max(0, len(items) - limit)
        # Slice the last `limit` items when over the cap; otherwise keep all.
        visible = items[-limit:] if hidden else items
        visible_ids = {id(it) for it in visible}

        # Truncation hint at the top of the window.
        self._hidden_count = hidden
        if hidden:
            self._hint_view.update(self._truncation_hint(hidden))
            self._hint_view.styles.display = "block"
        else:
            self._hint_view.styles.display = "none"

        # Drop widgets whose item left the visible window (truncated out, or
        # purged from the model by stream_end / plan_ready).
        for key in [k for k in self._widgets if k not in visible_ids]:
            widget = self._widgets.pop(key)
            widget.remove()

        # Walk the visible window in order: mount new items (at the end) and
        # re-render changed ones in place.
        for item in visible:
            key = id(item)
            widget = self._widgets.get(key)
            if widget is None:
                msg = _item_msg(item)
                widget = _ChatItemView(
                    self._render_for(item),
                    item=item,
                    signature=_signature(msg),
                )
                self._widgets[key] = widget
                self.mount(widget)
            else:
                widget.maybe_update(item, self._registry, self._reading_mode)

        if stick_to_bottom:
            self.scroll_end(animate=False)