src/logging_sink.py

"""
Configurable log sink for execution-trace / audit persistence.

The agent already accumulates a full execution trace in-memory
(`WorkflowEngine.trace_logs` + `message_history`, surfaced by
`CodeAgent.get_trace()`): prompts, tool invocations, generated code, and
dataset-load steps. This module abstracts *where* that trace JSON is written so
the destination is a CONFIG value, not hardcoded — prep for a future migration
to OHSU-managed AWS.

Selection is by the `LOG_SINK` env var (`local` | `hf` | `s3`), default
`local`. Every sink implements one method::

    persist_trace(run_id: str, trace: dict) -> str | None

returning a location string (path / URL) on success, or None on failure /
no-op. Callers wrap the call so a logging failure NEVER crashes a user request.

Env vars (all optional; sensible defaults):

  LOG_SINK               local | hf | s3            (default: local)

  local sink:
    LOG_SINK_LOCAL_DIR   output directory           (default: ./run_logs)

  hf sink (reproduces today's HuggingFace-dataset behavior):
    LOG_SINK_HF_DATASET  dataset repo id            (default: anne-voigt/decoupleRpy_results)
    decouplerpy_results_token   HF write token      (existing var; no token => no-op)

  s3 sink (OHSU-managed AWS migration target; see ADR-0009):
    LOG_SINK_S3_BUCKET   bucket name                (required; no bucket => error)
    LOG_SINK_S3_REGION   AWS region                 (default: us-west-2)
    LOG_SINK_S3_PREFIX   key prefix                 (default: runs)

  The s3 sink writes trace blobs only. At-rest encryption (SSE-KMS), versioning,
  lifecycle/retention, and access are enforced by the BUCKET (bucket policy /
  IAM / default encryption) per OHSU standard — the sink manages neither keys
  nor lifecycle (ADR-0008 posture: the store governs retention). Credentials
  come from the task's IAM role via boto3's default credential chain — never
  long-lived keys in env. boto3 is imported lazily so this module imports even
  where boto3 is absent (the HF-hosted `hf`/`local` posture does not need it;
  boto3 becomes a runtime requirement at AWS cutover).
"""

import json
import os
from datetime import datetime


# --------------------------------------------------------------------------- #
# Base interface
# --------------------------------------------------------------------------- #
class LogSink:
    """Common interface for trace persistence."""

    name = "base"

    def persist_trace(self, run_id: str, trace: dict) -> "str | None":
        """Persist the trace dict for `run_id`. Return a location string or None."""
        raise NotImplementedError


# --------------------------------------------------------------------------- #
# local (DEFAULT) — always-on, no external deps
# --------------------------------------------------------------------------- #
class LocalLogSink(LogSink):
    """Write the trace JSON to a configurable directory on local disk."""

    name = "local"

    def __init__(self, directory: str = None):
        self.directory = directory or os.environ.get("LOG_SINK_LOCAL_DIR", "./run_logs")

    def persist_trace(self, run_id: str, trace: dict) -> "str | None":
        os.makedirs(self.directory, exist_ok=True)
        path = os.path.join(self.directory, f"{run_id}_trace.json")
        with open(path, "w", encoding="utf-8") as f:
            json.dump(trace, f, indent=2, ensure_ascii=False, default=str)
        print(f"[log_sink:local] Trace saved: {path}")
        return path


# --------------------------------------------------------------------------- #
# hf — refactor of the existing HuggingFace-dataset behavior
# --------------------------------------------------------------------------- #
class HFLogSink(LogSink):
    """Persist the trace JSON to an HF Dataset repo under runs/<run_id>/trace.json.

    Preserves today's behavior: token comes from `decouplerpy_results_token`
    (same var `HFResultsStorage` uses); no token => no-op (returns None).
    huggingface_hub is imported lazily so a missing dep never breaks import.
    """

    name = "hf"

    def __init__(self, repo_id: str = None):
        self.repo_id = repo_id or os.environ.get(
            "LOG_SINK_HF_DATASET", "anne-voigt/decoupleRpy_results"
        )
        self._api = None

    def _get_api(self):
        if self._api is None:
            token = os.environ.get("decouplerpy_results_token")
            if not token:
                print("[log_sink:hf] decouplerpy_results_token not set — uploads disabled")
                return None
            from huggingface_hub import HfApi  # lazy

            self._api = HfApi(token=token)
        return self._api

    def persist_trace(self, run_id: str, trace: dict) -> "str | None":
        import tempfile

        api = self._get_api()
        if api is None:
            return None

        remote_path = f"runs/{run_id}/trace.json"
        tmp_path = None
        try:
            with tempfile.NamedTemporaryFile(
                mode="w", suffix=".json", delete=False, encoding="utf-8"
            ) as f:
                json.dump(trace, f, indent=2, ensure_ascii=False, default=str)
                tmp_path = f.name
            api.upload_file(
                path_or_fileobj=tmp_path,
                path_in_repo=remote_path,
                repo_id=self.repo_id,
                repo_type="dataset",
            )
        finally:
            if tmp_path and os.path.exists(tmp_path):
                os.unlink(tmp_path)

        url = f"https://huggingface.co/datasets/{self.repo_id}/blob/main/{remote_path}"
        print(f"[log_sink:hf] Trace saved: {url}")
        return url


# --------------------------------------------------------------------------- #
# s3 — OHSU-managed AWS migration target (ADR-0009)
# --------------------------------------------------------------------------- #
class S3LogSink(LogSink):
    """Write the trace JSON to S3 via boto3 `put_object`.

    Target key: ``{LOG_SINK_S3_PREFIX}/{run_id}/trace.json`` in
    ``LOG_SINK_S3_BUCKET``; returns the ``s3://…`` URI on success.

    boto3 is imported LAZILY (and the client cached), so this module imports
    cleanly where boto3 is absent — the HF-hosted `hf`/`local` posture never
    touches it. Encryption / retention / access are the bucket's job, not the
    sink's (ADR-0009): no SSE or ACL params are set here, so the bucket's
    default encryption and lifecycle policy govern the object. Credentials come
    from boto3's default chain (the task IAM role), never env keys.
    """

    name = "s3"

    def __init__(self):
        self.bucket = os.environ.get("LOG_SINK_S3_BUCKET")
        self.region = os.environ.get("LOG_SINK_S3_REGION", "us-west-2")
        # Normalize the prefix: drop surrounding slashes so key joining never
        # yields a leading/double slash (e.g. "runs/" or "/runs" -> "runs").
        self.prefix = os.environ.get("LOG_SINK_S3_PREFIX", "runs").strip("/")
        self._client = None

    def _get_client(self):
        if self._client is None:
            import boto3  # lazy — only needed on the s3 path

            self._client = boto3.client("s3", region_name=self.region)
        return self._client

    def persist_trace(self, run_id: str, trace: dict) -> "str | None":
        if not self.bucket:
            # LOG_SINK=s3 was selected but no bucket configured — a config
            # error. Raise loudly; persist_trace_safe still swallows it so a
            # run never crashes on a logging misconfiguration.
            raise ValueError(
                "LOG_SINK=s3 requires LOG_SINK_S3_BUCKET to be set "
                "(the OHSU-managed audit bucket). See ADR-0009."
            )

        key = f"{self.prefix}/{run_id}/trace.json" if self.prefix else f"{run_id}/trace.json"
        body = json.dumps(trace, indent=2, ensure_ascii=False, default=str).encode("utf-8")

        self._get_client().put_object(
            Bucket=self.bucket,
            Key=key,
            Body=body,
            ContentType="application/json",
        )

        uri = f"s3://{self.bucket}/{key}"
        print(f"[log_sink:s3] Trace saved: {uri}")
        return uri


# --------------------------------------------------------------------------- #
# Factory
# --------------------------------------------------------------------------- #
_SINKS = {
    "local": LocalLogSink,
    "hf": HFLogSink,
    "s3": S3LogSink,
}


def get_log_sink(kind: str = None) -> LogSink:
    """Return the configured LogSink.

    `kind` overrides the `LOG_SINK` env var; default is `local`. An unknown
    value raises ValueError (fail loud on misconfiguration, before any run).
    """
    kind = (kind or os.environ.get("LOG_SINK", "local")).strip().lower()
    if kind not in _SINKS:
        raise ValueError(f"Unknown LOG_SINK={kind!r}. Valid values: {sorted(_SINKS)}.")
    return _SINKS[kind]()


def persist_trace_safe(sink: LogSink, run_id: str, trace: dict) -> "str | None":
    """Persist a trace through `sink`, swallowing any error (warn, don't crash).

    A logging failure must NEVER take down a user request, so this wraps the
    write in try/except and returns None on any failure.
    """
    if run_id is None:
        run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    try:
        return sink.persist_trace(run_id, trace)
    except Exception as e:  # noqa: BLE001 — logging must never crash a request
        print(f"[log_sink:{getattr(sink, 'name', '?')}] persist_trace failed: {e}")
        return None