cowork-local/presentation/folder/code_editor.py

"""Ô soạn mã có đánh số dòng và tô màu cú pháp — R08-T12.

Dùng cho cả xem lẫn sửa file văn bản. Tô màu qua Pygments nếu có; không
có thì vẫn soạn được, chỉ mất màu.
"""
from __future__ import annotations

from .file_helpers import (
    _MAX_HIGHLIGHT_CHARS, _fmt,
)

import os
from PySide6.QtCore import QRect, QSize, Qt, QTimer
from PySide6.QtGui import QColor, QFont, QPainter, QSyntaxHighlighter, QTextCharFormat
from PySide6.QtWidgets import QPlainTextEdit, QWidget
from ...i18n import tr
from ...theme import current_palette


class PygmentsHighlighter(QSyntaxHighlighter):
    """Colour the whole document with Pygments and apply per-block. Re-lexes the
    full text (debounced) so multi-line strings/comments colour correctly."""

    def __init__(self, document):
        super().__init__(document)
        from pygments.lexers.special import TextLexer
        self._lexer = TextLexer(stripnl=False)
        self._ranges: list[tuple[int, int, QTextCharFormat]] = []
        self._rules = self._build_rules()
        self._timer = QTimer(self)
        self._timer.setSingleShot(True)
        self._timer.setInterval(250)
        self._timer.timeout.connect(self._retokenize)
        document.contentsChanged.connect(self._timer.start)

    @staticmethod
    def _build_rules():
        from pygments.token import (
            Comment, Error, Keyword, Name, Number, Operator, Punctuation, String,
        )
        p = current_palette()
        # Ordered specific → general: first matching token type wins.
        # Colours are resolved when the editor is built, so reopening a file
        # after a theme switch re-highlights it in the new theme.
        return [
            (Comment, _fmt(p.code_comment, italic=True)),
            (Keyword.Type, _fmt(p.code_type)),
            (Keyword, _fmt(p.code_keyword)),
            (Name.Function, _fmt(p.code_func)),
            (Name.Class, _fmt(p.code_type)),
            (Name.Decorator, _fmt(p.code_func)),
            (Name.Builtin, _fmt(p.code_type)),
            (Name.Tag, _fmt(p.code_keyword)),
            (Name.Attribute, _fmt(p.code_attr)),
            (String.Doc, _fmt(p.code_comment, italic=True)),
            (String, _fmt(p.code_string)),
            (Number, _fmt(p.code_number)),
            (Operator, _fmt(p.code_fg)),
            (Punctuation, _fmt(p.code_fg)),
            (Error, _fmt(p.code_error)),
        ]

    def set_filename(self, filename: str, text: str = "") -> None:
        from pygments.lexers import get_lexer_for_filename, guess_lexer
        from pygments.lexers.special import TextLexer
        from pygments.util import ClassNotFound
        try:
            self._lexer = get_lexer_for_filename(filename, stripnl=False)
        except ClassNotFound:
            try:
                self._lexer = guess_lexer(text) if text.strip() else TextLexer()
            except ClassNotFound:
                self._lexer = TextLexer(stripnl=False)
        self._retokenize()

    def _fmt_for(self, tok):
        for ttype, fmt in self._rules:
            if tok in ttype:
                return fmt
        return None

    def _retokenize(self) -> None:
        from pygments import lex
        text = self.document().toPlainText()
        self._ranges = []
        if len(text) <= _MAX_HIGHLIGHT_CHARS:
            pos = 0
            for tok, val in lex(text, self._lexer):
                fmt = self._fmt_for(tok)
                if fmt is not None and val:
                    self._ranges.append((pos, pos + len(val), fmt))
                pos += len(val)
        self.rehighlight()

    def highlightBlock(self, text: str) -> None:  # noqa: N802 - Qt override
        if not self._ranges:
            return
        bstart = self.currentBlock().position()
        bend = bstart + len(text)
        for start, end, fmt in self._ranges:
            if end <= bstart or start >= bend:
                continue
            s = max(start, bstart) - bstart
            e = min(end, bend) - bstart
            if e > s:
                self.setFormat(s, e - s, fmt)


class _LineNumbers(QWidget):
    def __init__(self, editor):
        super().__init__(editor)
        self._editor = editor

    def sizeHint(self) -> QSize:
        return QSize(self._editor.line_number_width(), 0)

    def paintEvent(self, event):  # noqa: N802
        self._editor.paint_line_numbers(event)


class CodeEditor(QPlainTextEdit):
    """A dark, monospaced editor with a line-number gutter + Pygments colouring —
    the Sublime/VS-Code look for viewing & editing source files."""

    def __init__(self):
        super().__init__()
        self.setObjectName("codeEditor")
        self.setLineWrapMode(QPlainTextEdit.NoWrap)
        self.setTabStopDistance(4 * self.fontMetrics().horizontalAdvance(" "))
        font = QFont("Consolas")
        font.setStyleHint(QFont.Monospace)
        font.setPointSize(10)
        self.setFont(font)
        # Surface comes from the central style sheet (#codeEditor) — see theme.py.
        self._gutter = _LineNumbers(self)
        self.blockCountChanged.connect(lambda _=0: self._update_gutter_width())
        self.updateRequest.connect(self._on_update_request)
        self._highlighter = PygmentsHighlighter(self.document())
        self._update_gutter_width()

    # ---- line-number gutter -------------------------------------------------
    def line_number_width(self) -> int:
        digits = max(2, len(str(max(1, self.blockCount()))))
        return 12 + self.fontMetrics().horizontalAdvance("9") * digits

    def _update_gutter_width(self) -> None:
        self.setViewportMargins(self.line_number_width(), 0, 0, 0)

    def _on_update_request(self, rect, dy: int) -> None:
        if dy:
            self._gutter.scroll(0, dy)
        else:
            self._gutter.update(0, rect.y(), self._gutter.width(), rect.height())
        if rect.contains(self.viewport().rect()):
            self._update_gutter_width()

    def resizeEvent(self, event):  # noqa: N802
        super().resizeEvent(event)
        cr = self.contentsRect()
        self._gutter.setGeometry(QRect(cr.left(), cr.top(), self.line_number_width(), cr.height()))

    def paint_line_numbers(self, event) -> None:
        p = current_palette()
        painter = QPainter(self._gutter)
        painter.fillRect(event.rect(), QColor(p.code_gutter_bg))
        block = self.firstVisibleBlock()
        num = block.blockNumber()
        top = self.blockBoundingGeometry(block).translated(self.contentOffset()).top()
        bottom = top + self.blockBoundingRect(block).height()
        painter.setPen(QColor(p.code_gutter_fg))
        while block.isValid() and top <= event.rect().bottom():
            if block.isVisible() and bottom >= event.rect().top():
                painter.drawText(0, int(top), self._gutter.width() - 6,
                                 self.fontMetrics().height(), Qt.AlignRight,
                                 str(num + 1))
            block = block.next()
            top = bottom
            bottom = top + self.blockBoundingRect(block).height()
            num += 1

    def load_file(self, path: str, text: str) -> None:
        self.setPlainText(text)
        self._highlighter.set_filename(path, text)