trace.py

#!/usr/bin/env python3
"""
Trace — a terminal code explorer.

Point it at a folder. It shows you the files, the functions and classes
inside them, and lets you peek at any definition without opening an editor.
Works over SSH, in Codespaces, in Docker — anywhere with a terminal.
"""

from __future__ import annotations

import sys
from dataclasses import dataclass, field
from pathlib import Path

from textual.app import App, ComposeResult, Screen
from textual.containers import Horizontal, Vertical, Container
from textual.widgets import Footer, Header, Static, Tree, ListView, ListItem, Label, Input
from textual.widgets.tree import TreeNode
from textual.reactive import reactive
from textual.binding import Binding
from rich.syntax import Syntax
from rich.text import Text

from tree_sitter_language_pack import get_parser

EXT_TO_LANG = {
    ".py": "python",
    ".js": "javascript",
    ".jsx": "javascript",
    ".mjs": "javascript",
    ".ts": "javascript",
    ".rs": "rust",
    ".go": "go",
}

DEFINITION_NODE_TYPES = {
    "python": {
        "function_definition": "function",
        "class_definition": "class",
    },
    "javascript": {
        "function_declaration": "function",
        "method_definition": "method",
        "class_declaration": "class",
        "arrow_function": "function",
    },
    "rust": {
        "function_item": "function",
        "struct_item": "struct",
        "trait_item": "trait",
    },
    "go": {
        "function_declaration": "function",
        "method_declaration": "method",
        "type_declaration": "type",
    },
}

CALL_NODE_TYPES = {
    "python": {"call"},
    "javascript": {"call_expression"},
    "rust": {"call_expression"},
    "go": {"call_expression"},
}

NAME_FIELD = {
    "python": "name",
    "javascript": "name",
    "rust": "name",
    "go": "name",
}


@dataclass
class Symbol:
    name: str
    kind: str
    file: Path
    start_line: int
    end_line: int
    lang: str


@dataclass
class CallSite:
    name: str
    file: Path
    line: int


@dataclass
class FileIndex:
    path: Path
    lang: str | None
    symbols: list[Symbol] = field(default_factory=list)
    calls: list[CallSite] = field(default_factory=list)
    source: str = ""
    error: str | None = None


class Indexer:
    """Walks a directory, parses every recognized source file with tree-sitter,
    and builds a symbol table plus a call graph."""

    IGNORE_DIRS = {
        ".git", "node_modules", "__pycache__", ".venv", "venv",
        "dist", "build", ".mypy_cache", ".pytest_cache", "target",
    }

    def __init__(self, root: Path):
        self.root = root
        self.files: dict[Path, FileIndex] = {}
        self.symbol_table: dict[str, list[Symbol]] = {}

    def build(self) -> None:
        for path in sorted(self.root.rglob("*")):
            if not path.is_file():
                continue
            if any(part in self.IGNORE_DIRS for part in path.parts):
                continue
            lang = EXT_TO_LANG.get(path.suffix)
            fi = self._index_file(path, lang)
            self.files[path] = fi
            for sym in fi.symbols:
                self.symbol_table.setdefault(sym.name, []).append(sym)

    def _index_file(self, path: Path, lang: str | None) -> FileIndex:
        try:
            source = path.read_text(encoding="utf-8", errors="replace")
        except Exception as e:
            return FileIndex(path=path, lang=lang, error=str(e))

        fi = FileIndex(path=path, lang=lang, source=source)
        if lang is None:
            return fi

        try:
            parser = get_parser(lang)
            tree = parser.parse(source.encode("utf-8"))
        except Exception as e:
            fi.error = f"parse error: {e}"
            return fi

        def_types = DEFINITION_NODE_TYPES.get(lang, {})
        call_types = CALL_NODE_TYPES.get(lang, set())
        name_field = NAME_FIELD.get(lang, "name")

        def walk(node, in_impl=False):
            if node.type == "impl_item":
                for child in node.children:
                    walk(child, in_impl=True)
                return

            if node.type == "arrow_function" and node.parent is not None and node.parent.type == "variable_declarator":
                name_node = node.parent.child_by_field_name("name")
                name = name_node.text.decode("utf-8") if name_node else "<anonymous>"
                fi.symbols.append(Symbol(
                    name=name,
                    kind="function",
                    file=path,
                    start_line=node.start_point[0] + 1,
                    end_line=node.end_point[0] + 1,
                    lang=lang,
                ))
            elif node.type in def_types:
                name_node = node.child_by_field_name(name_field)
                if name_node is None and node.type == "type_declaration":
                    for child in node.children:
                        if child.type == "type_spec":
                            name_node = child.child_by_field_name("name")
                            break
                name = name_node.text.decode("utf-8") if name_node else "<anonymous>"
                kind = def_types[node.type]
                if lang == "rust" and node.type == "function_item" and in_impl:
                    kind = "method"
                fi.symbols.append(Symbol(
                    name=name,
                    kind=kind,
                    file=path,
                    start_line=node.start_point[0] + 1,
                    end_line=node.end_point[0] + 1,
                    lang=lang,
                ))
            elif node.type in call_types:
                for child in node.children:
                    if child.type in ("identifier", "field_identifier"):
                        fi.calls.append(CallSite(
                            name=child.text.decode("utf-8"),
                            file=path,
                            line=node.start_point[0] + 1,
                        ))
                        break
            for child in node.children:
                walk(child, in_impl=in_impl)

        walk(tree.root_node)
        return fi

    def callers_of(self, symbol_name: str, exclude_file: Path) -> list[CallSite]:
        result = []
        for fi in self.files.values():
            if fi.path == exclude_file:
                continue
            for c in fi.calls:
                if c.name == symbol_name:
                    result.append(c)
        return result


ACCENT = "#7aa2f7"
DIM = "#565f89"

KIND_GLYPH = {
    "function": "\u0192",
    "method": "\u2022",
    "class": "\u25c6",
    "struct": "\u25c6",
    "trait": "\u25c7",
    "impl": "\u25b8",
    "type": "\u25b8",
}


class SymbolItem(ListItem):
    def __init__(self, symbol: Symbol):
        self.symbol = symbol
        glyph = KIND_GLYPH.get(symbol.kind, "\u00b7")
        label = Text()
        label.append(f" {glyph} ", style=ACCENT)
        label.append(symbol.name, style="bold")
        label.append(f"   {symbol.kind}", style=DIM)
        super().__init__(Label(label))


class SearchModal(Screen):
    """Modal for searching symbols in the current file."""

    CSS = """
    Screen {
        align: center middle;
    }
    #search_container {
        width: 60;
        height: auto;
        border: solid #7aa2f7;
        background: #16161e;
        padding: 1 2;
    }
    #search_input {
        width: 1fr;
        margin-bottom: 1;
    }
    #search_results {
        width: 1fr;
        height: auto;
        max-height: 15;
        border: solid #2a2e3f;
        background: #1a1b26;
    }
    #search_status {
        width: 1fr;
        margin-top: 1;
        color: #7aa2f7;
    }
    """

    BINDINGS = [
        Binding("escape", "close", "Close", show=False),
        Binding("enter", "next_match", "Next", show=False),
        Binding("n", "next_match", "Next Match", show=False),
        Binding("shift+n", "prev_match", "Prev Match", show=False),
    ]

    def __init__(self, symbols: list[Symbol]):
        super().__init__()
        self.all_symbols = symbols
        self.matches = []
        self.current_match_index = 0

    def compose(self) -> ComposeResult:
        with Container(id="search_container"):
            yield Static("Search symbols (Esc to close):", id="search_label")
            yield Input(id="search_input", placeholder="Type to filter...")
            yield ListView(id="search_results")
            yield Static("", id="search_status")

    def on_mount(self) -> None:
        """Focus the search input on open."""
        self.query_one("#search_input", Input).focus()
        self._update_results("")

    def on_input_changed(self, event: Input.Changed) -> None:
        """Filter symbols as user types."""
        self._update_results(event.value.lower())

    def _update_results(self, query: str) -> None:
        """Filter symbols and update the results list."""
        # Filter symbols by query (case-insensitive substring match)
        if not query:
            self.matches = self.all_symbols
        else:
            self.matches = [s for s in self.all_symbols if query in s.name.lower()]

        # Update results listview
        results_list = self.query_one("#search_results", ListView)
        results_list.clear()

        if not self.matches:
            results_list.append(ListItem(Label("[dim]No matches[/dim]")))
            self.query_one("#search_status", Static).update("")
            return

        for sym in self.matches:
            results_list.append(SymbolItem(sym))

        # Update status with match count
        query_text = self.query_one("#search_input", Input).value
        status = f"{len(self.matches)}/{len(self.all_symbols)} matches"
        if query_text:
            status += f" for '{query_text}'"
        self.query_one("#search_status", Static).update(status)
        self.current_match_index = 0

    def action_next_match(self) -> None:
        """Cycle to next match."""
        if not self.matches:
            return
        self.current_match_index = (self.current_match_index + 1) % len(self.matches)
        self._highlight_current()

    def action_prev_match(self) -> None:
        """Cycle to previous match."""
        if not self.matches:
            return
        self.current_match_index = (self.current_match_index - 1) % len(self.matches)
        self._highlight_current()

    def _highlight_current(self) -> None:
        """Highlight the current match in the results."""
        results_list = self.query_one("#search_results", ListView)
        if self.matches and self.current_match_index < len(results_list.children):
            results_list.index = self.current_match_index

    def action_close(self) -> None:
        """Close the search modal and return the selected symbol to the app."""
        if self.matches and self.current_match_index < len(self.matches):
            selected = self.matches[self.current_match_index]
            self.dismiss(selected)
        else:
            self.dismiss(None)


class HelpModal(Screen):
    """Modal showing all keybindings, navigation help, and symbol legend."""

    CSS = """
    Screen {
        align: center middle;
    }
    #help_container {
        width: 70;
        height: auto;
        max-height: 24;
        border: solid #7aa2f7;
        background: #16161e;
        padding: 1 2;
        overflow: auto;
    }
    #help_content {
        width: 1fr;
    }
    """

    BINDINGS = [
        Binding("escape", "close", "Close", show=False),
        Binding("q", "close", "Close", show=False),
    ]

    def compose(self) -> ComposeResult:
        with Container(id="help_container"):
            yield Static(id="help_content")

    def on_mount(self) -> None:
        """Generate and display help content."""
        help_text = self._generate_help_text()
        self.query_one("#help_content", Static).update(help_text)

    def _generate_help_text(self) -> str:
        """Generate formatted help text with keybindings and legend."""
        text = Text()

        # Title
        text.append("TRACE — Code Explorer Help\n", style="bold #7aa2f7")
        text.append("=" * 60 + "\n\n", style="dim")

        # Navigation Section
        text.append("NAVIGATION\n", style="bold #c0caf5")
        text.append("─" * 60 + "\n", style="dim")
        nav_help = [
            ("Arrow keys / hjkl", "Move through file tree or symbol list"),
            ("Tab / Shift+Tab", "Switch between panels (Files ← → Symbols ← → Preview)"),
            ("1, 2, 3", "Jump directly to Files, Symbols, or Preview panel"),
        ]
        for key, desc in nav_help:
            text.append(f"  {key:20} {desc}\n")
        text.append("\n")

        # Selection & Viewing Section
        text.append("SELECTION & VIEWING\n", style="bold #c0caf5")
        text.append("─" * 60 + "\n", style="dim")
        select_help = [
            ("Enter", "Select symbol and view in preview panel"),
            ("/", "Search symbols in current file (n/N to cycle)"),
            ("Escape", "Close search or any modal"),
            ("c", "Show all callers of selected symbol"),
            ("r", "Refresh workspace index (re-parse all files)"),
        ]
        for key, desc in select_help:
            text.append(f"  {key:20} {desc}\n")
        text.append("\n")

        # Legend Section
        text.append("SYMBOL LEGEND\n", style="bold #c0caf5")
        text.append("─" * 60 + "\n", style="dim")
        legend_items = [
            ("ƒ", "function", "A function or procedure"),
            ("•", "method", "A method (function in a class/struct)"),
            ("◆", "class", "A class definition"),
            ("◇", "trait", "A trait or interface"),
            ("▸", "impl", "An implementation block"),
            ("type", "type", "A type alias or definition"),
        ]
        for glyph, kind, desc in legend_items:
            text.append(f"  {glyph:3} {kind:12} {desc}\n")
        text.append("\n")

        # Status Bar Info
        text.append("STATUS BAR\n", style="bold #c0caf5")
        text.append("─" * 60 + "\n", style="dim")
        text.append("  Left:  Current file, symbols, and call information\n")
        text.append("  Right: Shortcut to quit (q)\n")
        text.append("\n")

        # Exit Section
        text.append("EXIT\n", style="bold #c0caf5")
        text.append("─" * 60 + "\n", style="dim")
        text.append("  q            Quit Trace\n")
        text.append("  Escape       Close this help (or any modal)\n")

        return text

    def action_close(self) -> None:
        """Close the help modal."""
        self.dismiss()


class TraceApp(App):
    CSS = """
    Screen {
        background: #1a1b26;
    }
    #file_tree {
        width: 32%;
        border-right: solid #2a2e3f;
        background: #16161e;
    }
    #symbol_panel {
        width: 26%;
        border-right: solid #2a2e3f;
        background: #1a1b26;
    }
    #symbol_panel > ListView {
        background: #1a1b26;
    }
    #code_panel {
        width: 42%;
        background: #16161e;
    }
    #status_bar {
        height: 1;
        background: #24283b;
    }
    #status_text {
        width: 1fr;
        color: #7aa2f7;
        padding: 0 1;
    }
    #panel_title {
        height: 1;
        background: #24283b;
        color: #c0caf5;
        padding: 0 1;
        text-style: bold;
    }
    Tree {
        background: #16161e;
    }
    ListView {
        border: none;
    }
    ListItem {
        padding: 0 1;
    }
    ListItem:hover {
        background: #292e42;
    }
    ListItem.--highlight {
        background: #364a82;
    }
    #empty_state {
        color: #565f89;
        padding: 2 2;
        text-align: center;
    }
    """

    BINDINGS = [
        # Navigation: arrows move within panels, Tab switches between panels
        Binding("q", "quit", "Quit", priority=True),
        Binding("escape", "close_modal", "Close", show=False),
        Binding("?", "show_help", "Help"),
        Binding("tab", "focus_next", "Next Panel", show=False),
        Binding("shift+tab", "focus_previous", "Prev Panel", show=False),
        Binding("1", "focus_tree", "Files", show=False),
        Binding("2", "focus_symbols", "Symbols", show=False),
        Binding("3", "focus_code", "Preview", show=False),
        # Actions
        Binding("r", "refresh", "Refresh"),
        Binding("/", "search_symbols", "Search"),
        Binding("c", "show_callers", "Who calls this"),
    ]

    TITLE = "Trace"

    current_file: reactive[Path | None] = reactive(None)
    current_symbol: reactive[Symbol | None] = reactive(None)

    def __init__(self, root: Path):
        super().__init__()
        self.root_path = root
        self.indexer = Indexer(root)

    def compose(self) -> ComposeResult:
        yield Header(show_clock=False)
        with Horizontal():
            with Vertical(id="file_tree"):
                yield Static(" FILES", id="panel_title")
                yield Tree(str(self.root_path.name), id="tree")
            with Vertical(id="symbol_panel"):
                yield Static(" SYMBOLS", id="panel_title")
                yield ListView(id="symbols")
            with Vertical(id="code_panel"):
                yield Static(" PREVIEW", id="panel_title")
                yield Static(id="code", expand=True)
        with Horizontal(id="status_bar"):
            yield Static("", id="status_text")
        yield Footer()

    def on_mount(self) -> None:
        self.set_status("Indexing workspace…")
        self.indexer.build()
        self._populate_tree()
        n_files = len(self.indexer.files)
        n_syms = sum(len(f.symbols) for f in self.indexer.files.values())
        self.set_status(f"{n_files} files \u00b7 {n_syms} symbols indexed")
        self.query_one("#code", Static).update(
            "[dim]Pick a file on the left, then a symbol in the middle to preview it here.\n"
            "Press [bold]?[/bold] for help and keybindings.[/dim]"
        )

    def set_status(self, text: str) -> None:
        self.query_one("#status_text", Static).update(text)

    def _populate_tree(self) -> None:
        tree = self.query_one("#tree", Tree)
        tree.root.expand()
        self._add_dir(tree.root, self.root_path)

    def _add_dir(self, node: TreeNode, path: Path) -> None:
        try:
            entries = sorted(
                path.iterdir(),
                key=lambda p: (p.is_file(), p.name.lower()),
            )
        except PermissionError:
            return
        for entry in entries:
            if entry.name in Indexer.IGNORE_DIRS or entry.name.startswith("."):
                continue
            if entry.is_dir():
                child = node.add(f"\U0001F4C1 {entry.name}", data=entry)
                self._add_dir(child, entry)
            else:
                # Check if file has a parse error and mark it
                fi = self.indexer.files.get(entry)
                if fi and fi.error:
                    node.add_leaf(f"\u26a0 {entry.name}", data=entry)
                else:
                    node.add_leaf(entry.name, data=entry)

    def on_tree_node_selected(self, event: Tree.NodeSelected) -> None:
        path = event.node.data
        if path is None or not path.is_file():
            return
        self.current_file = path
        self._populate_symbols(path)

    def _populate_symbols(self, path: Path) -> None:
        list_view = self.query_one("#symbols", ListView)
        list_view.clear()
        fi = self.indexer.files.get(path)
        if fi is None:
            return
        if fi.error:
            # Show error in preview panel instead of symbol list
            self._show_error(path, fi.error)
            return
        if not fi.symbols:
            list_view.append(ListItem(Label("[dim]No definitions found[/dim]")))
            self.query_one("#code", Static).update("[dim]No symbols to preview[/dim]")
            return
        for sym in fi.symbols:
            list_view.append(SymbolItem(sym))
        self.set_status(f"{path.name} \u00b7 {len(fi.symbols)} symbols")

    def _show_error(self, path: Path, error: str) -> None:
        """Display a parse error in the preview panel."""
        error_msg = (
            f"[red bold]⚠ Couldn't parse this file[/red bold]\n\n"
            f"[dim]{path.name}[/dim]\n\n"
            f"[red]{error}[/red]"
        )
        self.query_one("#code", Static).update(error_msg)
        self.set_status(f"⚠ {path.name} — parse error")

    def on_list_view_selected(self, event: ListView.Selected) -> None:
        item = event.item
        if isinstance(item, SymbolItem):
            self.current_symbol = item.symbol
            self._preview_symbol(item.symbol)

    def _preview_symbol(self, symbol: Symbol) -> None:
        fi = self.indexer.files.get(symbol.file)
        if fi is None:
            return
        lines = fi.source.splitlines()
        snippet = "\n".join(lines[symbol.start_line - 1: symbol.end_line])
        syntax = Syntax(
            snippet,
            symbol.lang or "text",
            theme="tokyo-night",
            line_numbers=True,
            start_line=symbol.start_line,
            word_wrap=False,
        )
        self.query_one("#code", Static).update(syntax)
        callers = self.indexer.callers_of(symbol.name, exclude_file=symbol.file)
        if callers:
            files = sorted({c.file.name for c in callers})
            self.set_status(
                f"{symbol.name} \u00b7 called from {len(callers)} place(s) in {', '.join(files[:3])}"
            )
        else:
            self.set_status(f"{symbol.name} \u00b7 no callers found elsewhere in the workspace")

    def action_focus_symbols(self) -> None:
        self.query_one("#symbols", ListView).focus()

    def action_show_callers(self) -> None:
        if not self.current_symbol:
            self.set_status("Select a symbol first")
            return
        sym = self.current_symbol
        callers = self.indexer.callers_of(sym.name, exclude_file=sym.file)
        if not callers:
            self.set_status(f"No other callers of '{sym.name}' found")
            return
        lines = [f"{c.file.relative_to(self.root_path)}:{c.line}" for c in callers]
        self.query_one("#code", Static).update(
            f"[bold]{sym.name}[/bold] is called from:\n\n" + "\n".join(f"  \u2192 {l}" for l in lines)
        )

    def action_close_modal(self) -> None:
        """Close any active modal. Safe no-op if none exists."""
        # Only pop if there's more than one screen on the stack (main + modal)
        if len(self.screen_stack) > 1:
            self.pop_screen()

    def action_search_symbols(self) -> None:
        """Open search modal for current file's symbols."""
        if not self.current_file:
            self.set_status("Select a file first")
            return

        fi = self.indexer.files.get(self.current_file)
        if not fi or not fi.symbols:
            self.set_status("No symbols to search")
            return

        def on_search_result(symbol: Symbol | None) -> None:
            """Handle search modal result."""
            if symbol:
                self.current_symbol = symbol
                self._preview_symbol(symbol)
            # If None or dismissed, just close and return to main view

        self.push_screen(SearchModal(fi.symbols), on_search_result)

    def action_show_help(self) -> None:
        """Open help modal with all keybindings and legend."""
        self.push_screen(HelpModal())

    def action_focus_next(self) -> None:
        """Focus next panel (Tab). Preview panel is display-only, not included in cycle."""
        focused_widget = self.focused
        if focused_widget and focused_widget.id == "tree":
            self.query_one("#symbols", ListView).focus()
        else:
            self.query_one("#tree", Tree).focus()

    def action_focus_previous(self) -> None:
        """Focus previous panel (Shift+Tab). Preview panel is display-only, not included in cycle."""
        focused_widget = self.focused
        if focused_widget and focused_widget.id == "tree":
            self.query_one("#symbols", ListView).focus()
        else:
            self.query_one("#tree", Tree).focus()

    def action_focus_tree(self) -> None:
        """Jump to FILES panel (1 key)."""
        self.query_one("#tree", Tree).focus()

    def action_focus_code(self) -> None:
        """Jump to PREVIEW panel (3 key)."""
        # Preview panel (#code) is display-only, not keyboard-focusable
        # Just update status to indicate the user is viewing the preview
        self.set_status("Viewing preview (read-only)")

    def action_refresh(self) -> None:
        """Refresh the workspace index (r key)."""
        self.set_status("Refreshing workspace…")
        self.indexer.build()
        n_files = len(self.indexer.files)
        n_syms = sum(len(f.symbols) for f in self.indexer.files.values())
        self.set_status(f"Refreshed: {n_files} files · {n_syms} symbols indexed")


def main():
    target = Path(sys.argv[1]).resolve() if len(sys.argv) > 1 else Path.cwd()
    if not target.is_dir():
        print(f"Not a directory: {target}")
        sys.exit(1)
    TraceApp(target).run()


if __name__ == "__main__":
    main()