plain2code_console.py

import logging
import os

from rich.console import Console
from rich.style import Style
from rich.tree import Tree

import plain2code_logger

CHARACTERS_TO_TOKENS_RULE_OF_THUMB_RATIO = 4

logger = logging.getLogger(plain2code_logger.LOGGER_NAME)

# Colors for log messages, applied by the terminal and the TUI via the "color"
# parameter of the console methods. The file log always receives plain text.
RETRY_COLOR = "#FFB454"  # Amber
SUCCESS_COLOR = "#79FC96"  # Green
MUTED_COLOR = "#888888"  # Grey


class Plain2CodeConsole(Console):
    INFO_STYLE = Style()
    WARNING_STYLE = Style(color="yellow", bold=True)
    ERROR_STYLE = Style(color="red", bold=True)
    INPUT_STYLE = Style(color="#4169E1")  # Royal Blue
    OUTPUT_STYLE = Style(color="green")
    DEBUG_STYLE = Style(color="purple")

    def __init__(self):
        super().__init__()
        try:
            import tiktoken

            self.llm_encoding = tiktoken.get_encoding("cl100k_base")
        except Exception as e:
            logger.warning(
                "Failed to import optional library tiktoken. Using approximate instead of exact token count."
            )
            logger.debug(f"Exception: {e}")
            self.llm_encoding = None

    def info(self, *args, color=None, **kwargs):
        self._log_and_print(logging.INFO, self.INFO_STYLE, args, color, kwargs)

    def warning(self, *args, color=None, **kwargs):
        self._log_and_print(logging.WARNING, self.WARNING_STYLE, args, color, kwargs)

    def error(self, *args, color=None, **kwargs):
        self._log_and_print(logging.ERROR, self.ERROR_STYLE, args, color, kwargs)

    def input(self, *args, color=None, **kwargs):
        # We also log input as info so it shows in the toggled view
        self._log_and_print(logging.INFO, self.INPUT_STYLE, args, color, kwargs)

    def output(self, *args, color=None, **kwargs):
        self._log_and_print(logging.INFO, self.OUTPUT_STYLE, args, color, kwargs)

    def debug(self, *args, color=None, **kwargs):
        self._log_and_print(logging.DEBUG, self.DEBUG_STYLE, args, color, kwargs)

    def _log_and_print(self, level, base_style, args, color, kwargs):
        """Log the plain message text, then print it styled to the terminal.

        The optional color is applied by the terminal (via style) and forwarded to
        the TUI as the "log_color" record attribute; the file log stays plain text.
        """
        logger.log(level, " ".join(map(str, args)), extra={"log_color": color})
        style = base_style + Style(color=color) if color else base_style
        # Log messages must render exactly as logged: don't interpret square brackets
        # in interpolated content (error texts, file names) as Rich markup.
        kwargs.setdefault("markup", False)
        super().print(*args, **kwargs, style=style)

    def print_list(self, items, style=None):
        for item in items:
            logger.debug(f"  {item}")
            super().print(f"{item}", style=style)

    def print_files(self, header, root_folder, files, style=None):
        if not files:
            return

        logger.debug(f"{header} {', '.join(files.keys())}")
        tree = self._create_tree_from_files(root_folder, files)
        super().print(f"\n{header}", style=style)

        super().print(tree, style=style)

        super().print()

    def _create_tree_from_files(self, root_folder, files):
        """
        Creates a Tree structure from a dictionary of files using the rich library.

        Args:
            files (dict): A dictionary where keys are file paths (strings)
                            and values are file content (strings).

        Returns:
            Tree: The root of the created tree structure.
        """
        tree = Tree(root_folder)
        for path, content in files.items():
            parts = path.split(os.sep)
            current_level = tree
            for part in parts:
                existing_level = None
                for child in current_level.children:
                    if child.label == part:
                        existing_level = child
                        break

                if existing_level is None:
                    if part == parts[-1]:
                        if files[path] is None:
                            current_level = current_level.add(f"{part} [red]deleted[/red]")
                        else:
                            file_lines = len(content.splitlines())
                            file_tokens = self._count_tokens(content)
                            current_level = current_level.add(f"{part} ({file_lines} lines, {file_tokens} tokens)")
                    else:
                        current_level = current_level.add(part)
                else:
                    current_level = existing_level

        return tree

    def _count_tokens(self, text):
        """Count tokens using tiktoken if available, otherwise estimate from character count."""
        if self.llm_encoding is not None:
            try:
                return len(self.llm_encoding.encode(text))
            except Exception:
                pass
        return len(text) // CHARACTERS_TO_TOKENS_RULE_OF_THUMB_RATIO

    def print_resources(self, resources_list, linked_resources):
        if len(resources_list) == 0:
            self.debug("Linked resources: None")
            return

        self.debug("Linked resources:")
        for resource_name in resources_list:
            if resource_name["target"] in linked_resources:
                file_tokens = self._count_tokens(linked_resources[resource_name["target"]])
                self.debug(f"- {resource_name['text']} ({resource_name['target']}, {file_tokens} tokens)")

        self.input()


console = Plain2CodeConsole()