"""Write the workload status file that the host reads after docker wait.

The file is written to /run/pc-status/done.json — a directory mounted by the
Perpetual Compute host into every workload container. The done secret never
enters the container; billing stops on actual container exit, not on this call.

This module is the single source of truth for the file path and schema.
"""

from __future__ import annotations

import json
import os
from typing import Dict, Optional

# Default path; can be overridden for testing via the environment variable.
_STATUS_DIR = os.environ.get("PC_STATUS_DIR", "/run/pc-status")
_STATUS_FILE = os.path.join(_STATUS_DIR, "done.json")
_METRICS_FILE = os.path.join(_STATUS_DIR, "metrics.json")

# Set to True once a terminal status has been written to prevent double-writes.
_written: bool = False

# In-memory metrics accumulator. Written live to metrics.json and included in
# done.json at container exit so the portal can display them.
_metrics: Dict[str, float] = {}


def log_metric(key: str, value: float) -> None:
    """Record a numeric metric visible in the Perpetual Compute portal.

    Values overwrite any previous value for the same key.
    The metric is persisted live to /run/pc-status/metrics.json so the portal
    can surface it even if the container crashes before calling done().

    Example::

        import perpetual_compute as pc
        pc.log_metric("epoch", 42)
        pc.log_metric("loss", 0.312)
        pc.log_metric("throughput_samples_per_sec", 1024.0)
    """
    _metrics[key] = float(value)
    _flush_metrics()


def _flush_metrics() -> None:
    """Persist current metrics to metrics.json (best-effort, non-blocking)."""
    if not _metrics:
        return
    try:
        os.makedirs(_STATUS_DIR, exist_ok=True)
        with open(_METRICS_FILE, "w", encoding="utf-8") as f:
            json.dump(_metrics, f)
    except OSError:
        pass


def write_status(
    status: str,
    reason: Optional[str] = None,
    exit_code: Optional[int] = None,
) -> None:
    """Write the status file and mark as written.

    Idempotent: subsequent calls are no-ops once written.
    Automatically includes any accumulated metrics from log_metric().
    """
    global _written
    if _written:
        return
    _written = True

    payload: dict = {"status": status}
    if reason:
        payload["reason"] = reason
    if exit_code is not None:
        payload["exit_code"] = exit_code
    if _metrics:
        payload["metrics"] = dict(_metrics)

    try:
        os.makedirs(_STATUS_DIR, exist_ok=True)
        with open(_STATUS_FILE, "w", encoding="utf-8") as f:
            json.dump(payload, f)
    except OSError:
        # Best-effort. If /run/pc-status isn't mounted (e.g. local dev), we
        # don't want to crash the workload. The host falls back to exit code.
        pass


def is_written() -> bool:
    return _written


def reset() -> None:
    """Reset written state and metrics — for testing only."""
    global _written, _metrics
    _written = False
    _metrics = {}
