"""Hierarchical trace viewer for training rollouts.
Navigable tree structure:
    Step → Group (prompt) → Sample (rollout) → Trace (messages)
Based on pyvimdiff pattern: each level is a picker that opens the next level.
Usage:
    from ..tui.traces import StepPicker, TraceData
    data = TraceData.from_jsonl("results/rl/run_xxx/rollouts.jsonl")
    picker = StepPicker(data)
    picker.run()
"""
from __future__ import annotations
import json
import time
from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from .terminal import Terminal
# Colors (ANSI)
DIM = "\x1b[90m"
WHITE = "\x1b[37m"
CYAN = "\x1b[36m"
GREEN = "\x1b[32m"
YELLOW = "\x1b[33m"
RED = "\x1b[31m"
MAGENTA = "\x1b[35m"
BOLD = "\x1b[1m"
RESET = "\x1b[0m"
BG_HEADER = "\x1b[48;2;40;44;52m"
# ─────────────────────────────────────────────────────────────────────────────
# Data Types
# ─────────────────────────────────────────────────────────────────────────────
@dataclass
class Message:
    """A single message in a trace."""
    role: str  # "user", "assistant", "tool", "system"
    content: str
@dataclass
class Rollout:
    """A single rollout (sample)."""
    step: int
    group_id: int  # Which prompt group
    sample_id: int  # Which sample within group
    prompt: str
    response: str
    reward: float
    messages: list[Message] = field(default_factory=list)
    is_streaming: bool = False  # True if still receiving updates
    sample_id_str: str = ""  # Original string sample ID for matching updates
@dataclass
class Group:
    """A group of rollouts for the same prompt."""
    step: int
    group_id: int
    prompt: str
    rollouts: list[Rollout] = field(default_factory=list)
    @property
    def avg_reward(self) -> float:
        if not self.rollouts:
            return 0.0
        return sum(r.reward for r in self.rollouts) / len(self.rollouts)
    @property
    def max_reward(self) -> float:
        if not self.rollouts:
            return 0.0
        return max(r.reward for r in self.rollouts)
    @property
    def min_reward(self) -> float:
        if not self.rollouts:
            return 0.0
        return min(r.reward for r in self.rollouts)
@dataclass
class Step:
    """A training step containing multiple groups."""
    step: int
    groups: list[Group] = field(default_factory=list)
    @property
    def avg_reward(self) -> float:
        all_rewards = [r.reward for g in self.groups for r in g.rollouts]
        if not all_rewards:
            return 0.0
        return sum(all_rewards) / len(all_rewards)
    @property
    def num_rollouts(self) -> int:
        return sum(len(g.rollouts) for g in self.groups)
@dataclass
class TraceData:
    """All trace data from a training run."""
    steps: list[Step] = field(default_factory=list)
    @classmethod
    def from_jsonl(cls, path: str | Path) -> TraceData:
        """Load trace data from ..jsonl file."""
        path = Path(path)
        if not path.exists():
            return cls()
        step_groups: dict[int, dict[int, list[dict]]] = defaultdict(lambda: defaultdict(list))
        with open(path) as f:
            for line in f:
                if not line.strip():
                    continue
                try:
                    record = json.loads(line)
                    step = record.get("step", 0)
                    group_id = record.get("group_index")
                    if group_id is None:
                        prompt = record.get("prompt", "")
                        group_id = hash(str(prompt)) % 10000
                    step_groups[step][group_id].append(record)
                except json.JSONDecodeError:
                    continue
        steps = []
        for step_num in sorted(step_groups.keys()):
            groups = []
            for group_id, records in step_groups[step_num].items():
                if not records:
                    continue
                prompt = records[0].get("prompt", "")
                rollouts = []
                for i, record in enumerate(records):
                    messages = []
                    raw_messages = record.get("messages") or record.get("metadata", {}).get(
                        "messages", []
                    )
                    for msg in raw_messages:
                        messages.append(
                            Message(
                                role=msg.get("role", "unknown"),
                                content=msg.get("content", ""),
                            )
                        )
                    rollouts.append(
                        Rollout(
                            step=step_num,
                            group_id=group_id,
                            sample_id=i,
                            prompt=prompt,
                            response=record.get("response", ""),
                            reward=record.get("reward", 0.0),
                            messages=messages,
                        )
                    )
                groups.append(
                    Group(
                        step=step_num,
                        group_id=group_id,
                        prompt=prompt,
                        rollouts=rollouts,
                    )
                )
            steps.append(Step(step=step_num, groups=groups))
        return cls(steps=steps)
    def add_record(self, record: dict) -> None:
        """Add or update a rollout record (from log stream).
        For streaming rollouts, updates existing rollout if sample_id matches.
        """
        sample_id_str = str(record.get("step", ""))  # The string ID like "sample_0003"
        status = record.get("status", "")
        is_streaming = status == "streaming"
        group_id = record.get("group_index")
        if group_id is None:
            prompt = record.get("prompt", "")
            group_id = hash(str(prompt)) % 10000
        messages = []
        raw_messages = record.get("messages") or record.get("metadata", {}).get("messages", [])
        for msg in raw_messages:
            messages.append(
                Message(
                    role=msg.get("role", "unknown"),
                    content=msg.get("content", ""),
                )
            )
        existing_rollout = self.find_rollout_by_sample_id(sample_id_str)
        if existing_rollout:

            existing_rollout.response = record.get("response", "")
            existing_rollout.reward = record.get("reward", 0.0)
            existing_rollout.messages = messages
            existing_rollout.is_streaming = is_streaming
            return
        step_num = 0

        step = None
        for s in self.steps:
            if s.step == step_num:
                step = s
                break
        if step is None:
            step = Step(step=step_num, groups=[])
            self.steps.append(step)
            self.steps.sort(key=lambda s: s.step)

        group = None
        for g in step.groups:
            if g.group_id == group_id:
                group = g
                break
        if group is None:
            group = Group(
                step=step_num,
                group_id=group_id,
                prompt=record.get("prompt", ""),
                rollouts=[],
            )
            step.groups.append(group)
        rollout = Rollout(
            step=step_num,
            group_id=group_id,
            sample_id=len(group.rollouts),
            prompt=record.get("prompt", ""),
            response=record.get("response", ""),
            reward=record.get("reward", 0.0),
            messages=messages,
            is_streaming=is_streaming,
            sample_id_str=sample_id_str,
        )
        group.rollouts.append(rollout)
    def find_rollout_by_sample_id(self, sample_id_str: str) -> Rollout | None:
        """Find a rollout by its string sample ID."""
        for step in self.steps:
            for group in step.groups:
                for rollout in group.rollouts:
                    if rollout.sample_id_str == sample_id_str:
                        return rollout
        return None
    def handle_stream_event(self, sample_id: str, event_type: str, event_data: dict) -> None:
        """Handle a StreamEvent for live streaming display.
        Uses AgentRenderer-style event handling - builds messages from
        TextEnd, ThinkingEnd, ToolCallEnd events (not manual accumulation).
        """

        rollout = self.find_rollout_by_sample_id(sample_id)
        if rollout is None:
            step_num = 0
            group_id = hash(sample_id) % 10000

            step = None
            for s in self.steps:
                if s.step == step_num:
                    step = s
                    break
            if step is None:
                step = Step(step=step_num, groups=[])
                self.steps.append(step)

            group = None
            for g in step.groups:
                if g.group_id == group_id:
                    group = g
                    break
            if group is None:
                group = Group(step=step_num, group_id=group_id, prompt="", rollouts=[])
                step.groups.append(group)
            rollout = Rollout(
                step=step_num,
                group_id=group_id,
                sample_id=len(group.rollouts),
                prompt="",
                response="",
                reward=0.0,
                messages=[],
                is_streaming=True,
                sample_id_str=sample_id,
            )
            group.rollouts.append(rollout)
        # sample_start contains initial messages including user prompt
        if event_type == "sample_start":
            messages = event_data.get("messages", [])
            for msg in messages:
                rollout.messages.append(
                    Message(role=msg.get("role", "user"), content=msg.get("content", ""))
                )
            # Also set prompt for fallback display
            if messages:
                rollout.prompt = messages[0].get("content", "")
        elif event_type == "TextEnd":
            content = event_data.get("content", "")
            self._update_assistant_text(rollout, content)
        elif event_type == "ThinkingEnd":
            content = event_data.get("content", "")
            self._update_assistant_thinking(rollout, content)
        elif event_type == "ToolCallEnd":
            tool_call = event_data.get("tool_call", {})
            self._add_tool_call(rollout, tool_call)
        elif event_type == "ToolResultReceived":
            tool_call_id = event_data.get("tool_call_id", "")
            content = event_data.get("content", "")
            is_error = event_data.get("is_error", False)
            self._add_tool_result(rollout, tool_call_id, content, is_error)
        elif event_type == "StreamDone":
            # Mark streaming as complete
            rollout.is_streaming = False
        # For delta events, accumulate text for live preview
        elif event_type == "TextDelta":
            delta = event_data.get("delta", "")
            rollout.response += delta
        elif event_type == "ThinkingDelta":
            # Could accumulate for live thinking preview if needed
            pass
    def _update_assistant_text(self, rollout: Rollout, content: str) -> None:
        """Update or create assistant message with text content."""

        if rollout.messages and rollout.messages[-1].role == "assistant":
            rollout.messages[-1] = Message(role="assistant", content=content)
        else:
            rollout.messages.append(Message(role="assistant", content=content))
        rollout.response = content
    def _update_assistant_thinking(self, rollout: Rollout, content: str) -> None:
        """Add thinking content to assistant message."""
        # For now, prepend thinking to content with a marker
        if rollout.messages and rollout.messages[-1].role == "assistant":
            existing = rollout.messages[-1].content
            rollout.messages[-1] = Message(
                role="assistant", content=f"[thinking]\n{content}\n[/thinking]\n\n{existing}"
            )
    def _add_tool_call(self, rollout: Rollout, tool_call: dict) -> None:
        """Add tool call info to display."""
        name = tool_call.get("name", "unknown")
        args = tool_call.get("args", {})
        # Format as text for now
        tool_text = f"[tool_call: {name}]\n{args}"
        if rollout.messages and rollout.messages[-1].role == "assistant":
            existing = rollout.messages[-1].content
            rollout.messages[-1] = Message(role="assistant", content=f"{existing}\n\n{tool_text}")
    def _add_tool_result(
        self, rollout: Rollout, tool_call_id: str, content: str, is_error: bool
    ) -> None:
        """Add tool result message."""
        prefix = "[tool_error]" if is_error else "[tool_result]"
        rollout.messages.append(Message(role="tool", content=f"{prefix}\n{content}"))
# ─────────────────────────────────────────────────────────────────────────────
# Trace Viewer (deepest level - shows message trace)
# ─────────────────────────────────────────────────────────────────────────────
class TraceViewer:
    """View a single rollout's message trace with nvim-style folding.
    Navigation:
    - j/k: scroll lines (skips over folded content)
    - Space/Enter on fold header: toggle fold
    - za/zo/zc: toggle/open/close fold at current line
    - zM/zR: close/open all folds
    """
    def __init__(self, rollout: Rollout) -> None:
        self.rollout = rollout
        self.scroll = 0  # Current line position (cursor line)
        self.h_scroll = 0  # Horizontal scroll
        self.wrap = True  # Wrap mode (False = truncate + h-scroll)
        self.terminal = Terminal(use_alternate_screen=True)
        self._running = False
        self._needs_redraw = True
        self._rendered_lines: list[str] = []
        # Fold state: track which messages are folded (by index)
        self._folds: dict[int, bool] = {}  # msg_index -> is_folded
        self._init_folds()
        # Line metadata for fold operations
        self._line_to_msg: dict[int, int] = {}  # line_index -> msg_index (for fold headers only)
        self._fold_header_lines: set[int] = set()  # lines that are fold headers
    def _init_folds(self) -> None:
        """Initialize fold state - fold all non-user messages by default."""
        messages = self._get_all_messages()
        for i, msg in enumerate(messages):
            # Fold everything except user messages by default
            self._folds[i] = msg.role != "user"
    def _get_all_messages(self) -> list[Message]:
        """Get all messages including fallback assistant response."""
        messages = list(self.rollout.messages)
        has_assistant = any(m.role == "assistant" for m in messages)
        if not has_assistant and self.rollout.response:
            messages.append(Message(role="assistant", content=self.rollout.response))
        return messages
    def run(self) -> None:
        self._running = True
        self._render_messages()  # Pre-render message lines
        self.terminal.start(on_input=lambda x: None, on_resize=self._on_resize)
        try:
            self._main_loop()
        finally:
            self.terminal.stop()
    def _on_resize(self) -> None:
        self._needs_redraw = True
    def _main_loop(self) -> None:
        while self._running:
            if self._needs_redraw:
                self._render()
                self._needs_redraw = False
            data = self.terminal.read_input()
            if data:
                self._handle_input(data)
            time.sleep(0.01)
    def _handle_input(self, data: str) -> None:
        if data == "q":
            self._running = False
            return
        height = self.terminal.rows
        content_height = height - 2
        max_scroll = max(0, len(self._rendered_lines) - content_height)
        # j/k scroll by line
        if data in ("j", "\x1b[B"):
            self.scroll = min(max_scroll, self.scroll + 1)
            self._needs_redraw = True
        elif data in ("k", "\x1b[A"):
            self.scroll = max(0, self.scroll - 1)
            self._needs_redraw = True
        elif data == "\x04":  # Ctrl+D - half page down
            self.scroll = min(max_scroll, self.scroll + content_height // 2)
            self._needs_redraw = True
        elif data == "\x15":  # Ctrl+U - half page up
            self.scroll = max(0, self.scroll - content_height // 2)
            self._needs_redraw = True
        elif data == "g":
            next_char = self._wait_for_char()
            if next_char == "g":
                self.scroll = 0
                self._needs_redraw = True
        elif data == "G":
            self.scroll = max_scroll
            self._needs_redraw = True
        # Horizontal scrolling
        elif data in ("h", "\x1b[D"):  # Left
            self.h_scroll = max(0, self.h_scroll - 20)
            self.wrap = False
            if self.h_scroll == 0:
                self.wrap = True
            self._needs_redraw = True
        elif data in ("l", "\x1b[C"):  # Right
            self.h_scroll += 20
            self.wrap = False
            self._needs_redraw = True
        elif data == "0":  # Beginning
            self.h_scroll = 0
            self._needs_redraw = True
        elif data == "w":  # Toggle wrap
            self.wrap = not self.wrap
            if self.wrap:
                self.h_scroll = 0
            self._needs_redraw = True
        # Fold commands (nvim-style)
        elif data == "z":
            self._handle_fold_command()
        # Tab/Space/Enter to toggle fold if on a fold header
        elif data in ("\t", " ", "\r", "\n"):
            self._toggle_fold_at_scroll()
            self._needs_redraw = True
    def _handle_fold_command(self) -> None:
        """Handle nvim-style fold commands (za, zo, zc, zM, zR)."""
        next_char = self._wait_for_char()
        if next_char is None:
            return
        if next_char == "a":  # za - toggle fold at current line
            self._toggle_fold_at_scroll()
        elif next_char == "o":  # zo - open fold at current line
            self._set_fold_at_scroll(False)
        elif next_char == "c":  # zc - close fold at current line
            self._set_fold_at_scroll(True)
        elif next_char == "M":  # zM - close all folds
            for i in self._folds:
                self._folds[i] = True
            self._render_messages()
        elif next_char == "R":  # zR - open all folds
            for i in self._folds:
                self._folds[i] = False
            self._render_messages()
        self._needs_redraw = True
    def _get_msg_at_scroll(self) -> int | None:
        """Get the message index for the fold header at or before current scroll."""
        if self.scroll in self._line_to_msg:
            return self._line_to_msg[self.scroll]

        for line_idx in range(self.scroll, -1, -1):
            if line_idx in self._line_to_msg:
                return self._line_to_msg[line_idx]
        return None
    def _toggle_fold_at_scroll(self) -> None:
        """Toggle fold for the message at current scroll position."""
        msg_idx = self._get_msg_at_scroll()
        if msg_idx is not None:
            self._folds[msg_idx] = not self._folds.get(msg_idx, False)
            self._render_messages()
    def _set_fold_at_scroll(self, folded: bool) -> None:
        """Set fold state for the message at current scroll position."""
        msg_idx = self._get_msg_at_scroll()
        if msg_idx is not None:
            self._folds[msg_idx] = folded
            self._render_messages()
    def _open_in_nvim(self) -> None:
        """Open trace content in neovim for full vim navigation."""
        import os
        import subprocess
        import tempfile
        content = self._get_plain_text()
        fd, path = tempfile.mkstemp(suffix=".md")
        os.write(fd, content.encode())
        os.close(fd)
        self.terminal.stop()
        try:
            init_cmd = " | ".join([
                "set filetype=markdown",
                "set number",
                "set cursorline",
                "set nomodifiable",
                "nnoremap <silent> q :qa!<CR>",
                "set laststatus=2",
                "set statusline=%f\\ [TRACE]\\ %l/%L",
            ])
            subprocess.run(["nvim", "-R", f"+{init_cmd}", path])
        finally:
            os.unlink(path)
            self.terminal.start(on_input=lambda x: None, on_resize=self._on_resize)
    def _get_plain_text(self) -> str:
        """Get trace content as plain text."""
        lines = []
        for msg in self.rollout.messages:
            lines.append(f"# [{msg.role}]")
            lines.append("")
            for line in msg.content.split("\n"):
                lines.append(line)
            lines.append("")
        return "\n".join(lines)
    def _wait_for_char(self, timeout: float = 0.5) -> str | None:
        start = time.time()
        while time.time() - start < timeout:
            data = self.terminal.read_input()
            if data:
                return data
            time.sleep(0.01)
        return None
    def _render_messages(self) -> None:
        """Pre-render all message lines with fold support.
        Builds _rendered_lines and tracks which lines are fold headers.
        The cursor highlight is applied at render time based on scroll position.
        """
        self._rendered_lines = []
        self._line_to_msg = {}
        self._fold_header_lines = set()
        role_colors = {
            "user": CYAN,
            "assistant": GREEN,
            "tool": YELLOW,
            "system": MAGENTA,
        }
        messages = self._get_all_messages()
        for msg_idx, msg in enumerate(messages):
            color = role_colors.get(msg.role, WHITE)
            is_folded = self._folds.get(msg_idx, False)
            content_lines = msg.content.split("\n")
            num_lines = len(content_lines)

            line_idx = len(self._rendered_lines)
            self._line_to_msg[line_idx] = msg_idx
            self._fold_header_lines.add(line_idx)
            if is_folded:
                # Show folded summary: ▶ [role] --- N lines ---
                fold_marker = "▶"
                preview = content_lines[0][:50] if content_lines else ""
                if content_lines and len(content_lines[0]) > 50:
                    preview += "..."
                header = f"  {fold_marker} {color}{BOLD}[{msg.role}]{RESET} {DIM}--- {num_lines} lines: {preview}{RESET}"
                self._rendered_lines.append(header)
            else:
                # Show expanded: ▼ [role] followed by content
                fold_marker = "▼"
                header = f"  {fold_marker} {color}{BOLD}[{msg.role}]{RESET}"
                self._rendered_lines.append(header)
                for line in content_lines:
                    self._rendered_lines.append(f"    {line}")
            self._rendered_lines.append("")
    def _render(self) -> None:
        width = self.terminal.columns
        height = self.terminal.rows
        content_height = height - 2
        self.terminal.write("\x1b[2J\x1b[H")
        output = []
        # Header
        output.append(self._render_header(width))
        # Content - either wrap or truncate based on mode
        visible = self._rendered_lines[self.scroll : self.scroll + content_height]
        for ln in visible:
            if self.wrap:
                # Wrap long lines
                while ln and len(output) < content_height + 1:
                    chunk = ln[: width - 1]
                    ln = ln[width - 1 :]
                    output.append(chunk + " " * max(0, width - len(chunk)))
                if not ln:
                    continue
            else:

                display_ln = ln
                if self.h_scroll > 0:
                    display_ln = (
                        display_ln[self.h_scroll :] if self.h_scroll < len(display_ln) else ""
                    )
                if len(display_ln) > width - 1:
                    display_ln = display_ln[: width - 4] + "..."
                output.append(display_ln + " " * max(0, width - len(display_ln)))
        # Pad
        while len(output) < height - 1:
            output.append(" " * width)
        # Footer
        output.append(self._render_footer(width))
        for i, line in enumerate(output):
            self.terminal.write(f"\x1b[{i + 1};1H{line}")
    def _render_header(self, width: int) -> str:
        r = self.rollout
        reward_color = GREEN if r.reward > 0 else RED if r.reward < 0 else WHITE
        left = f" Step {r.step} / Group {r.group_id} / Sample {r.sample_id}"
        right = f"reward: {reward_color}{r.reward:.3f}{RESET}{BG_HEADER} "
        padding = width - len(left) - len(f"reward: {r.reward:.3f} ")
        return f"{BG_HEADER}{BOLD}{left}{RESET}{BG_HEADER}{' ' * max(0, padding)}{right}{RESET}"
    def _render_footer(self, width: int) -> str:
        wrap_hint = "w:truncate" if self.wrap else "h/l:scroll  w:wrap"
        hints = f"j/k:move  Space:fold  zR/zM:all  {wrap_hint}  q:back"
        messages = self._get_all_messages()
        num_messages = len(messages)
        pos = f"msg {self.cursor + 1}/{num_messages}"
        padding = width - len(hints) - len(pos) - 4
        return (
            f"{BG_HEADER} {DIM}{hints}{RESET}{BG_HEADER}{' ' * max(0, padding)}{WHITE}{pos} {RESET}"
        )
# ─────────────────────────────────────────────────────────────────────────────
# Streaming Trace Viewer (live updating)
# ─────────────────────────────────────────────────────────────────────────────
class TraceStreamingViewer:
    """View a streaming rollout with live token updates, nvim-style folding and cursor.
    Similar to TraceViewer but polls the Rollout object for updates and
    auto-scrolls as new content arrives.
    """
    def __init__(
        self,
        rollout: Rollout,
    ) -> None:
        """Initialize streaming viewer.
        Args:
            rollout: The Rollout object (gets updated by handle_stream_event)
        """
        self.rollout = rollout
        self.sample_id = rollout.sample_id_str or str(rollout.sample_id)
        self.scroll = 0
        self.h_scroll = 0
        self.wrap = True
        self.auto_scroll = True  # Follow tail (also moves cursor to last message)
        self.terminal = Terminal(use_alternate_screen=True)
        self._running = False
        self._needs_redraw = True
        self._last_response_len = 0  # Track changes
        self._rendered_lines: list[str] = []
        # Cursor: which message is currently selected
        self.cursor = 0  # Index into messages list
        # Fold state: track which messages are folded (by index)
        # Start with all messages folded except user messages
        self._folds: dict[int, bool] = {}  # msg_index -> is_folded
        self._msg_to_line: dict[int, int] = {}  # msg_index -> first rendered line
        self._last_msg_count = 0  # Track when new messages arrive
    def _init_folds_for_new_messages(self, messages: list) -> None:
        """Initialize fold state for any new messages."""
        for i, msg in enumerate(messages):
            if i not in self._folds:
                # Fold everything except user messages by default
                role = msg.role if hasattr(msg, "role") else msg.get("role", "")
                self._folds[i] = role != "user"
    def run(self) -> None:
        self._running = True
        self._render_content()
        self.terminal.start(on_input=lambda x: None, on_resize=self._on_resize)
        try:
            self._main_loop()
        finally:
            self.terminal.stop()
    def _on_resize(self) -> None:
        self._needs_redraw = True
    def _main_loop(self) -> None:
        while self._running:
            current_len = len(self.rollout.response)
            if current_len != self._last_response_len:
                self._render_content()
                self._last_response_len = current_len
                self._needs_redraw = True
                # Auto-scroll: move cursor to last message and scroll to it
                if self.auto_scroll:
                    messages = self._get_all_messages()
                    if messages:
                        self.cursor = len(messages) - 1
                    height = self.terminal.rows
                    content_height = height - 2
                    self._scroll_to_cursor(content_height)
            if self._needs_redraw:
                self._render()
                self._needs_redraw = False
            data = self.terminal.read_input()
            if data:
                self._handle_input(data)
            time.sleep(0.02)  # 50fps for smooth streaming
    def _get_all_messages(self) -> list[Message]:
        """Get all messages including prompt fallback and streaming response."""
        messages: list[Message] = []
        # Show prompt as user message if no messages yet
        has_user_message = any(msg.role == "user" for msg in self.rollout.messages)
        if not has_user_message and self.rollout.prompt:
            messages.append(Message(role="user", content=self.rollout.prompt))
        messages.extend(self.rollout.messages)
        if self.rollout.response:
            messages.append(Message(role="assistant", content=self.rollout.response))
        return messages
    def _scroll_to_cursor(self, content_height: int) -> None:
        """Ensure cursor is visible by adjusting scroll."""
        if self.cursor not in self._msg_to_line:
            return
        cursor_line = self._msg_to_line[self.cursor]
        # If cursor is above visible area, scroll up
        if cursor_line < self.scroll:
            self.scroll = cursor_line
        # If cursor is below visible area, scroll down
        elif cursor_line >= self.scroll + content_height:
            self.scroll = cursor_line - content_height + 1
    def _handle_input(self, data: str) -> None:
        if data == "q":
            self._running = False
            return
        messages = self._get_all_messages()
        num_messages = len(messages)
        height = self.terminal.rows
        content_height = height - 2
        # j/k move cursor between messages
        if data in ("j", "\x1b[B"):
            if self.cursor < num_messages - 1:
                self.cursor += 1
                self._scroll_to_cursor(content_height)
                self.auto_scroll = self.cursor >= num_messages - 1
            self._needs_redraw = True
        elif data in ("k", "\x1b[A"):
            if self.cursor > 0:
                self.cursor -= 1
                self._scroll_to_cursor(content_height)
                self.auto_scroll = False
            self._needs_redraw = True
        elif data == "\x04":  # Ctrl+D - move cursor down by half page worth of messages
            self.cursor = min(num_messages - 1, self.cursor + 5)
            self._scroll_to_cursor(content_height)
            self.auto_scroll = self.cursor >= num_messages - 1
            self._needs_redraw = True
        elif data == "\x15":  # Ctrl+U - move cursor up by half page worth of messages
            self.cursor = max(0, self.cursor - 5)
            self._scroll_to_cursor(content_height)
            self.auto_scroll = False
            self._needs_redraw = True
        elif data == "g":
            next_char = self._wait_for_char()
            if next_char == "g":
                self.cursor = 0
                self._scroll_to_cursor(content_height)
                self.auto_scroll = False
                self._needs_redraw = True
        elif data == "G":
            self.cursor = num_messages - 1
            self._scroll_to_cursor(content_height)
            self.auto_scroll = True
            self._needs_redraw = True
        elif data in ("h", "\x1b[D"):
            self.h_scroll = max(0, self.h_scroll - 20)
            self.wrap = self.h_scroll == 0
            self._needs_redraw = True
        elif data in ("l", "\x1b[C"):
            self.h_scroll += 20
            self.wrap = False
            self._needs_redraw = True
        elif data == "0":
            self.h_scroll = 0
            self._needs_redraw = True
        elif data == "w":
            self.wrap = not self.wrap
            if self.wrap:
                self.h_scroll = 0
            self._needs_redraw = True
        # Fold commands (nvim-style)
        elif data == "z":
            self._handle_fold_command()
        # Tab/Space/Enter to toggle fold at cursor
        elif data in ("\t", " ", "\r", "\n"):
            self._toggle_fold_at_cursor()
            self._needs_redraw = True
    def _handle_fold_command(self) -> None:
        """Handle nvim-style fold commands (za, zo, zc, zM, zR)."""
        next_char = self._wait_for_char()
        if next_char is None:
            return
        if next_char == "a":  # za - toggle fold at cursor
            self._toggle_fold_at_cursor()
        elif next_char == "o":  # zo - open fold at cursor
            self._set_fold_at_cursor(False)
        elif next_char == "c":  # zc - close fold at cursor
            self._set_fold_at_cursor(True)
        elif next_char == "M":  # zM - close all folds
            for i in self._folds:
                self._folds[i] = True
            self._render_content()
        elif next_char == "R":  # zR - open all folds
            for i in self._folds:
                self._folds[i] = False
            self._render_content()
        self._needs_redraw = True
    def _toggle_fold_at_cursor(self) -> None:
        """Toggle fold for the message at cursor."""
        self._folds[self.cursor] = not self._folds.get(self.cursor, False)
        self._render_content()
    def _set_fold_at_cursor(self, folded: bool) -> None:
        """Set fold state for the message at cursor."""
        self._folds[self.cursor] = folded
        self._render_content()
    def _get_plain_text(self) -> str:
        """Get trace content as plain text for nvim."""
        lines = []
        for msg in self._get_messages():
            role = msg.get("role", "unknown")
            content = msg.get("content", "")
            lines.append(f"[{role}]")
            lines.append("")
            for line in content.split("\n"):
                lines.append(line)
            lines.append("")
        return "\n".join(lines)
    def _open_in_nvim(self) -> None:
        """Open streaming trace in neovim with auto-reload."""
        import os
        import subprocess
        import tempfile
        import threading
        content = self._get_plain_text()
        fd, path = tempfile.mkstemp(suffix=".md")
        os.write(fd, content.encode())
        os.close(fd)
        self.terminal.stop()
        stop_streaming = threading.Event()
        last_hash = [hash(content)]  # Use list to allow mutation in closure
        def stream_updates() -> None:
            while not stop_streaming.is_set() and self._is_streaming:
                new_content = self._get_plain_text()
                new_hash = hash(new_content)
                if new_hash != last_hash[0]:
                    with open(path, "w") as f:
                        f.write(new_content)
                    last_hash[0] = new_hash
                time.sleep(0.1)
        thread = threading.Thread(target=stream_updates, daemon=True)
        thread.start()
        try:
            init_cmd = " | ".join([
                "set filetype=markdown",
                "set number",
                "set cursorline",
                "set nomodifiable",
                "set autoread",
                # Timer-based auto-reload (every 200ms, even while scrolling)
                "let g:stream_timer = timer_start(200, {-> execute('silent! checktime')}, {'repeat': -1})",
                "autocmd FileChangedShellPost * set modifiable | silent! %d | silent! read | 1d | set nomodifiable | normal! G",
                "nnoremap <silent> q :call timer_stop(g:stream_timer) \\| qa!<CR>",
                "set laststatus=2",
                "set statusline=%f\\ [STREAMING]\\ %l/%L",
                "normal! G",
            ])
            subprocess.run(["nvim", "-R", f"+{init_cmd}", path])
        finally:
            stop_streaming.set()
            os.unlink(path)
            self.terminal.start(on_input=lambda x: None, on_resize=self._on_resize)
    def _wait_for_char(self, timeout: float = 0.5) -> str | None:
        start = time.time()
        while time.time() - start < timeout:
            data = self.terminal.read_input()
            if data:
                return data
            time.sleep(0.01)
        return None
    def _render_content(self) -> None:
        """Render content lines from rollout with fold support and cursor highlight."""
        self._rendered_lines = []
        self._msg_to_line = {}
        role_colors = {
            "user": CYAN,
            "assistant": GREEN,
            "tool": YELLOW,
            "system": MAGENTA,
        }
        # Cursor styling
        BG_CURSOR = "\x1b[48;2;50;50;30m"  # Muted olive background
        CURSOR_FG = "\x1b[38;2;255;200;0m"  # Bright yellow/gold foreground
        messages = self._get_all_messages()
        self._init_folds_for_new_messages(messages)
        # Render each message with fold support and cursor
        for msg_idx, msg in enumerate(messages):
            color = role_colors.get(msg.role, WHITE)
            is_folded = self._folds.get(msg_idx, False)
            is_cursor = msg_idx == self.cursor
            content_lines = msg.content.split("\n")
            num_lines = len(content_lines)

            line_idx = len(self._rendered_lines)
            self._msg_to_line[msg_idx] = line_idx
            # Cursor indicator - bright arrow when selected, space otherwise
            if is_cursor:
                cursor_marker = f"{CURSOR_FG}{BOLD}▸{RESET}{BG_CURSOR}"
                bg = BG_CURSOR
                bg_reset = RESET
            else:
                cursor_marker = " "
                bg = ""
                bg_reset = ""
            if is_folded:
                # Show folded summary: ▸ ▶ [role] --- N lines ---
                fold_marker = "▶"
                preview = content_lines[0][:40] if content_lines else ""
                if len(content_lines[0]) > 40 if content_lines else False:
                    preview += "..."
                summary = f"{bg}{cursor_marker} {fold_marker} {color}{BOLD}[{msg.role}]{RESET}{bg} {DIM}--- {num_lines} lines: {preview}{RESET}{bg_reset}"
                self._rendered_lines.append(summary)
            else:
                # Show expanded: ▸ ▼ [role] followed by content
                fold_marker = "▼"
                header = (
                    f"{bg}{cursor_marker} {fold_marker} {color}{BOLD}[{msg.role}]{RESET}{bg_reset}"
                )
                self._rendered_lines.append(header)
                for line in content_lines:
                    self._rendered_lines.append(f"    {line}")
            self._rendered_lines.append("")
    def _render(self) -> None:
        width = self.terminal.columns
        height = self.terminal.rows
        content_height = height - 2
        self.terminal.write("\x1b[2J\x1b[H")
        output = []
        output.append(self._render_header(width))
        visible = self._rendered_lines[self.scroll : self.scroll + content_height]
        for ln in visible:
            if self.wrap:
                while ln and len(output) < content_height + 1:
                    chunk = ln[: width - 1]
                    ln = ln[width - 1 :]
                    output.append(chunk + " " * max(0, width - len(chunk)))
                if not ln:
                    continue
            else:
                display_ln = ln
                if self.h_scroll > 0:
                    display_ln = (
                        display_ln[self.h_scroll :] if self.h_scroll < len(display_ln) else ""
                    )
                if len(display_ln) > width - 1:
                    display_ln = display_ln[: width - 4] + "..."
                output.append(display_ln + " " * max(0, width - len(display_ln)))
        while len(output) < height - 1:
            output.append(" " * width)
        output.append(self._render_footer(width))
        for i, line in enumerate(output):
            self.terminal.write(f"\x1b[{i + 1};1H{line}")
    def _render_header(self, width: int) -> str:
        sample_id = self.sample_id
        char_count = len(self.rollout.response)
        if self.rollout.is_streaming:
            status = f"{YELLOW}● STREAMING{RESET}{BG_HEADER}"
        else:
            reward = self.rollout.reward or 0.0
            reward_color = GREEN if reward > 0 else RED if reward < 0 else WHITE
            status = f"{reward_color}reward: {reward:.3f}{RESET}{BG_HEADER}"
        left = f" [{sample_id}] {status} ({char_count} chars)"
        padding = width - len(f" [{sample_id}]  STREAMING ({char_count} chars)")
        return f"{BG_HEADER}{BOLD}{left}{RESET}{BG_HEADER}{' ' * max(0, padding)}{RESET}"
    def _render_footer(self, width: int) -> str:
        wrap_hint = "w:truncate" if self.wrap else "h/l:scroll  w:wrap"
        hints = f"j/k:move  Space:fold  zR/zM:all  {wrap_hint}  q:back"
        messages = self._get_all_messages()
        num_messages = len(messages)
        pos = f"msg {self.cursor + 1}/{num_messages}"
        if self.auto_scroll and self.rollout.is_streaming:
            pos += " [FOLLOW]"
        padding = width - len(hints) - len(pos) - 6
        return (
            f"{BG_HEADER} {DIM}{hints}{RESET}{BG_HEADER}{' ' * max(0, padding)}{WHITE}{pos} {RESET}"
        )
# ─────────────────────────────────────────────────────────────────────────────
# Rollout Picker (samples within a group)
# ─────────────────────────────────────────────────────────────────────────────
class RolloutPicker:
    """Pick a rollout from a group."""
    def __init__(self, group: Group, line_queue: list[str] | None = None) -> None:
        self.group = group
        self.line_queue = line_queue
        self.cursor = 0
        self.scroll = 0
        self.terminal = Terminal(use_alternate_screen=True)
        self._running = False
        self._needs_redraw = True
    def run(self) -> None:
        if not self.group.rollouts:
            return
        self._running = True
        self.terminal.start(on_input=lambda x: None, on_resize=self._on_resize)
        try:
            self._main_loop()
        finally:
            self.terminal.stop()
    def _on_resize(self) -> None:
        self._needs_redraw = True
    def _main_loop(self) -> None:
        while self._running:
            if self._needs_redraw:
                self._render()
                self._needs_redraw = False
            data = self.terminal.read_input()
            if data:
                self._handle_input(data)
            time.sleep(0.01)
    def _handle_input(self, data: str) -> None:
        if data == "q":
            self._running = False
            return
        height = self.terminal.rows
        content_height = height - 2
        if data in ("j", "\x1b[B"):
            if self.cursor < len(self.group.rollouts) - 1:
                self.cursor += 1
                if self.cursor >= self.scroll + content_height:
                    self.scroll = self.cursor - content_height + 1
                self._needs_redraw = True
        elif data in ("k", "\x1b[A"):
            if self.cursor > 0:
                self.cursor -= 1
                if self.cursor < self.scroll:
                    self.scroll = self.cursor
                self._needs_redraw = True
        elif data == "g":
            next_char = self._wait_for_char()
            if next_char == "g":
                self.cursor = 0
                self.scroll = 0
                self._needs_redraw = True
        elif data == "G":
            self.cursor = len(self.group.rollouts) - 1
            self.scroll = max(0, self.cursor - content_height + 1)
            self._needs_redraw = True
        elif data in ("\r", "\n"):
            self._open_trace()
    def _wait_for_char(self, timeout: float = 0.5) -> str | None:
        start = time.time()
        while time.time() - start < timeout:
            data = self.terminal.read_input()
            if data:
                return data
            time.sleep(0.01)
        return None
    def _open_trace(self) -> None:
        rollout = self.group.rollouts[self.cursor]
        self.terminal.stop()
        if rollout.is_streaming:
            viewer = TraceStreamingViewer(rollout)
        else:
            viewer = TraceViewer(rollout)
        viewer.run()
        self.terminal.start(on_input=lambda x: None, on_resize=self._on_resize)
        self._needs_redraw = True
    def _render(self) -> None:
        width = self.terminal.columns
        height = self.terminal.rows
        content_height = height - 2
        self.terminal.write("\x1b[2J\x1b[H")
        output = []
        output.append(self._render_header(width))
        visible = self.group.rollouts[self.scroll : self.scroll + content_height]
        for i, rollout in enumerate(visible):
            idx = self.scroll + i
            output.append(self._render_rollout_line(rollout, idx == self.cursor, width))
        while len(output) < height - 1:
            output.append(" " * width)
        output.append(self._render_footer(width))
        for i, line in enumerate(output):
            self.terminal.write(f"\x1b[{i + 1};1H{line}")
    def _render_header(self, width: int) -> str:
        g = self.group

        prompt_preview = g.prompt[:50] + "..." if len(g.prompt) > 50 else g.prompt
        prompt_preview = prompt_preview.replace("\n", " ")
        left = f" Step {g.step} / Group {g.group_id}: {prompt_preview}"
        if len(left) > width - 20:
            left = left[: width - 23] + "..."
        padding = width - len(left)
        return f"{BG_HEADER}{BOLD}{left}{RESET}{BG_HEADER}{' ' * max(0, padding)}{RESET}"
    def _render_rollout_line(self, r: Rollout, selected: bool, width: int) -> str:
        cursor = "> " if selected else "  "
        # Show streaming indicator or reward
        if r.is_streaming:
            char_count = len(r.response)
            status_str = f"{YELLOW}● streaming ({char_count} chars){RESET}"
        else:
            reward_color = GREEN if r.reward > 0.5 else YELLOW if r.reward > 0 else RED
            status_str = f"reward={reward_color}{r.reward:.3f}{RESET}"
        # Response preview
        response_preview = r.response[:50].replace("\n", " ")
        if len(r.response) > 50:
            response_preview += "..."
        sample_label = r.sample_id_str if r.sample_id_str else f"Sample {r.sample_id}"
        if selected:
            line = f"{BOLD}{cursor}{sample_label}  {status_str}  {response_preview}{RESET}"
        else:
            line = (
                f"{DIM}{cursor}{sample_label}  {RESET}{status_str}{DIM}  {response_preview}{RESET}"
            )
        return line
    def _render_footer(self, width: int) -> str:
        hints = "j/k:move  Enter:view  q:back"
        pos = f"{self.cursor + 1}/{len(self.group.rollouts)}"
        padding = width - len(hints) - len(pos) - 4
        return (
            f"{BG_HEADER} {DIM}{hints}{RESET}{BG_HEADER}{' ' * max(0, padding)}{WHITE}{pos} {RESET}"
        )
# ─────────────────────────────────────────────────────────────────────────────
# ─────────────────────────────────────────────────────────────────────────────
class GroupPicker:
    """Pick a group from a step."""
    def __init__(self, step: Step, line_queue: list[str] | None = None) -> None:
        self.step = step
        self.line_queue = line_queue
        self.cursor = 0
        self.scroll = 0
        self.terminal = Terminal(use_alternate_screen=True)
        self._running = False
        self._needs_redraw = True
    def run(self) -> None:
        if not self.step.groups:
            return
        self._running = True
        self.terminal.start(on_input=lambda x: None, on_resize=self._on_resize)
        try:
            self._main_loop()
        finally:
            self.terminal.stop()
    def _on_resize(self) -> None:
        self._needs_redraw = True
    def _main_loop(self) -> None:
        while self._running:
            if self._needs_redraw:
                self._render()
                self._needs_redraw = False
            data = self.terminal.read_input()
            if data:
                self._handle_input(data)
            time.sleep(0.01)
    def _handle_input(self, data: str) -> None:
        if data == "q":
            self._running = False
            return
        height = self.terminal.rows
        content_height = height - 2
        if data in ("j", "\x1b[B"):
            if self.cursor < len(self.step.groups) - 1:
                self.cursor += 1
                if self.cursor >= self.scroll + content_height:
                    self.scroll = self.cursor - content_height + 1
                self._needs_redraw = True
        elif data in ("k", "\x1b[A"):
            if self.cursor > 0:
                self.cursor -= 1
                if self.cursor < self.scroll:
                    self.scroll = self.cursor
                self._needs_redraw = True
        elif data == "g":
            next_char = self._wait_for_char()
            if next_char == "g":
                self.cursor = 0
                self.scroll = 0
                self._needs_redraw = True
        elif data == "G":
            self.cursor = len(self.step.groups) - 1
            self.scroll = max(0, self.cursor - content_height + 1)
            self._needs_redraw = True
        elif data in ("\r", "\n"):
            self._open_group()
    def _wait_for_char(self, timeout: float = 0.5) -> str | None:
        start = time.time()
        while time.time() - start < timeout:
            data = self.terminal.read_input()
            if data:
                return data
            time.sleep(0.01)
        return None
    def _open_group(self) -> None:
        group = self.step.groups[self.cursor]
        self.terminal.stop()
        picker = RolloutPicker(group, line_queue=self.line_queue)
        picker.run()
        self.terminal.start(on_input=lambda x: None, on_resize=self._on_resize)
        self._needs_redraw = True
    def _render(self) -> None:
        width = self.terminal.columns
        height = self.terminal.rows
        content_height = height - 2
        self.terminal.write("\x1b[2J\x1b[H")
        output = []
        output.append(self._render_header(width))
        visible = self.step.groups[self.scroll : self.scroll + content_height]
        for i, group in enumerate(visible):
            idx = self.scroll + i
            output.append(self._render_group_line(group, idx == self.cursor, width))
        while len(output) < height - 1:
            output.append(" " * width)
        output.append(self._render_footer(width))
        for i, line in enumerate(output):
            self.terminal.write(f"\x1b[{i + 1};1H{line}")
    def _render_header(self, width: int) -> str:
        s = self.step
        left = f" Step {s.step} - {len(s.groups)} groups, {s.num_rollouts} rollouts"
        reward_color = GREEN if s.avg_reward > 0.5 else YELLOW if s.avg_reward > 0 else WHITE
        right = f"avg reward: {reward_color}{s.avg_reward:.3f}{RESET}{BG_HEADER} "
        padding = width - len(left) - len(f"avg reward: {s.avg_reward:.3f} ")
        return f"{BG_HEADER}{BOLD}{left}{RESET}{BG_HEADER}{' ' * max(0, padding)}{right}{RESET}"
    def _render_group_line(self, g: Group, selected: bool, width: int) -> str:
        cursor = "> " if selected else "  "
        # Reward range
        reward_color = GREEN if g.avg_reward > 0.5 else YELLOW if g.avg_reward > 0 else RED
        reward_str = f"{reward_color}{g.avg_reward:.3f}{RESET}"
        # Prompt preview
        prompt_preview = g.prompt[:50].replace("\n", " ")
        if len(g.prompt) > 50:
            prompt_preview += "..."
        n_samples = len(g.rollouts)
        if selected:
            line = f"{BOLD}{cursor}Group {g.group_id}  ({n_samples} samples)  avg={reward_str}  {prompt_preview}{RESET}"
        else:
            line = f"{DIM}{cursor}Group {g.group_id}  ({n_samples} samples)  avg={RESET}{reward_str}{DIM}  {prompt_preview}{RESET}"
        return line
    def _render_footer(self, width: int) -> str:
        hints = "j/k:move  Enter:view  q:back"
        pos = f"{self.cursor + 1}/{len(self.step.groups)}"
        padding = width - len(hints) - len(pos) - 4
        return (
            f"{BG_HEADER} {DIM}{hints}{RESET}{BG_HEADER}{' ' * max(0, padding)}{WHITE}{pos} {RESET}"
        )
# ─────────────────────────────────────────────────────────────────────────────
# Step Picker (top level)
# ─────────────────────────────────────────────────────────────────────────────
class StepPicker:
    """Pick a training step."""
    def __init__(self, data: TraceData, line_queue: list[str] | None = None) -> None:
        self.data = data
        self.line_queue = line_queue
        self.cursor = 0
        self.scroll = 0
        self.terminal = Terminal(use_alternate_screen=True)
        self._running = False
        self._needs_redraw = True
    def run(self) -> None:
        if not self.data.steps:
            print("No trace data.")
            return
        self._running = True
        self.terminal.start(on_input=lambda x: None, on_resize=self._on_resize)
        try:
            self._main_loop()
        finally:
            self.terminal.stop()
    def _on_resize(self) -> None:
        self._needs_redraw = True
    def _main_loop(self) -> None:
        while self._running:
            if self._needs_redraw:
                self._render()
                self._needs_redraw = False
            data = self.terminal.read_input()
            if data:
                self._handle_input(data)
            time.sleep(0.01)
    def _handle_input(self, data: str) -> None:
        if data == "q":
            self._running = False
            return
        height = self.terminal.rows
        content_height = height - 2
        if data in ("j", "\x1b[B"):
            if self.cursor < len(self.data.steps) - 1:
                self.cursor += 1
                if self.cursor >= self.scroll + content_height:
                    self.scroll = self.cursor - content_height + 1
                self._needs_redraw = True
        elif data in ("k", "\x1b[A"):
            if self.cursor > 0:
                self.cursor -= 1
                if self.cursor < self.scroll:
                    self.scroll = self.cursor
                self._needs_redraw = True
        elif data == "g":
            next_char = self._wait_for_char()
            if next_char == "g":
                self.cursor = 0
                self.scroll = 0
                self._needs_redraw = True
        elif data == "G":
            self.cursor = len(self.data.steps) - 1
            self.scroll = max(0, self.cursor - content_height + 1)
            self._needs_redraw = True
        elif data in ("\r", "\n"):
            self._open_step()
    def _wait_for_char(self, timeout: float = 0.5) -> str | None:
        start = time.time()
        while time.time() - start < timeout:
            data = self.terminal.read_input()
            if data:
                return data
            time.sleep(0.01)
        return None
    def _open_step(self) -> None:
        step = self.data.steps[self.cursor]
        self.terminal.stop()
        picker = GroupPicker(step, line_queue=self.line_queue)
        picker.run()
        self.terminal.start(on_input=lambda x: None, on_resize=self._on_resize)
        self._needs_redraw = True
    def _render(self) -> None:
        width = self.terminal.columns
        height = self.terminal.rows
        content_height = height - 2
        self.terminal.write("\x1b[2J\x1b[H")
        output = []
        output.append(self._render_header(width))
        visible = self.data.steps[self.scroll : self.scroll + content_height]
        for i, step in enumerate(visible):
            idx = self.scroll + i
            output.append(self._render_step_line(step, idx == self.cursor, width))
        while len(output) < height - 1:
            output.append(" " * width)
        output.append(self._render_footer(width))
        for i, line in enumerate(output):
            self.terminal.write(f"\x1b[{i + 1};1H{line}")
    def _render_header(self, width: int) -> str:
        total_rollouts = sum(s.num_rollouts for s in self.data.steps)
        left = f" Training Steps ({len(self.data.steps)} steps, {total_rollouts} rollouts)"
        padding = width - len(left)
        return f"{BG_HEADER}{BOLD}{left}{RESET}{BG_HEADER}{' ' * max(0, padding)}{RESET}"
    def _render_step_line(self, s: Step, selected: bool, width: int) -> str:
        cursor = "> " if selected else "  "
        # Reward with color
        reward_color = GREEN if s.avg_reward > 0.5 else YELLOW if s.avg_reward > 0 else RED
        reward_str = f"{reward_color}{s.avg_reward:.3f}{RESET}"
        # Stats
        stats = f"{len(s.groups)} groups, {s.num_rollouts} rollouts"
        if selected:
            line = f"{BOLD}{cursor}Step {s.step:>3}  avg_reward={reward_str}  {stats}{RESET}"
        else:
            line = f"{DIM}{cursor}Step {s.step:>3}  avg_reward={RESET}{reward_str}{DIM}  {stats}{RESET}"
        return line
    def _render_footer(self, width: int) -> str:
        hints = "j/k:move  Enter:view  q:quit"
        pos = f"{self.cursor + 1}/{len(self.data.steps)}"
        padding = width - len(hints) - len(pos) - 4
        return (
            f"{BG_HEADER} {DIM}{hints}{RESET}{BG_HEADER}{' ' * max(0, padding)}{WHITE}{pos} {RESET}"
        )
# ─────────────────────────────────────────────────────────────────────────────
# CLI Entry Point
# ─────────────────────────────────────────────────────────────────────────────
def main() -> None:
    """CLI entry point."""
    import sys
    if len(sys.argv) < 2:
        print("Usage: python -m rollouts.tui.traces <rollouts.jsonl>")
        sys.exit(1)
    path = sys.argv[1]
    data = TraceData.from_jsonl(path)
    if not data.steps:
        print(f"No traces found in {path}")
        sys.exit(1)
    picker = StepPicker(data)
    picker.run()
if __name__ == "__main__":
    main()
