#!/usr/bin/env python3
"""Read-only WSL host/guest capacity check. Exit 0/1/2/3: OK/warn/critical/unknown."""

import argparse
from datetime import datetime, timezone
import json
import math
from pathlib import Path
import shutil

GIB = 1024 ** 3
LEVELS = {"ok": 0, "warning": 1, "critical": 2, "unknown": 3}


def nonnegative_gib(value):
    number = float(value)
    if not math.isfinite(number) or number < 0:
        raise argparse.ArgumentTypeError("reserve must be finite and nonnegative")
    return number


def check_volume(path, warning_gib, critical_gib, reserve_bytes, require_mount=False):
    result = {"path": str(path), "warning_bytes": warning_gib * GIB,
              "critical_bytes": critical_gib * GIB, "reserve_bytes": reserve_bytes}
    try:
        if require_mount and not Path(path).is_mount():
            raise OSError("host drive is not mounted")
        usage = shutil.disk_usage(path)
        remaining = usage.free - reserve_bytes
        result.update(total_bytes=usage.total, free_bytes=usage.free,
                      free_after_reserve_bytes=remaining)
        result["status"] = ("critical" if remaining < critical_gib * GIB else
                            "warning" if remaining < warning_gib * GIB else "ok")
    except OSError:
        result.update(status="unknown", error="capacity_unavailable_or_drive_unmounted")
    return result


def collect(reserve_bytes=0):
    volumes = [check_volume("/mnt/c", 60, 30, reserve_bytes, require_mount=True),
               check_volume("/", 40, 20, reserve_bytes)]
    status = max((v["status"] for v in volumes), key=LEVELS.get)
    return {"schema": "platform.workstation-disk-guard.v1",
            "observed_at": datetime.now(timezone.utc).isoformat(),
            "status": status, "volumes": volumes}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--reserve-gib", type=nonnegative_gib, default=0,
                        help="expected additional build use; subtract from BOTH host and guest free space")
    args = parser.parse_args()
    report = collect(int(args.reserve_gib * GIB))
    print(json.dumps(report, indent=2, sort_keys=True))
    return LEVELS[report["status"]]


if __name__ == "__main__":
    raise SystemExit(main())