"""Persistent conservative cost accounting, including interrupted sessions."""

import json
import time
from pathlib import Path

from .common import atomic_json


class Budget:
    def __init__(self, cfg, run_dir, preparation_cost=0.0, clock=time.monotonic, extension=None):
        self.cfg = dict(cfg["modal"])
        if extension is not None:
            from .continuation import LIMITS

            if extension != LIMITS:
                raise ValueError("Budget extension exceeds the explicitly authorized limits")
            self.cfg.update(extension)
        self.path = Path(run_dir) / "budget.json"
        self.clock = clock
        self.started = clock()
        self.shutdown_requested = False
        self.rate = (
            self.cfg["gpu_usd_second"]
            + self.cfg["cpu"] * self.cfg["cpu_usd_core_second"]
            + self.cfg["memory_mib"] / 1024 * self.cfg["memory_usd_gib_second"]
        ) * self.cfg["cost_safety_multiplier"]
        ledger = json.loads(self.path.read_text()) if self.path.exists() else {}
        if (
            ledger.get("settings")
            and ledger["settings"] != self.cfg
            and not (extension is not None and ledger["settings"] == cfg["modal"])
        ):
            raise ValueError("Cannot change budget/resource settings on resume")
        self.prior_seconds = ledger.get("worker_seconds", 0.0)
        if ledger and not ledger.get("clean_exit", False):
            # Heartbeats are committed at least once a minute; charge an extra two minutes
            # after interruption instead of treating unrecorded work as free.
            self.prior_seconds += 120
        self.preparation_cost = max(preparation_cost, ledger.get("preparation_cost_usd", 0.0))
        self.persist()

    @property
    def seconds(self):
        return self.prior_seconds + self.clock() - self.started

    @property
    def cost(self):
        return self.preparation_cost + self.seconds * self.rate

    def stop(self, training=False):
        if self.shutdown_requested:
            return "shutdown"
        reserve = self.cfg["finalization_reserve_seconds"] if training else 0
        if self.seconds >= self.cfg["max_worker_seconds"] - reserve:
            return "runtime_budget"
        if self.cost >= self.cfg["stop_cost_usd"] - reserve * self.rate:
            return "cost_budget"
        return None

    def metrics(self):
        return {
            "budget/worker_seconds": self.seconds,
            "budget/estimated_usd": self.cost,
            "budget/remaining_usd": max(0, self.cfg["budget_usd"] - self.cost),
        }

    def persist(self, clean=False):
        atomic_json(
            self.path,
            {
                "settings": self.cfg,
                "worker_seconds": self.seconds,
                "preparation_cost_usd": self.preparation_cost,
                "estimated_usd": self.cost,
                "clean_exit": clean,
            },
        )