codeview.py

import pygments
from pygments.lexers import PythonLexer
from pygments.token import Token

from rich.style import Style as RichStyle
from rich.text import Text
from rich.cells import cell_len

from textual.app import ComposeResult
from textual.containers import ScrollableContainer
from textual.widgets import Static

from themes import ORANGE_SYNTAX, PURPLE_SYNTAX, ANSI_SYNTAX


# Map Pygments token types to palette keys
TOKEN_TO_PALETTE_KEY = {
    Token.Keyword: "keyword",
    Token.Keyword.Constant: "keyword",
    Token.Name: "default",
    Token.Name.Function: "function",
    Token.Name.Class: "class",
    Token.Name.Builtin: "builtin",
    Token.Name.Namespace: "default",
    Token.String: "string",
    Token.String.Doc: "string",
    Token.Number: "number",
    Token.Number.Integer: "number",
    Token.Number.Float: "number",
    Token.Operator: "operator",
    Token.Comment: "comment",
    Token.Punctuation: "default",
    Token.Keyword.Namespace: "keyword",
    Token.Literal: "string",
    Token.Error: "operator",
}


def _get_syntax_palette(app):
    """Get the syntax color palette for the current theme."""
    if app is None:
        return ORANGE_SYNTAX
    theme_name = getattr(app, "theme", "pytutor-orange")
    if theme_name == "pytutor-purple":
        return PURPLE_SYNTAX
    elif theme_name == "ansi":
        return ANSI_SYNTAX
    return ORANGE_SYNTAX


def _pygments_to_rich(style_str):
    parts = style_str.split()
    color = None
    bgcolor = None
    italic = False
    bold = False
    for p in parts:
        if p.startswith("#"):
            if color is None:
                color = p
            else:
                bgcolor = p
        elif p == "italic":
            italic = True
        elif p == "bold":
            bold = True
    kwargs = {}
    if color:
        kwargs["color"] = color
    if bgcolor:
        kwargs["bgcolor"] = bgcolor
    if italic:
        kwargs["italic"] = True
    if bold:
        kwargs["bold"] = True
    return RichStyle(**kwargs) if kwargs else None


def _build_style_map(syntax_palette):
    """Build a rich style map from a syntax palette."""
    mapping = {}
    for token, palette_key in TOKEN_TO_PALETTE_KEY.items():
        color = syntax_palette.get(palette_key)
        if color:
            rs = _pygments_to_rich(color)
            if rs:
                mapping[token] = rs
    return mapping


class CodeView(ScrollableContainer):
    """Syntax-highlighted, scrollable source viewer with a current-line band."""

    # Keep arrow keys for TutorScreen step bindings; mouse-wheel still scrolls.
    can_focus = False
    BINDINGS = []

    DEFAULT_CSS = """
    CodeView {
        height: 1fr;
        overflow-x: auto;
        overflow-y: auto;
        padding: 0 1;
    }
    .code-line {
        height: 1;
        width: auto;
        text-wrap: nowrap;
    }
    .code-line.current {
        text-style: bold;
    }
    """

    def __init__(self, source_lines, **kwargs):
        super().__init__(**kwargs)
        self.source_lines = source_lines or [""]
        self.current_line = 1
        self._style_map = _build_style_map(ORANGE_SYNTAX)
        # Content-space Y of the highlighted line (updated after layout).
        self._highlight_y = 0
        self._highlight_height = 1

    def compose(self) -> ComposeResult:
        for i in range(1, len(self.source_lines) + 1):
            yield Static("", id=f"code_line_{i}", classes="code-line")

    def _highlight_line_text(self, raw_line: str) -> Text:
        text = Text()
        if not raw_line:
            return text
        try:
            tokens = pygments.lex(raw_line, PythonLexer())
        except Exception:
            return Text(raw_line)
        for token, value in tokens:
            style = None
            t = token
            while t is not None:
                if t in self._style_map:
                    style = self._style_map[t]
                    break
                t = t.parent
            text.append(value, style=style)
        return text

    def _line_widget(self, line_no: int) -> Static:
        return self.query_one(f"#code_line_{line_no}", Static)

    def _paint_line(self, line_no: int) -> None:
        """Paint a single source line into its Static widget."""
        if line_no < 1 or line_no > len(self.source_lines):
            return

        palette = _get_syntax_palette(self.app)
        self._style_map = _build_style_map(palette)
        highlight_color = palette.get("keyword", "#e8751a")

        gutter_w = len(str(len(self.source_lines)))
        max_w = max((cell_len(l) for l in self.source_lines), default=0)
        raw = self.source_lines[line_no - 1]
        line_text = self._highlight_line_text(raw)
        pad = max(0, max_w - cell_len(raw))
        is_current = line_no == self.current_line

        if is_current:
            prefix = Text(
                f"{'▶':>{gutter_w}} ",
                style=RichStyle(bgcolor=highlight_color, color="#ffffff"),
            )
            line_text.stylize(RichStyle(bgcolor=highlight_color, color="#ffffff"))
        else:
            prefix = Text(f"{line_no:>{gutter_w}} ", style=RichStyle(color="#6e7681"))

        row = Text()
        row.append_text(prefix)
        row.append_text(line_text)
        if pad > 0:
            row.append(
                " " * pad,
                style=RichStyle(bgcolor=highlight_color if is_current else None),
            )

        widget = self._line_widget(line_no)
        widget.update(row)
        widget.set_class(is_current, "current")

    def _paint_all(self) -> None:
        for i in range(1, len(self.source_lines) + 1):
            self._paint_line(i)

    def on_mount(self) -> None:
        self._paint_all()
        self.call_after_refresh(self._ensure_line_visible)

    def highlight(self, line_no: int) -> None:
        prev = self.current_line
        self.current_line = max(1, min(line_no, len(self.source_lines)))
        if prev != self.current_line:
            self._paint_line(prev)
        self._paint_line(self.current_line)
        self.call_after_refresh(self._ensure_line_visible)

    def _sync_highlight_geometry(self) -> bool:
        """Read the highlighted line's content-space position after layout."""
        try:
            line = self._line_widget(self.current_line)
        except Exception:
            return False
        region = line.virtual_region
        self._highlight_y = region.y
        self._highlight_height = max(1, region.height)
        return True

    def _ensure_line_visible(self) -> None:
        """Keep the highlighted line in view, like an editor caret.

        Scrolls only when the highlight is outside the viewport. When scrolling
        down, places the line near the bottom with a small margin; when scrolling
        up, places it near the top — same behavior as typical code editors.
        """
        if not self._sync_highlight_geometry():
            return

        view_h = self.scrollable_content_region.height
        if view_h <= 0:
            return

        y = self._highlight_y
        h = self._highlight_height
        view_top = self.scroll_y
        view_bottom = view_top + view_h

        # Already fully visible — do nothing (preserve manual mouse-wheel position).
        if y >= view_top and (y + h) <= view_bottom:
            return

        margin = 1
        if y < view_top:
            # Highlight above viewport → scroll up so it sits near the top.
            target = max(0, y - margin)
        else:
            # Highlight below viewport → scroll down so it sits near the bottom.
            target = max(0, y + h - view_h + margin)

        self.scroll_to(y=target, animate=False, force=True, immediate=True)
        # Re-sync after scroll in case layout rounded differently.
        self._sync_highlight_geometry()