src/observability.py

from __future__ import annotations

import logging
import sys
import threading
import time
from collections import deque
from contextlib import contextmanager
from pathlib import Path
from typing import Callable, Iterator


LOG_LEVELS = ["INFO", "DEBUG"]
DEFAULT_LOG_LEVEL = "INFO"
LOG_TAIL_LINES = 160
LOG_TAIL_CHARS = 24_000
INFO_PROGRESS_INTERVAL_SECONDS = 10.0
INFO_PROGRESS_CHECKPOINTS = (0, 25, 50, 75, 100)


def normalize_log_level(value) -> str:
    text = str(value or DEFAULT_LOG_LEVEL).strip().upper()
    return text if text in LOG_LEVELS else DEFAULT_LOG_LEVEL


def log_level_number(value) -> int:
    return getattr(logging, normalize_log_level(value), logging.INFO)


def _short_message(value: str, limit: int = 180) -> str:
    text = " ".join(str(value or "").split())
    if len(text) <= limit:
        return text
    return text[: limit - 1] + "…"


class ProgressBridge:
    """Small adapter around Gradio Progress, with safe no-op behavior in tests.

    This class is only a display sink. Runtime logs are produced independently by
    :class:`ProgressReporter`, so a Gradio implementation change cannot remove the
    diagnostic record.
    """

    def __init__(self, progress=None):
        self.progress = progress
        self.value = 0.0
        self._lock = threading.Lock()

    def update(self, value: float, desc: str, *, force: bool = True) -> None:
        del force  # Compatibility with older callers; display failures are always ignored.
        try:
            numeric = max(0.0, min(1.0, float(value)))
        except (TypeError, ValueError):
            numeric = self.value
        with self._lock:
            self.value = numeric
        if self.progress is None:
            return
        try:
            self.progress(numeric, desc=_short_message(desc))
        except Exception:
            pass

    def finish(self, desc: str = "Completed") -> None:
        self.update(1.0, desc)


class ProgressReporter:
    """Single source of truth for application progress.

    Every application progress event can be sent to both the UI and the runtime
    log from this one object. INFO remains compact: explicit stages, quarter
    checkpoints, completion, and ten-second heartbeats. DEBUG additionally gets
    each application update. Third-party tqdm detail is handled separately by
    ``StructuredTqdmCapture`` and is enabled only for DEBUG jobs.
    """

    def __init__(
        self,
        bridge: ProgressBridge,
        emit_info: Callable[..., None],
        emit_debug: Callable[..., None] | None = None,
        *,
        level=DEFAULT_LOG_LEVEL,
        info_interval: float = INFO_PROGRESS_INTERVAL_SECONDS,
        initial_value: float = 0.0,
    ):
        self.bridge = bridge
        self.emit_info = emit_info
        self.emit_debug = emit_debug or emit_info
        self.level_name = normalize_log_level(level)
        self.info_interval = max(0.1, float(info_interval))
        self.value = max(0.0, min(1.0, float(initial_value)))
        self._last_info_emit = 0.0
        self._last_checkpoint = -1
        self._lock = threading.RLock()

    @staticmethod
    def _checkpoint(percent: float) -> int:
        reached = 0
        for value in INFO_PROGRESS_CHECKPOINTS:
            if percent + 1e-9 >= value:
                reached = value
            else:
                break
        return reached

    def update(
        self,
        value: float,
        desc: str,
        *,
        stage: str = "pipeline",
        force_info: bool = False,
        current=None,
        total=None,
        unit: str | None = None,
        source: str = "application",
        ui: bool = True,
    ) -> None:
        try:
            numeric = max(0.0, min(1.0, float(value)))
        except (TypeError, ValueError):
            numeric = self.value
        message = _short_message(desc)
        percent = numeric * 100.0
        now = time.monotonic()
        checkpoint = self._checkpoint(percent)

        if ui:
            self.bridge.update(numeric, message)

        fields = {
            "progress_stage": stage,
            "percent": f"{percent:.1f}",
            "checkpoint": checkpoint,
            "current": current,
            "total": total,
            "unit": unit,
            "source": source,
            "desc": message,
        }

        with self._lock:
            self.value = numeric
            if self.level_name == "DEBUG":
                try:
                    self.emit_debug("progress-update", **fields)
                except Exception:
                    pass

            should_info = (
                force_info
                or self._last_info_emit == 0.0
                or checkpoint > self._last_checkpoint
                or numeric >= 1.0
                or now - self._last_info_emit >= self.info_interval
            )
            if should_info:
                try:
                    self.emit_info("progress", **fields)
                except Exception:
                    pass
                self._last_info_emit = now
                self._last_checkpoint = max(self._last_checkpoint, checkpoint)

    def heartbeat(
        self,
        *,
        stage: str,
        desc: str,
        started_at: float,
        source: str = "application-heartbeat",
    ) -> None:
        elapsed = max(0.0, time.monotonic() - started_at)
        with self._lock:
            numeric = self.value
        # Heartbeats are intentionally log-first. Updating Gradio from a helper
        # thread is unnecessary and less stable than the runtime log record.
        try:
            self.emit_info(
                "progress-heartbeat",
                progress_stage=stage,
                percent=f"{numeric * 100.0:.1f}",
                elapsed_seconds=f"{elapsed:.1f}",
                source=source,
                desc=_short_message(desc),
            )
        except Exception:
            pass

    @contextmanager
    def long_operation(
        self,
        value: float,
        desc: str,
        *,
        stage: str,
        heartbeat_interval: float | None = None,
    ) -> Iterator[None]:
        interval = max(0.1, float(heartbeat_interval or self.info_interval))
        self.update(value, desc, stage=stage, force_info=True)
        stop = threading.Event()
        started_at = time.monotonic()

        def heartbeat_loop() -> None:
            while not stop.wait(interval):
                self.heartbeat(stage=stage, desc=desc, started_at=started_at)

        thread = threading.Thread(
            target=heartbeat_loop,
            name=f"sesa-progress-{stage}",
            daemon=True,
        )
        thread.start()
        try:
            yield
        finally:
            stop.set()
            thread.join(timeout=min(1.0, interval))

    def finish(self, desc: str = "Completed", *, stage: str = "pipeline") -> None:
        self.update(1.0, desc, stage=stage, force_info=True)


class StructuredTqdmCapture:
    """Mirror throttled third-party tqdm updates into DEBUG runtime logs.

    The application progress log does not depend on this hook. It is diagnostic
    detail only and should be enabled for DEBUG jobs. Patching ``update`` and
    ``close`` is more reliable than observing terminal ``display`` calls and also
    works when a UI library replaces tqdm's rendering behavior.
    """

    _patch_lock = threading.RLock()

    def __init__(
        self,
        emit_debug: Callable[..., None],
        *,
        stage: str = "tqdm-progress",
        min_percent_delta: float = 1.0,
        min_interval: float = 0.5,
    ):
        self.emit_debug = emit_debug
        self.stage = str(stage)
        self.min_percent_delta = max(0.0, float(min_percent_delta))
        self.min_interval = max(0.0, float(min_interval))
        self._owner_thread = threading.get_ident()
        self._patched: list[tuple[type, Callable, Callable]] = []
        self._state: dict[int, tuple[float | None, float, float | None]] = {}

    def __enter__(self) -> "StructuredTqdmCapture":
        self._patch_lock.acquire()
        classes: list[type] = []
        try:
            import tqdm as tqdm_module
            from tqdm import auto as tqdm_auto
            from tqdm import std as tqdm_std

            for candidate in (
                getattr(tqdm_module, "tqdm", None),
                getattr(tqdm_auto, "tqdm", None),
                getattr(tqdm_std, "tqdm", None),
            ):
                if isinstance(candidate, type) and candidate not in classes:
                    classes.append(candidate)

            for cls in classes:
                original_update = cls.update
                original_close = cls.close
                capture = self

                def update(instance, n=1, _original=original_update, _capture=capture):
                    result = _original(instance, n=n)
                    _capture._observe(instance)
                    return result

                def close(instance, _original=original_close, _capture=capture):
                    _capture._observe(instance, force=True)
                    return _original(instance)

                cls.update = update
                cls.close = close
                self._patched.append((cls, original_update, original_close))
        except Exception:
            self.__exit__(None, None, None)
            raise
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        for cls, original_update, original_close in reversed(self._patched):
            try:
                cls.update = original_update
                cls.close = original_close
            except Exception:
                pass
        self._patched.clear()
        try:
            self._patch_lock.release()
        except RuntimeError:
            pass

    def _observe(self, bar, *, force: bool = False) -> None:
        if threading.get_ident() != self._owner_thread:
            return
        try:
            current = float(getattr(bar, "n", 0.0) or 0.0)
            total_raw = getattr(bar, "total", None)
            total = float(total_raw) if total_raw not in (None, 0) else None
            percent = (100.0 * current / total) if total else None
            desc = _short_message(getattr(bar, "desc", "") or "tqdm", limit=120)
        except Exception:
            return

        now = time.monotonic()
        key = id(bar)
        last_percent, last_time, last_current = self._state.get(key, (None, 0.0, None))
        completed = total is not None and current >= total
        first = last_current is None
        advanced_enough = (
            percent is not None
            and (last_percent is None or percent - last_percent >= self.min_percent_delta)
        )
        timed_out = now - last_time >= self.min_interval
        if not (force or first or completed or advanced_enough or timed_out):
            return
        if last_current == current and not (force or completed or first):
            return

        self._state[key] = (percent, now, current)
        fields = {
            "current": f"{current:g}",
            "total": f"{total:g}" if total is not None else "unknown",
            "percent": f"{percent:.1f}" if percent is not None else "unknown",
            "desc": desc,
            "source": "third-party-tqdm",
        }
        try:
            self.emit_debug(self.stage, **fields)
        except Exception:
            pass


class _CurrentThreadFilter(logging.Filter):
    def __init__(self, thread_id: int, job_id: str):
        super().__init__()
        self.thread_id = thread_id
        self.job_id = str(job_id)

    def filter(self, record: logging.LogRecord) -> bool:
        # Third-party records are accepted only from the serialized GPU worker
        # thread. SESA's own structured events may also come from the heartbeat
        # helper thread and carry an explicit job id.
        return record.thread == self.thread_id or getattr(record, "sesa_job_id", None) == self.job_id


class _DeduplicatingFileHandler(logging.FileHandler):
    def __init__(self, filename: str):
        super().__init__(filename, mode="a", encoding="utf-8", delay=False)
        self._seen_ids: deque[int] = deque(maxlen=2048)
        self._seen_set: set[int] = set()

    def emit(self, record: logging.LogRecord) -> None:
        record_id = id(record)
        if record_id in self._seen_set:
            return
        if len(self._seen_ids) == self._seen_ids.maxlen:
            oldest = self._seen_ids.popleft()
            self._seen_set.discard(oldest)
        self._seen_ids.append(record_id)
        self._seen_set.add(record_id)
        super().emit(record)


class _UtcFormatter(logging.Formatter):
    converter = time.gmtime


class JobLogCapture:
    """Capture application and audio-separator logs for one GPU job/thread."""

    _console_lock = threading.Lock()

    LOGGER_NAMES = (
        "",  # root, for propagated third-party records
        "sesa",
        "audio_separator",
        "separator",
        "backend",
        "common_separator",
        "mdx_separator",
        "mdxc_separator",
        "vr_separator",
        "demucs_separator",
    )

    def __init__(
        self,
        path: Path,
        job_id: str,
        level=DEFAULT_LOG_LEVEL,
        progress: ProgressBridge | None = None,
    ):
        # ``progress`` is retained for source compatibility but deliberately not
        # wired to arbitrary log records. ProgressReporter owns UI updates.
        del progress
        self.path = Path(path)
        self.job_id = str(job_id)
        self.level_name = normalize_log_level(level)
        self.level = log_level_number(self.level_name)
        self.handler: _DeduplicatingFileHandler | None = None
        self._loggers: list[logging.Logger] = []
        self._previous_levels: dict[logging.Logger, int] = {}

    def __enter__(self) -> "JobLogCapture":
        self.path.parent.mkdir(parents=True, exist_ok=True)
        handler = _DeduplicatingFileHandler(str(self.path))
        handler.setLevel(self.level)
        handler.addFilter(_CurrentThreadFilter(threading.get_ident(), self.job_id))
        handler.setFormatter(
            _UtcFormatter(
                fmt=(
                    "%(asctime)sZ | %(levelname)s | %(name)s | %(module)s:%(lineno)d | "
                    f"job={self.job_id} | %(message)s"
                ),
                datefmt="%Y-%m-%dT%H:%M:%S",
            )
        )
        self.handler = handler

        seen: set[int] = set()
        for name in self.LOGGER_NAMES:
            logger = logging.getLogger(name)
            if id(logger) in seen:
                continue
            seen.add(id(logger))
            self._loggers.append(logger)
            self._previous_levels[logger] = logger.level
            if name:
                logger.setLevel(self.level)
            logger.addHandler(handler)

        self.event("job-log-open", level=self.level_name, path=self.path.name)
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        if exc is not None:
            self.event("job-log-exit-exception", error_type=type(exc).__name__, error=str(exc))
        self.event("job-log-close")
        if self.handler is not None:
            try:
                self.handler.flush()
            except Exception:
                pass
        for logger in self._loggers:
            try:
                if self.handler is not None:
                    logger.removeHandler(self.handler)
                logger.setLevel(self._previous_levels.get(logger, logger.level))
            except Exception:
                pass
        if self.handler is not None:
            try:
                self.handler.close()
            except Exception:
                pass

    @staticmethod
    def _render_event(stage: str, fields: dict) -> str:
        parts = [f"stage={stage}"]
        for key, value in fields.items():
            if value is None:
                continue
            safe = str(value).replace("\n", " ").replace("\r", " ")
            parts.append(f"{key}={safe}")
        return " | ".join(parts)

    def event_at(self, severity: int, stage: str, **fields) -> None:
        rendered = self._render_event(stage, fields)
        # Guarantee that structured application events reach Space Run logs even
        # if Gradio/Uvicorn changes the root logging configuration. Duplicate
        # console lines are preferable to missing diagnostics.
        if severity >= self.level:
            timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
            level_name = logging.getLevelName(severity)
            with self._console_lock:
                print(
                    f"{timestamp} | {level_name} | sesa | job={self.job_id} | {rendered}",
                    file=sys.stderr,
                    flush=True,
                )
        logging.getLogger("sesa").log(
            severity,
            rendered,
            extra={"sesa_job_id": self.job_id},
            stacklevel=2,
        )

    def event(self, stage: str, **fields) -> None:
        self.event_at(logging.INFO, stage, **fields)

    def debug_event(self, stage: str, **fields) -> None:
        self.event_at(logging.DEBUG, stage, **fields)

    def warning_event(self, stage: str, **fields) -> None:
        self.event_at(logging.WARNING, stage, **fields)

    def error_event(self, stage: str, **fields) -> None:
        self.event_at(logging.ERROR, stage, **fields)

    def tail(self, max_lines: int = LOG_TAIL_LINES, max_chars: int = LOG_TAIL_CHARS) -> str:
        try:
            lines = self.path.read_text(encoding="utf-8", errors="replace").splitlines()
        except OSError:
            return ""
        text = "\n".join(lines[-max_lines:])
        if len(text) > max_chars:
            text = "…\n" + text[-max_chars:]
        return text


def read_log_tail(path: Path, max_lines: int = LOG_TAIL_LINES, max_chars: int = LOG_TAIL_CHARS) -> str:
    try:
        lines = Path(path).read_text(encoding="utf-8", errors="replace").splitlines()
    except OSError:
        return ""
    text = "\n".join(lines[-max_lines:])
    if len(text) > max_chars:
        text = "…\n" + text[-max_chars:]
    return text