model/log.py

"""
Shared structured logger for the Luminary HF backend.
Outputs readable, sectioned logs that are easy to scan in the HF Space container view.
"""

from __future__ import annotations
import logging
import sys
import time
from datetime import datetime, timezone

# ── ANSI colour palette ──────────────────────────────────────────────────────
_R  = "\033[0m"          # reset
_B  = "\033[1m"          # bold
_DIM = "\033[2m"         # dim
_GRN = "\033[32m"        # green
_CYN = "\033[36m"        # cyan
_YLW = "\033[33m"        # yellow
_RED = "\033[31m"        # red
_MAG = "\033[35m"        # magenta
_BLU = "\033[34m"        # blue
_WHT = "\033[97m"        # bright white

_TAG_COLORS: dict[str, str] = {
    "STARTUP":  _CYN,
    "MODEL":    _MAG,
    "UPLOAD":   _BLU,
    "PROCESS":  _BLU,
    "CHAT":     _GRN,
    "QUIZ":     _GRN,
    "SWITCH":   _YLW,
    "ERROR":    _RED,
    "HEALTH":   _DIM,
}


class _FmtHandler(logging.StreamHandler):
    """Formatter that wraps log records into readable tag-prefixed lines."""

    def emit(self, record: logging.LogRecord) -> None:
        try:
            tag   = getattr(record, "tag", record.levelname)
            msg   = record.getMessage()
            color = _TAG_COLORS.get(tag, _WHT)
            ts    = datetime.now(timezone.utc).strftime("%H:%M:%S")
            prefix = f"{_DIM}{ts}{_R} {color}{_B}[{tag}]{_R}"
            # indent continuation lines
            lines  = msg.splitlines()
            out    = prefix + "  " + lines[0]
            for line in lines[1:]:
                out += "\n" + (" " * (len(ts) + len(tag) + 5)) + line
            sys.stdout.write(out + "\n")
            sys.stdout.flush()
        except Exception:
            self.handleError(record)


# ── Module-level logger setup ─────────────────────────────────────────────────
_handler = _FmtHandler()
_handler.setFormatter(logging.Formatter("%(message)s"))

log = logging.getLogger("luminary")
log.setLevel(logging.DEBUG)
if not log.handlers:
    log.addHandler(_handler)
log.propagate = False


# ── Convenience helpers ───────────────────────────────────────────────────────

def _tag(tag: str) -> dict:
    return {"extra": {"tag": tag}}


def banner(title: str, width: int = 58) -> None:
    """Print a prominent box banner (e.g. at startup)."""
    bar   = "━" * width
    inner = title.center(width)
    sys.stdout.write(
        f"\n{_B}{_CYN}┌{bar}┐\n"
        f"│{_WHT}{_B}{inner}{_CYN}│\n"
        f"└{bar}┘{_R}\n\n"
    )
    sys.stdout.flush()


def section(tag: str, msg: str) -> None:
    """Print a thin divider line with an annotation."""
    color  = _TAG_COLORS.get(tag, _WHT)
    ts     = datetime.now(timezone.utc).strftime("%H:%M:%S")
    width  = max(0, 58 - len(tag) - len(msg) - 4)
    bar    = "─" * width
    sys.stdout.write(f"{_DIM}{ts}{_R} {color}{_B}[{tag}]{_R} {_DIM}{msg} {bar}{_R}\n")
    sys.stdout.flush()


def ok(tag: str, msg: str) -> None:
    color = _TAG_COLORS.get(tag, _WHT)
    ts    = datetime.now(timezone.utc).strftime("%H:%M:%S")
    sys.stdout.write(f"{_DIM}{ts}{_R} {color}{_B}[{tag}]{_R}  {_GRN}✓{_R}  {msg}\n")
    sys.stdout.flush()


def step(tag: str, msg: str) -> None:
    color = _TAG_COLORS.get(tag, _WHT)
    ts    = datetime.now(timezone.utc).strftime("%H:%M:%S")
    sys.stdout.write(f"{_DIM}{ts}{_R} {color}{_B}[{tag}]{_R}  {_DIM}→{_R}  {msg}\n")
    sys.stdout.flush()


def warn(tag: str, msg: str) -> None:
    color = _TAG_COLORS.get(tag, _WHT)
    ts    = datetime.now(timezone.utc).strftime("%H:%M:%S")
    sys.stdout.write(f"{_DIM}{ts}{_R} {color}{_B}[{tag}]{_R}  {_YLW}⚠{_R}  {_YLW}{msg}{_R}\n")
    sys.stdout.flush()


def error(tag: str, msg: str) -> None:
    ts = datetime.now(timezone.utc).strftime("%H:%M:%S")
    sys.stdout.write(
        f"{_DIM}{ts}{_R} {_RED}{_B}[{tag}]{_R}  "
        f"{_RED}✕  Error{_R}\n"
        f"{' ' * (len(ts) + len(tag) + 5)}{_RED}{msg}{_R}\n"
    )
    sys.stdout.flush()


class Timer:
    """Context manager / manual stopwatch with labelled output."""
    def __init__(self, tag: str, label: str) -> None:
        self.tag   = tag
        self.label = label
        self._t0: float = 0.0

    def start(self) -> "Timer":
        self._t0 = time.perf_counter()
        return self

    def elapsed(self) -> float:
        return time.perf_counter() - self._t0

    def done(self, extra: str = "") -> float:
        secs = self.elapsed()
        msg  = f"{self.label}  {_DIM}({secs:.2f}s){_R}"
        if extra:
            msg += f"  {_DIM}{extra}{_R}"
        ok(self.tag, msg)
        return secs

    def __enter__(self) -> "Timer":
        return self.start()

    def __exit__(self, *_) -> None:
        self.done()