"""Fail-closed local cost reservations for serialized Modal reference jobs."""

import fcntl
import json
import math
import os
import time
import uuid
from contextlib import contextmanager
from pathlib import Path


class Budget:
    def __init__(self, path, cap=250.0):
        self.path, self.cap = Path(path), cap
        if not math.isfinite(cap) or not 0 < cap <= 250:
            raise ValueError("The authorized cap is at most US$250")

    @contextmanager
    def reserve(self, amount, label):
        if not math.isfinite(amount) or amount <= 0:
            raise ValueError("Positive finite reservation required")
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with self.path.with_suffix(".lock").open("a") as lock:
            # Held until the synchronous remote call finishes. Refuse concurrent jobs.
            try:
                fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
            except BlockingIOError:
                raise RuntimeError("Another reference job already owns the GPU budget") from None
            ledger = (
                json.loads(self.path.read_text()) if self.path.exists() else {"cap": self.cap, "entries": []}
            )
            if ledger["cap"] != self.cap:
                raise ValueError("Budget cap changed")
            if any(e["status"] == "reserved" for e in ledger["entries"]):
                raise RuntimeError(
                    "Unresolved prior dispatch: reconcile its remote status before another GPU job"
                )
            if any(
                type(e.get("reserved_usd")) not in (int, float)
                or not math.isfinite(e["reserved_usd"])
                or e["reserved_usd"] <= 0
                or e.get("status") not in ("reserved", "completed_conservative_charge")
                for e in ledger["entries"]
            ):
                raise ValueError("Corrupt budget ledger")
            used = sum(e["reserved_usd"] for e in ledger["entries"])
            if used + amount > self.cap:
                raise RuntimeError("Modal budget exhausted before dispatch")
            entry = {
                "id": uuid.uuid4().hex,
                "label": label,
                "reserved_usd": amount,
                "status": "reserved",
                "time": time.time(),
            }
            ledger["entries"].append(entry)
            self._write(ledger)
            # If interrupted, leave the reservation unresolved.
            yield entry
            entry["status"] = "completed_conservative_charge"
            entry["completed"] = time.time()
            self._write(ledger)

    def _write(self, value):
        tmp = self.path.with_suffix(".tmp")
        with tmp.open("w") as file:
            file.write(json.dumps(value, indent=2))
            file.flush()
            os.fsync(file.fileno())
        tmp.replace(self.path)