alnn_rewrite/util/ui_logger.py

import logging
import threading
import time

from textual.app import App
from textual.widgets import RichLog


class PipelineLogger:
    """
    Manages centralized logging for both a standard system log file
    and a Textual RichLog UI widget. Automatically handles thread-safety
    and task-specific context tracking.
    """

    def __init__(
        self,
        app: App,
        widget_id: str = "#live_log",
        log_file: str = "alnn_rewrite.log",
        source: str = "SYSTEM",
        file_logger: logging.Logger | None = None,
    ):
        self.app = app
        self.widget_id = widget_id
        self.source = source

        # Share the underlying file logger if spawning a sub-logger
        if file_logger is None:
            logging.basicConfig(
                filename=log_file,
                level=logging.INFO,
                format="%(asctime)s - %(levelname)s - %(message)s",
            )
            self.file_logger = logging.getLogger("CNNHarness")
        else:
            self.file_logger = file_logger

    def get_task_logger(self, task_name: str) -> "PipelineLogger":
        """Spawns a new instance of this logger hardcoded to a specific task's source name."""
        return PipelineLogger(
            app=self.app,
            widget_id=self.widget_id,
            source=task_name,
            file_logger=self.file_logger,
        )

    def write(self, message: str, error: bool = False) -> None:
        """Core write method that handles both file and UI logging intelligently."""

        # 1. System file logging
        log_msg = f"[{self.source}] {message}"
        if error:
            self.file_logger.error(log_msg)
        else:
            self.file_logger.info(log_msg)

        # 2. UI logging (Conditional Thread-Safe Execution)
        def _do_write():
            try:
                timestamp = time.strftime("%H:%M:%S")
                color = (
                    "red"
                    if error
                    else ("magenta" if self.source == "SYSTEM" else "cyan")
                )
                prefix = f"[dim]{timestamp}[/dim] [[bold {color}]{self.source}[/bold {color}]]"

                log_widget = self.app.query_one(self.widget_id, RichLog)
                log_widget.write(f"{prefix} {message}")
            except Exception:
                pass

        # Textual throws an error if call_from_thread is used on the main thread.
        # This safely executes _do_write() based on where it's being called from.
        if threading.current_thread() is threading.main_thread():
            _do_write()
        else:
            self.app.call_from_thread(_do_write)

    def info(self, message: str) -> None:
        """Helper for standard informational logging."""
        self.write(message, error=False)

    def error(self, message: str) -> None:
        """Helper for error logging."""
        self.write(message, error=True)