SimurgArray/src/simurg/core/log.py

"""Structured message reporting — spec CLI-007 (logging), CLI-008 (rate limiting), CLI-009 (run summary).

Messages carry a stable catalogue ``code``. Human-readable lines go to stderr; with ``--log-json`` each record is
one JSON object per line (schema ``simurg.log/1``). A log file, when configured, always receives JSON Lines.
"""

from __future__ import annotations

import datetime as _dt
import json
import sys
import traceback
from collections import Counter
from dataclasses import dataclass, field
from pathlib import Path
from typing import IO, Any

from simurg.core.errors import CATALOGUE

LOG_SCHEMA = "simurg.log/1"
LEVELS = {"debug": 10, "info": 20, "warning": 30, "error": 40}
MAX_PER_SAMPLE = 20


def _now() -> str:
    return _dt.datetime.now(_dt.UTC).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"


@dataclass
class Reporter:
    command: str = ""
    run_id: str = ""
    level: str = "info"
    json_stderr: bool = False
    debug: bool = False
    quiet: bool = False
    log_file: Path | None = None
    stream: IO[str] = field(default_factory=lambda: sys.stderr)
    counts: Counter[str] = field(default_factory=Counter)
    _per_sample: Counter[tuple[str, str]] = field(default_factory=Counter)
    _suppressed: Counter[tuple[str, str]] = field(default_factory=Counter)
    _fh: IO[str] | None = None

    def open(self) -> None:
        if self.log_file is not None:
            self.log_file.parent.mkdir(parents=True, exist_ok=True)
            self._fh = self.log_file.open("a", encoding="utf-8")

    def close(self) -> None:
        self._flush_suppressed()
        if self._fh is not None:
            self._fh.close()
            self._fh = None

    # -- core ---------------------------------------------------------------------------------------------
    def emit(self, level: str, code: str | None, msg: str, *, exc: BaseException | None = None, **ctx: Any) -> None:
        if code:
            self.counts[code] += 1
        sample = ctx.get("sample_id") or ""
        if code and level == "warning":
            key = (code, sample)
            self._per_sample[key] += 1
            if self._per_sample[key] > MAX_PER_SAMPLE:
                self._suppressed[key] += 1
                self._write_file(level, code, msg, exc, ctx)
                return
        record_exc = "".join(traceback.format_exception(exc)) if (exc is not None and self.debug) else None
        self._write_file(level, code, msg, exc, ctx)
        if LEVELS.get(level, 20) < LEVELS.get(self.level, 20):
            return
        if self.quiet and level in ("info", "debug"):
            return
        if self.json_stderr:
            self.stream.write(json.dumps(self._record(level, code, msg, record_exc, ctx), ensure_ascii=False) + "\n")
        else:
            prefix = {"error": "ERROR", "warning": "WARNING", "info": "INFO", "debug": "DEBUG"}.get(level, level)
            code_s = f" [{code}]" if code else ""
            self.stream.write(f"{prefix}{code_s}: {msg}\n")
            if record_exc:
                self.stream.write(record_exc)
        self.stream.flush()

    def _record(self, level: str, code: str | None, msg: str, exc: str | None, ctx: dict[str, Any]) -> dict[str, Any]:
        rec: dict[str, Any] = {
            "schema": LOG_SCHEMA,
            "ts": _now(),
            "level": level,
            "code": code,
            "msg": msg,
            "command": self.command,
            "run_id": self.run_id,
            "step": ctx.pop("step", None),
            "sample_id": ctx.pop("sample_id", None),
            "locus": ctx.pop("locus", None),
            "file": ctx.pop("file", None),
            "count": ctx.pop("count", 1),
            "elapsed_ms": ctx.pop("elapsed_ms", None),
            "exc": exc,
        }
        rec.update(ctx)
        return rec

    def _write_file(
        self, level: str, code: str | None, msg: str, exc: BaseException | None, ctx: dict[str, Any]
    ) -> None:
        if self._fh is None:
            return
        ex = "".join(traceback.format_exception(exc)) if exc is not None else None
        self._fh.write(json.dumps(self._record(level, code, msg, ex, dict(ctx)), ensure_ascii=False) + "\n")

    def _flush_suppressed(self) -> None:
        for (code, sample), n in sorted(self._suppressed.items()):
            where = f" for {sample}" if sample else ""
            self.emit("info", None, f"{code} repeated {n} more times{where}")
        self._suppressed.clear()

    # -- helpers ------------------------------------------------------------------------------------------
    def warn(self, code: str, msg: str, **ctx: Any) -> None:
        self.emit(CATALOGUE[code].level if code in CATALOGUE else "warning", code, msg, **ctx)

    def info(self, msg: str, code: str | None = None, **ctx: Any) -> None:
        self.emit("info", code, msg, **ctx)

    def error(self, code: str, msg: str, *, exc: BaseException | None = None, **ctx: Any) -> None:
        self.emit("error", code, msg, exc=exc, **ctx)

    def warning_counts(self) -> dict[str, int]:
        return dict(sorted(self.counts.items()))