#!/usr/bin/env python3
"""fzf-ai-highlight: highlight code blocks in session preview text.

Reads text from stdin, detects fenced code blocks (```lang ... ```),
applies Pygments syntax highlighting, and writes ANSI-highlighted output
to stdout. Everything outside code blocks is passed through unchanged.

Usage:
    fzf-ai-preview ... | fzf-ai-highlight
    cat session.txt | fzf-ai-highlight
"""

from __future__ import annotations

import re
import sys

# Lazy import Pygments — only pay the cost when there's actual code to highlight.
# We wrap in try/except so the previewer degrades gracefully if pygments is
# not installed (though it's now a hard dependency).
try:
    import pygments
    import pygments.formatters
    import pygments.lexers
    import pygments.util
    _HAS_PYGMENTS = True
except ImportError:
    _HAS_PYGMENTS = False

# Regex for fenced code blocks: ```lang ... ``` (with optional lang).
# The DOTALL flag lets us match across lines inside the block.
_BLOCK_RE = re.compile(
    r"```([\w+#.-]*)\s*\n(.*?)```",
    re.DOTALL,
)

# Inline code detection: `text` — rendered as bold for contrast.
_INLINE_RE = re.compile(r"`([^`]+)`")

# Safe ANSI reset
_RESET = "\033[0m"


def _guess_lexer(lang: str, code: str) -> pygments.lexers.Lexer | None:
    """Guess a Pygments lexer from the language tag or code content."""
    if not _HAS_PYGMENTS:
        return None
    lang = lang.strip().lower()
    if lang:
        try:
            return pygments.lexers.get_lexer_by_name(lang)
        except pygments.util.ClassNotFound:
            pass
    # Fallback: guess from content
    try:
        return pygments.lexers.guess_lexer(code)
    except pygments.util.ClassNotFound:
        return None


def highlight_code_block(code: str, lang: str = "") -> str:
    """Highlight a code block with Pygments, returning ANSI-escaped text."""
    if not _HAS_PYGMENTS or not code.strip():
        return code

    lexer = _guess_lexer(lang, code)
    if lexer is None:
        return code

    try:
        formatter = pygments.formatters.Terminal256Formatter(style="monokai")
        return pygments.highlight(code, lexer, formatter).rstrip("\n")
    except Exception:
        return code


def highlight_inline(text: str) -> str:
    """Apply ANSI to inline ``code`` spans (bold + dim)."""
    return _INLINE_RE.sub(r"\033[1m`\1`\033[0m", text)


def highlight_text(text: str) -> str:
    """Process text: highlight fenced code blocks and inline code.

    Non-code content is passed through with inline code highlighting.
    """
    if not _HAS_PYGMENTS:
        return text

    parts: list[str] = []
    last_end = 0

    for match in _BLOCK_RE.finditer(text):
        start = match.start()
        # Emit non-code text before this block (with inline highlighting)
        if start > last_end:
            parts.append(highlight_inline(text[last_end:start]))

        lang = match.group(1)
        code = match.group(2)
        # Decorate the opening fence with language info
        header = f"\033[2m```{lang}\033[0m" if lang else "\033[2m```\033[0m"
        highlighted = highlight_code_block(code, lang)
        parts.append(header)
        parts.append(highlighted)
        parts.append("\033[2m```\033[0m")

        last_end = match.end()

    # Emit remaining text after last block
    if last_end < len(text):
        parts.append(highlight_inline(text[last_end:]))

    return "\n".join(parts) if parts else text


def main() -> int:
    text = sys.stdin.read()
    if not text:
        return 0
    try:
        sys.stdout.write(highlight_text(text))
        sys.stdout.write("\n")
    except BrokenPipeError:
        pass
    return 0


if __name__ == "__main__":
    sys.exit(main())
