execution_tree.py

import types

from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, ScrollableContainer
from textual.screen import ModalScreen
from textual.widgets import Footer, Header, Static

from rich.text import Text

from tracer import TraceStep


# Nerd font icons
ICON_FUNCTION = "\uf120"
ICON_MODULE = "\uf405"
ICON_VARIABLE = "\uf481"
ICON_CURRENT = "\uf04b"
ICON_STEP = "\uf051"
ICON_ARROW_R = "\uf054"
ICON_LINE = "\uf016"
ICON_CHECK = "\uf00c"

# Colors
C_VAR = "bold green"
C_REF = "bold blue"
C_TYPE = "bold yellow"
C_STR_NUM = "white"
C_NONE = "dim"
C_DIM = "dim"

PRIMITIVE_TYPES = (int, float, str, bool, type(None), bytes)
CALLABLE_TYPES = (
    types.FunctionType, types.MethodType,
    types.GeneratorType, types.CoroutineType, types.AsyncGeneratorType,
    type, types.LambdaType,
    types.BuiltinFunctionType, types.BuiltinMethodType,
)


def _format_primitive(obj) -> str:
    if obj is None:
        return "None"
    if isinstance(obj, bool):
        return "True" if obj else "False"
    if isinstance(obj, str):
        return repr(obj)
    if isinstance(obj, bytes):
        return repr(obj)
    return repr(obj)


def _walk_vars_objects(vars_dict, id_registry, counter, parent_ids=None):
    """Walk objects in a variable dict assigning stable hex IDs."""
    if parent_ids is None:
        parent_ids = set()
    for val in vars_dict.values():
        _walk_val(val, id_registry, counter, parent_ids)


def _walk_val(obj, id_registry, counter, parent_ids):
    obj_id = id(obj)
    if isinstance(obj, PRIMITIVE_TYPES) or isinstance(obj, CALLABLE_TYPES):
        return
    if obj_id in id_registry:
        return
    if obj_id in parent_ids:
        return
    counter[0] += 1
    id_registry[obj_id] = counter[0]
    new_parents = parent_ids | {obj_id}
    if isinstance(obj, dict):
        for v in obj.values():
            _walk_val(v, id_registry, counter, new_parents)
    elif isinstance(obj, (list, tuple, set, frozenset)):
        for v in obj:
            _walk_val(v, id_registry, counter, new_parents)
    elif hasattr(obj, "__dict__"):
        for v in vars(obj).values():
            _walk_val(v, id_registry, counter, new_parents)


def _format_var_line(name, val, id_registry) -> Text:
    """Format a variable line with primitives inline and refs as → 0x{id}."""
    if isinstance(val, PRIMITIVE_TYPES):
        return Text.assemble(
            Text(name, style=C_VAR),
            Text(" = ", style=C_DIM),
            Text(_format_primitive(val), style=C_STR_NUM),
        )
    if isinstance(val, CALLABLE_TYPES):
        val_repr = repr(val)
        return Text.assemble(
            Text(name, style=C_VAR),
            Text(" = ", style=C_DIM),
            Text(val_repr[:45] + "..." if len(val_repr) > 45 else val_repr, style="purple"),
        )
    ref_id = id_registry.get(id(val))
    if ref_id is not None:
        return Text.assemble(
            Text(name, style=C_VAR),
            Text(f" → 0x{ref_id:x}", style=C_REF),
        )
    return Text.assemble(
        Text(name, style=C_VAR),
        Text(" = ...", style=C_DIM),
    )


NODE_WIDTH = 72
NODE_HEIGHT = 16


class VarsBar(Static):
    """Compact horizontal bar showing all current variables."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._vars: dict[str, object] = {}

    def update_vars(self, locals_dict: dict, globals_dict: dict):
        self._vars = {}
        for name, val in globals_dict.items():
            if name.startswith("__") or name in ("__builtins__",):
                continue
            if name in locals_dict:
                continue
            self._vars[name] = val
        for name, val in locals_dict.items():
            if name.startswith("__"):
                continue
            self._vars[name] = val

        if not self._vars:
            self.update("[dim]No variables in scope[/dim]")
            return

        parts = []
        for name, val in self._vars.items():
            val_str = repr(val)
            if len(val_str) > 14:
                val_str = val_str[:12] + ".."
            parts.append(f"[bold green]{name}[/bold green] = [white]{val_str}[/white]")

        self.update(f" {ICON_VARIABLE}  " + "  [dim]│[/dim]  ".join(parts))


def _get_code_context(source_lines: list[str], line: int) -> list[str]:
    """Get 3 lines of context around the given line."""
    if not source_lines:
        return ["", "", ""]
    
    start = max(1, line - 1)
    end = min(len(source_lines), line + 1)
    
    context = []
    for l in range(start, end + 1):
        if 1 <= l <= len(source_lines):
            prefix = "▶ " if l == line else "  "
            line_content = source_lines[l-1]
            max_code_len = NODE_WIDTH - 10
            if len(line_content) > max_code_len:
                line_content = line_content[:max_code_len - 3] + "..."
            context.append(f"{prefix}{l:4d} {line_content}")
    
    while len(context) < 3:
        context.append("")
    
    return context[:3]


def _build_node_text(step: TraceStep, index: int, is_current: bool, source_lines: list[str]) -> Text:
    """Build node content as Rich Text to avoid markup parsing issues."""
    text = Text()
    
    # Determine node type icon and color
    if step.event == "call":
        icon = ICON_FUNCTION
        icon_color = "yellow"
        event_text = "call"
    elif step.func_name == "<module>":
        icon = ICON_MODULE
        icon_color = "blue"
        event_text = "module"
    elif step.event == "return":
        icon = ICON_CHECK
        icon_color = "green"
        event_text = "return"
    elif step.event == "exception":
        icon = "!"
        icon_color = "red"
        event_text = "exception"
    else:
        icon = ICON_VARIABLE
        icon_color = "cyan"
        event_text = "line"

    # Header
    if is_current:
        text.append(f"{ICON_CURRENT} STEP {index + 1}", style="bold yellow")
    else:
        text.append(f"{icon} ", style=icon_color)
        text.append(f"Step {index + 1}", style="bold")
    text.append("\n")

    # Location info
    text.append("  ", style="")
    text.append(f"{ICON_LINE} ", style="dim")
    text.append(f"Line {step.line}", style="dim")
    text.append(" │ ", style="dim")
    text.append(step.func_name, style="italic")
    text.append(" │ ", style="dim")
    text.append(event_text, style="dim")
    text.append("\n")

    # Separator
    text.append(" " + "─" * (NODE_WIDTH - 4) + "\n", style="dim")

    # Code context
    context = _get_code_context(source_lines, step.line)
    for ctx_line in context:
        text.append("  ")
        text.append(ctx_line)
        text.append("\n")

    # Separator
    text.append(" " + "─" * (NODE_WIDTH - 4) + "\n", style="dim")

    # Variables
    all_vars = {}
    for name, val in step.globals_copy.items():
        if name.startswith("__") or name in ("__builtins__",):
            continue
        if name in step.locals_copy:
            continue
        all_vars[name] = val
    for name, val in step.locals_copy.items():
        if name.startswith("__"):
            continue
        all_vars[name] = val

    # Walk objects for reference IDs
    id_reg: dict[int, int] = {}
    counter = [0]
    _walk_vars_objects(all_vars, id_reg, counter)

    var_count = 0
    for name in sorted(all_vars.keys()):
        val = all_vars[name]
        if var_count >= 4:
            text.append("  ...\n", style=C_DIM)
            break

        line = _format_var_line(name, val, id_reg)
        text.append_text(line)
        text.append("\n")
        var_count += 1

    if var_count == 0:
        text.append("  (no variables)\n", style="dim")

    # Pad to minimum height
    current_lines = len(text.plain.split("\n"))
    while current_lines < NODE_HEIGHT:
        text.append("\n")
        current_lines += 1

    return text


def _make_node(step: TraceStep, index: int, is_current: bool, source_lines: list[str]) -> Static:
    node_text = _build_node_text(step, index, is_current, source_lines)
    css_class = "tree-node current" if is_current else "tree-node"
    return Static(node_text, classes=css_class)


def _make_arrow() -> Static:
    lines = [""] * NODE_HEIGHT
    mid = NODE_HEIGHT // 2
    lines[mid - 1] = " ╱"
    lines[mid] = f" ──{ICON_ARROW_R}──"
    lines[mid + 1] = " ╲"
    return Static("\n".join(lines), classes="tree-arrow")


class TreeScreen(ModalScreen):
    """Full-screen overlay showing execution tree."""

    CSS = """
    TreeScreen {
        background: $background;
        layout: vertical;
    }

    #tree_vars_bar {
        height: 3;
        background: $panel;
        padding: 0 1;
        overflow-x: auto;
        overflow-y: hidden;
        border-bottom: solid $primary;
    }

    #tree_info {
        height: 3;
        background: $accent 15%;
        padding: 0 1;
        content-align: left middle;
    }

    #tree_container {
        height: 1fr;
        padding: 1 2;
        overflow-x: auto;
        overflow-y: hidden;
    }

    #tree_nodes {
        height: auto;
        width: auto;
        padding: 1 0;
        layout: horizontal;
    }

    #tree_hint {
        height: 1;
        background: $panel;
        color: $text-muted;
        padding: 0 2;
        border-top: solid $secondary;
    }

    .tree-node {
        width: 72;
        min-height: 16;
        max-height: 16;
        border: tall $secondary;
        padding: 0 1;
        margin: 0 1;
        overflow: hidden;
        background: $surface;
    }

    .tree-node.current {
        border: tall $primary;
        background: $primary 10%;
    }

    .tree-arrow {
        width: 8;
        min-height: 16;
        max-height: 16;
        content-align: center middle;
        color: $text-muted;
    }
    """

    BINDINGS = [
        Binding("left", "step_backward", "Step back", show=True, priority=True),
        Binding("right", "step_forward", "Step forward", show=True, priority=True),
        Binding("escape", "close", "Close", show=True, priority=True),
        Binding("t", "close", "Close", show=True, priority=True),
    ]

    def __init__(self, steps: list[TraceStep], current: int = 0, source_lines: list[str] = None, **kwargs):
        super().__init__(**kwargs)
        self.steps = steps
        self.current = current
        self.source_lines = source_lines or []
        self._on_step_change = None

    def set_step_callback(self, callback):
        self._on_step_change = callback

    def compose(self) -> ComposeResult:
        yield Header()
        yield VarsBar(id="tree_vars_bar")
        yield Static(id="tree_info")
        with ScrollableContainer(id="tree_container", can_focus=False):
            with Horizontal(id="tree_nodes"):
                for i, step in enumerate(self.steps):
                    yield _make_node(step, i, i == self.current, self.source_lines)
                    if i < len(self.steps) - 1:
                        yield _make_arrow()
        yield Static(" ←/→ Step │ Esc Close ", id="tree_hint")
        yield Footer()

    def on_mount(self) -> None:
        # Ensure modal owns key handling (no leftover focus on CodeView below).
        self.set_focus(None)
        self._update_info()
        self.call_after_refresh(self._scroll_to_current)

    def _update_info(self):
        info = self.query_one("#tree_info", Static)
        step = self.steps[self.current]

        if step.event == "call":
            event_icon = ICON_FUNCTION
        elif step.event == "return":
            event_icon = ICON_CHECK
        elif step.event == "exception":
            event_icon = "!"
        else:
            event_icon = ICON_LINE

        info.update(
            f" {ICON_STEP} [bold]Step {self.current + 1}[/bold] [dim]/ {len(self.steps)}[/dim]"
            f" {ICON_LINE} Line [bold]{step.line}[/bold]"
            f" [dim]│[/dim] {step.func_name}"
            f" [dim]│[/dim] {event_icon} {step.event}"
        )

        vars_bar = self.query_one("#tree_vars_bar", VarsBar)
        vars_bar.update_vars(step.locals_copy, step.globals_copy)

    def _scroll_to_current(self):
        try:
            container = self.query_one("#tree_container", ScrollableContainer)
            nodes = list(self.query(".tree-node"))
            if self.current >= len(nodes):
                return
            node = nodes[self.current]
            container.scroll_to_widget(
                node,
                animate=False,
                center=True,
                force=True,
                immediate=True,
            )
        except Exception:
            pass

    def action_step_forward(self):
        if self.current < len(self.steps) - 1:
            self.current += 1
            self._highlight_current()
            self._update_info()
            self.call_after_refresh(self._scroll_to_current)
            if self._on_step_change:
                self._on_step_change(self.current)

    def action_step_backward(self):
        if self.current > 0:
            self.current -= 1
            self._highlight_current()
            self._update_info()
            self.call_after_refresh(self._scroll_to_current)
            if self._on_step_change:
                self._on_step_change(self.current)

    def _highlight_current(self):
        nodes = list(self.query(".tree-node"))
        for i, n in enumerate(nodes):
            want = i == self.current
            has = n.has_class("current")
            if want == has:
                continue
            if want:
                n.add_class("current")
            else:
                n.remove_class("current")
            n.update(_build_node_text(self.steps[i], i, want, self.source_lines))

    def action_close(self):
        self.app.pop_screen()