"""Worker-owned, conservative spend reservations; no provider or secret access.

One private SQLite file per envelope, shared by all admitted processes on one
host. Missing state is an error: only the explicit operator init creates it.
"""

from __future__ import annotations

import hashlib
import json
import os
import re
import sqlite3
import stat
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import UTC, datetime
from decimal import ROUND_CEILING, ROUND_FLOOR, Decimal, InvalidOperation, localcontext
from pathlib import Path
from typing import Any
from zoneinfo import ZoneInfo


class SpendAdmissionError(RuntimeError):
    """A bounded refusal; never contains prompt, credential or provider output."""


def digest(value: Any) -> str:
    return hashlib.sha256(
        json.dumps(
            value, sort_keys=True, separators=(",", ":"), allow_nan=False
        ).encode()
    ).hexdigest()


def amount(value: Any, *, positive: bool = True) -> Decimal:
    if isinstance(value, bool) or not isinstance(value, (str, int, float)):
        raise SpendAdmissionError("invalid monetary amount")
    if len(str(value)) > 500:
        raise SpendAdmissionError("invalid monetary amount")
    try:
        result = Decimal(str(value))
    except InvalidOperation:
        raise SpendAdmissionError("invalid monetary amount") from None
    if (
        not result.is_finite()
        or result < 0
        or (positive and result == 0)
        or result > 1_000_000
        or abs(result.as_tuple().exponent) > 100
    ):
        raise SpendAdmissionError("invalid monetary amount")
    return result


def micros(value: Decimal) -> int:
    with localcontext() as context:
        context.prec = 256
        return int((value * 1_000_000).to_integral_value(rounding=ROUND_CEILING))


def cap_micros(value: str) -> int:
    with localcontext() as context:
        context.prec = 256
        return int((amount(value) * 1_000_000).to_integral_value(rounding=ROUND_FLOOR))


def converted_micros(usd: Any, eur_per_usd: str) -> int:
    with localcontext() as context:
        context.prec = 256
        return micros(amount(usd, positive=False) * amount(eur_per_usd))


def timestamp(value: str) -> datetime:
    try:
        result = datetime.fromisoformat(value)
        if result.tzinfo is None:
            raise ValueError
        return result.astimezone(UTC)
    except (TypeError, ValueError):
        raise SpendAdmissionError("invalid policy timestamp") from None


@dataclass(frozen=True)
class SpendPolicy:
    version: str
    envelope_id: str
    authority_ref: str
    valid_from: str
    expires_at: str
    timezone: str
    worker_id: str
    activity_definition_id: str
    target_repo: str
    project: str
    profile_ref: str
    profile_sha256: str
    descriptor_sha256: str
    repository_grant_id: str
    max_budget_usd: str
    max_liability_usd: str
    max_turns: int
    eur_per_usd: str
    per_run_eur: str
    daily_eur: str
    total_eur: str

    def __post_init__(self) -> None:
        if self.version != "1":
            raise SpendAdmissionError("unsupported spend policy version")
        for name in self.__dataclass_fields__:
            value = getattr(self, name)
            if name == "max_turns":
                if type(value) is not int or not 1 <= value <= 10000:
                    raise SpendAdmissionError("invalid turn limit")
            elif (
                not isinstance(value, str)
                or not value
                or len(value) > 500
                or any(ord(c) < 32 for c in value)
            ):
                raise SpendAdmissionError("invalid spend policy field")
        for name in ("profile_sha256", "descriptor_sha256"):
            if not re.fullmatch(r"[0-9a-f]{64}", getattr(self, name)):
                raise SpendAdmissionError("invalid runtime digest")
        if not Path(self.target_repo).is_absolute() or "@" not in self.profile_ref:
            raise SpendAdmissionError(
                "spend policy needs an absolute target and versioned profile"
            )
        try:
            ZoneInfo(self.timezone)
        except (ValueError, KeyError):
            raise SpendAdmissionError("invalid budget timezone") from None
        if timestamp(self.valid_from) >= timestamp(self.expires_at):
            raise SpendAdmissionError("empty policy validity interval")
        for name in (
            "max_budget_usd",
            "max_liability_usd",
            "eur_per_usd",
            "per_run_eur",
            "daily_eur",
            "total_eur",
        ):
            amount(getattr(self, name))
        if amount(self.max_budget_usd) > amount(self.max_liability_usd):
            raise SpendAdmissionError("native limit exceeds declared liability")
        if self.reservation > cap_micros(self.per_run_eur) or not amount(
            self.per_run_eur
        ) <= amount(self.daily_eur) <= amount(self.total_eur):
            raise SpendAdmissionError("liability exceeds operating envelope")

    @property
    def reservation(self) -> int:
        return converted_micros(self.max_liability_usd, self.eur_per_usd)

    @property
    def sha256(self) -> str:
        return digest(self.__dict__)

    @classmethod
    def load(cls, path: Path) -> SpendPolicy:
        _private_file(path)
        try:
            if path.stat().st_size > 16384:
                raise ValueError

            def unique(pairs):
                result = {}
                for key, value in pairs:
                    if key in result:
                        raise ValueError
                    result[key] = value
                return result

            policy = cls(**json.loads(path.read_text(), object_pairs_hook=unique))
        except (OSError, ValueError, TypeError):
            raise SpendAdmissionError("invalid spend policy file") from None
        if path.resolve().is_relative_to(Path(policy.target_repo).resolve()):
            raise SpendAdmissionError("spend policy must be outside target repository")
        return policy


def _private_file(path: Path) -> None:
    try:
        info = path.lstat()
    except OSError:
        raise SpendAdmissionError("required private spend file unavailable") from None
    if (
        not stat.S_ISREG(info.st_mode)
        or info.st_nlink != 1
        or info.st_uid != os.getuid()
        or info.st_mode & 0o077
    ):
        raise SpendAdmissionError(
            "spend file must be private, regular and worker-owned"
        )


class SpendLedger:
    def __init__(self, path: Path, policy: SpendPolicy) -> None:
        self.path, self.policy = path, policy
        parent = path.parent
        if not path.is_absolute() or path.resolve().is_relative_to(
            Path(policy.target_repo).resolve()
        ):
            raise SpendAdmissionError(
                "ledger must be absolute and outside target repository"
            )
        try:
            info = parent.lstat()
        except OSError:
            raise SpendAdmissionError("private ledger directory unavailable") from None
        if (
            not stat.S_ISDIR(info.st_mode)
            or info.st_uid != os.getuid()
            or info.st_mode & 0o077
        ):
            raise SpendAdmissionError(
                "ledger directory must be private and worker-owned"
            )

    def initialize(self) -> None:
        """Explicit provisioning only. Never truncate or recreate existing state."""
        try:
            fd = os.open(self.path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
            os.close(fd)
            with self._db(check=False) as db:
                db.execute(
                    "CREATE TABLE envelope (policy TEXT NOT NULL, last_seen TEXT NOT NULL, breached INTEGER NOT NULL CHECK(breached IN (0,1)))"
                )
                db.execute("""CREATE TABLE reservations (
                        run_id TEXT PRIMARY KEY, definition_id TEXT NOT NULL,
                        idempotency_key TEXT NOT NULL, attempt INTEGER NOT NULL,
                        state TEXT NOT NULL CHECK(state IN ('held','charged')),
                        liability INTEGER NOT NULL CHECK(liability>0),
                        start_day TEXT NOT NULL, end_day TEXT,
                        observed_usd TEXT, receipt TEXT,
                        UNIQUE(definition_id, idempotency_key),
                        CHECK((state='held' AND end_day IS NULL) OR
                              (state='charged' AND end_day IS NOT NULL AND end_day>=start_day)));
                """)
                db.execute(
                    "INSERT INTO envelope VALUES (?, ?, 0)",
                    (self.policy.sha256, self.policy.valid_from),
                )
            fd = os.open(self.path.parent, os.O_RDONLY | os.O_DIRECTORY)
            try:
                os.fsync(fd)
            finally:
                os.close(fd)
        except OSError:
            raise SpendAdmissionError("spend ledger initialization refused") from None

    @contextmanager
    def _db(self, *, check: bool = True) -> Iterator[sqlite3.Connection]:
        _private_file(self.path)
        db = None
        try:
            db = sqlite3.connect(
                self.path.as_uri() + "?mode=rw",
                uri=True,
                timeout=10,
                isolation_level=None,
            )
            db.row_factory = sqlite3.Row
            db.execute("PRAGMA synchronous=FULL")
            db.execute("BEGIN IMMEDIATE")
            if check:
                rows = db.execute("SELECT * FROM envelope").fetchall()
                if len(rows) != 1 or rows[0]["policy"] != self.policy.sha256:
                    raise SpendAdmissionError("spend ledger policy mismatch")
            yield db
            db.commit()
        except sqlite3.Error:
            raise SpendAdmissionError(
                "spend ledger unavailable or inconsistent"
            ) from None
        finally:
            if db is not None:
                db.close()  # rolls back any incomplete mutation

    def _clock(self, db: sqlite3.Connection, now: datetime, *, admission: bool) -> str:
        if now.tzinfo is None:
            raise SpendAdmissionError("budget clock must be timezone-aware")
        now = now.astimezone(UTC)
        previous = timestamp(db.execute("SELECT last_seen FROM envelope").fetchone()[0])
        if now < previous:
            raise SpendAdmissionError("budget clock moved backwards")
        if admission and not timestamp(self.policy.valid_from) <= now < timestamp(
            self.policy.expires_at
        ):
            raise SpendAdmissionError("spend policy is not currently valid")
        db.execute("UPDATE envelope SET last_seen=?", (now.isoformat(),))
        return now.astimezone(ZoneInfo(self.policy.timezone)).date().isoformat()

    def _capacity(self, db: sqlite3.Connection, day: str) -> None:
        if db.execute("SELECT breached FROM envelope").fetchone()[0]:
            raise SpendAdmissionError(
                "spend envelope breached; new owner admission required"
            )
        if db.execute("SELECT 1 FROM reservations WHERE state='held'").fetchone():
            raise SpendAdmissionError(
                "unresolved spend reservation; reconcile before dispatch"
            )
        total = db.execute(
            "SELECT COALESCE(SUM(liability),0) FROM reservations"
        ).fetchone()[0]
        daily = db.execute(
            "SELECT COALESCE(SUM(liability),0) FROM reservations WHERE start_day<=? AND end_day>=?",
            (day, day),
        ).fetchone()[0]
        if total + self.policy.reservation > cap_micros(self.policy.total_eur):
            raise SpendAdmissionError("total spend capacity exhausted")
        if daily + self.policy.reservation > cap_micros(self.policy.daily_eur):
            raise SpendAdmissionError("daily spend capacity exhausted")

    def preflight(self, *, now: datetime | None = None) -> None:
        with self._db() as db:
            day = self._clock(db, now or datetime.now(UTC), admission=True)
            self._capacity(db, day)

    def validate_dispatch(
        self, run: Any, config: Any, request: Any, profile: Any, descriptor: Any
    ) -> None:
        p = self.policy
        if (
            config.worker_id != p.worker_id
            or run.claim_owner != p.worker_id
            or run.activity_definition_id != p.activity_definition_id
            or request.project != p.project
            or request.actor != "agt"
            or Path(request.repo).resolve() != Path(p.target_repo).resolve()
            or run.harness_profile_ref != p.profile_ref
            or request.harness_profile_ref != p.profile_ref
            or request.request_id != run.id
            or not run.repository_grant
            or run.repository_grant.grant_id != p.repository_grant_id
            or digest(profile.model_dump(mode="json")) != p.profile_sha256
            or digest(descriptor.model_dump(mode="json")) != p.descriptor_sha256
        ):
            raise SpendAdmissionError("dispatch does not match admitted spend scope")
        if (
            profile.rein.id != "rein-aharness"
            or amount(profile.limits.max_budget_usd) != amount(p.max_budget_usd)
            or profile.limits.max_turns != p.max_turns
        ):
            raise SpendAdmissionError(
                "resolved native limits do not match spend policy"
            )

    def reserve(self, run: Any, *, now: datetime | None = None) -> None:
        for value in (run.id, run.activity_definition_id, run.idempotency_key):
            if not isinstance(value, str) or not value or len(value) > 500:
                raise SpendAdmissionError("invalid spend run identity")
        if type(run.attempt) is not int or run.attempt < 1:
            raise SpendAdmissionError("invalid spend attempt")
        with self._db() as db:
            day = self._clock(db, now or datetime.now(UTC), admission=True)
            if db.execute(
                "SELECT 1 FROM reservations WHERE run_id=? OR (definition_id=? AND idempotency_key=?)",
                (run.id, run.activity_definition_id, run.idempotency_key),
            ).fetchone():
                raise SpendAdmissionError(
                    "run already admitted; workload replay refused"
                )
            self._capacity(db, day)
            db.execute(
                "INSERT INTO reservations VALUES (?, ?, ?, ?, 'held', ?, ?, NULL, NULL, NULL)",
                (
                    run.id,
                    run.activity_definition_id,
                    run.idempotency_key,
                    run.attempt,
                    self.policy.reservation,
                    day,
                ),
            )

    def observe(
        self, run_id: str, raw: dict[str, Any], *, now: datetime | None = None
    ) -> bool:
        """Charge full liability on complete success; never refund self-reported cost."""
        evidence = raw.get("evidence", {})
        try:
            cost = amount(evidence.get("cost_usd"), positive=False)
        except SpendAdmissionError:
            with self._db() as db:
                self._close_requests(db, run_id)
            return False  # Unknown remains held, even for apparent success.
        with self._db() as db:
            day = self._clock(db, now or datetime.now(UTC), admission=False)
            requests_resolved = self._close_requests(db, run_id)
            row = db.execute(
                "SELECT * FROM reservations WHERE run_id=?", (run_id,)
            ).fetchone()
            if row is None or row["state"] != "held":
                raise SpendAdmissionError("reservation observation conflict")
            liability = max(
                row["liability"], converted_micros(str(cost), self.policy.eur_per_usd)
            )
            breached = cost > amount(self.policy.max_liability_usd)
            complete = (
                raw.get("ok") is True
                and evidence.get("outcome") == "succeeded"
                and evidence.get("request_id") == run_id
                and evidence.get("profile_ref") == self.policy.profile_ref
                and evidence.get("session_cleanup") == "succeeded"
                and evidence.get("sandbox_destroy") == "succeeded"
                and not breached
                and requests_resolved
            )
            db.execute(
                "UPDATE reservations SET liability=?, observed_usd=?, state=?, end_day=? WHERE run_id=?",
                (
                    liability,
                    str(cost),
                    "charged" if complete else "held",
                    day if complete else None,
                    run_id,
                ),
            )
            if breached:
                db.execute("UPDATE envelope SET breached=1")
        return complete

    @staticmethod
    def _close_requests(
        db: sqlite3.Connection, run_id: str, *, reconcile: bool = False
    ) -> bool:
        # Existing ledgers without the opt-in request extension remain compatible.
        if not db.execute(
            "SELECT 1 FROM sqlite_master WHERE type='table' AND name='request_routes'"
        ).fetchone():
            return True
        db.execute("UPDATE request_routes SET revoked=1 WHERE run_id=?", (run_id,))
        if reconcile:
            # The existing operator reconciliation attests provider termination
            # and final accounting. Full child liabilities are retained.
            db.execute(
                "UPDATE request_reservations SET state='charged' WHERE run_id=?",
                (run_id,),
            )
        return not db.execute(
            "SELECT 1 FROM request_reservations WHERE run_id=? AND state='held'",
            (run_id,),
        ).fetchone()

    def reconcile(
        self, run_id: str, *, cost_usd: str, receipt: str, now: datetime | None = None
    ) -> None:
        """Operator attests provider termination + final accounting; no retry grant."""
        cost = amount(cost_usd, positive=False)
        if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._:/@+-]{0,199}", receipt):
            raise SpendAdmissionError("bounded reconciliation receipt required")
        with self._db() as db:
            day = self._clock(db, now or datetime.now(UTC), admission=False)
            row = db.execute(
                "SELECT * FROM reservations WHERE run_id=?", (run_id,)
            ).fetchone()
            if row is None:
                raise SpendAdmissionError("unknown reservation")
            if row["state"] != "held":
                if row["receipt"] == receipt and row["observed_usd"] == str(cost):
                    return
                raise SpendAdmissionError(
                    "reconciliation conflicts with settled reservation"
                )
            if row["observed_usd"] is not None and cost < amount(
                row["observed_usd"], positive=False
            ):
                raise SpendAdmissionError(
                    "reconciliation cannot discard observed liability"
                )
            self._close_requests(db, run_id, reconcile=True)
            liability = max(
                row["liability"], converted_micros(str(cost), self.policy.eur_per_usd)
            )
            db.execute(
                "UPDATE reservations SET state='charged', liability=?, observed_usd=?, end_day=?, receipt=? WHERE run_id=?",
                (liability, str(cost), day, receipt, run_id),
            )
            if cost > amount(self.policy.max_liability_usd):
                db.execute("UPDATE envelope SET breached=1")

    def status(self) -> dict[str, Any]:
        with self._db() as db:
            return {
                "envelope_id": self.policy.envelope_id,
                "policy_sha256": self.policy.sha256,
                "breached": bool(
                    db.execute("SELECT breached FROM envelope").fetchone()[0]
                ),
                "reservations": [
                    dict(row)
                    for row in db.execute(
                        "SELECT * FROM reservations ORDER BY start_day, run_id"
                    )
                ],
            }


def worker_spend(config: Any) -> SpendLedger | None:
    if config.spend_policy_path is None and config.spend_ledger_path is None:
        if config.require_spend_admission:
            raise SpendAdmissionError("required spend admission is not configured")
        return None
    if not config.spend_policy_path or not config.spend_ledger_path:
        raise SpendAdmissionError("both spend policy and ledger must be configured")
    policy = SpendPolicy.load(Path(config.spend_policy_path))
    if (
        config.worker_id != policy.worker_id
        or config.execution_project != policy.project
    ):
        raise SpendAdmissionError("worker identity does not match spend policy")
    return SpendLedger(Path(config.spend_ledger_path), policy)