# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Append-only JSONL cost ledger with a run-level spend cap. Never stores prompts or credentials."""
from __future__ import annotations

import json
import threading
import time
from pathlib import Path

from ..config import RESULTS, max_spend_usd, model_info


class SpendCapExceeded(RuntimeError):
    pass


class CostLedger:
    def __init__(self, path: Path | None = None, cap_usd: float | None = None, run_id: str = "adhoc"):
        self.path = path or (RESULTS / "ledger" / f"{run_id}.jsonl")
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.cap = cap_usd if cap_usd is not None else max_spend_usd()
        self.run_id = run_id
        self._lock = threading.Lock()
        self.total_usd = 0.0
        self.calls = 0
        if self.path.exists():
            for line in self.path.read_text().splitlines():
                try:
                    self.total_usd += float(json.loads(line).get("cost_usd", 0.0))
                    self.calls += 1
                except Exception:
                    pass

    @staticmethod
    def estimate_cost(model: str, input_tokens: int, output_tokens: int, cached_tokens: int = 0) -> float:
        try:
            m = model_info(model)
        except KeyError:
            return 0.0
        uncached = max(input_tokens - cached_tokens, 0)
        return (
            uncached * m["input_per_m"] + cached_tokens * m.get("cached_input_per_m", m["input_per_m"]) + output_tokens * m["output_per_m"]
        ) / 1e6

    def record(self, **rec) -> None:
        with self._lock:
            rec.setdefault("ts", time.time())
            rec["run_id"] = self.run_id
            self.total_usd += float(rec.get("cost_usd", 0.0))
            self.calls += 1
            with open(self.path, "a", encoding="utf-8") as fh:
                fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
            if self.total_usd > self.cap:
                raise SpendCapExceeded(f"run {self.run_id}: estimated spend {self.total_usd:.2f} USD exceeds cap {self.cap:.2f}")