from __future__ import annotations

import os
import sys
from datetime import datetime

from one_run.models import RunRuntime


def _total_ram_gb() -> float | None:
    if not hasattr(os, "sysconf"):
        return None
    keys = os.sysconf_names
    if "SC_PHYS_PAGES" not in keys or "SC_PAGE_SIZE" not in keys:
        return None
    try:
        pages = os.sysconf("SC_PHYS_PAGES")
        page_size = os.sysconf("SC_PAGE_SIZE")
    except (ValueError, OSError):
        return None
    if pages <= 0 or page_size <= 0:
        return None
    return round((pages * page_size) / (1024**3), 2)


def _ru_maxrss_to_mb(value: int) -> float:
    if sys.platform == "darwin":
        return round(value / (1024**2), 2)
    return round(value / 1024, 2)


def local_runtime_snapshot(duration_sec: float | None, peak_rss_raw: int | None) -> RunRuntime:
    peak_rss_mb = _ru_maxrss_to_mb(peak_rss_raw) if peak_rss_raw is not None else None
    return RunRuntime(
        collected_at=datetime.now(),
        duration_sec=round(duration_sec, 3) if duration_sec is not None else None,
        peak_rss_mb=peak_rss_mb,
        cpu_count=os.cpu_count(),
        total_ram_gb=_total_ram_gb(),
    )


def slurm_runtime_snapshot(
    *,
    elapsed: str | None,
    max_rss: str | None,
    alloc_cpus: int | None,
) -> RunRuntime:
    return RunRuntime(
        collected_at=datetime.now(),
        slurm_elapsed=elapsed,
        slurm_max_rss=max_rss,
        slurm_alloc_cpus=alloc_cpus,
    )
