"""Cumulative Oxylabs allowance shared by every probe execution (v2.2.16).

The ledger is one JSON document. On Azure it is the blob probe/cumulative-usage.json; the probe reads and writes that
blob itself with the VM's managed identity (--ledger blob:<url>), so no local copy exists that could go stale or be
uploaded over a newer version. A local file (--ledger <path>) gets the same protocol with a lock file.

Every change is one read-modify-write:
  1. take the exclusive lock (blob: a 30 s lease; file: an O_EXCL lock file), waiting at most `lock_wait` seconds;
  2. read the document and its ETag (file: SHA-256 of the bytes);
  3. write the new document only if the ETag is unchanged (blob: If-Match plus the lease id);
  4. read it back and compare byte for byte. Any difference raises LedgerMismatch: the probe stops, nothing further
     is sent, and the ledger is left for a person to resolve.

Reservation before dispatch: the probe calls reserve() before its first request. The reservation (status "reserved",
reserved_jobs = this run's cap) is durable and verified before anything is sent. finalize() later replaces it, matched
by reservation_id on a fresh read (never on the copy read at the start), with the jobs actually dispatched.

Reconciliation on restart: reserve() first looks for entries still "reserved" (a run that was interrupted, killed, or
whose final update could not be verified). They become "unfinished" and stay counted at their full reservation. This
is never lowered automatically; only a person who has the evidence may correct it.

Any ledger error stops the probe before dispatch (exit 2). A missing blob is an error, never a fresh start.
"""
from __future__ import annotations

import hashlib
import json
import os
import socket
import time
import uuid
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path


class LedgerError(Exception):
    """The ledger could not be read, locked, changed or verified. The probe stops; nothing more is sent."""


class LedgerMismatch(LedgerError):
    """What was read back is not what was written. The probe stops and the ledger stays unresolved."""


def utc() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def jobs_counted(ledger: dict) -> int:
    """Jobs counted so far: dispatched where recorded, else the full reservation. An entry with neither stops."""
    total = 0
    for i, e in enumerate(ledger.get("executions") or []):
        jobs = e.get("oxylabs_jobs_submitted")
        if jobs is None:
            jobs = e.get("reserved_jobs")
        if jobs is None:
            raise LedgerError(f"ledger entry {i} ({e.get('folder')}) has neither dispatched jobs nor a reservation")
        total += int(jobs)
    return total


def encode(data: dict) -> bytes:
    return (json.dumps(data, indent=1, ensure_ascii=False) + "\n").encode("utf-8")


def sha256(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()


def _status(exc: Exception):
    return getattr(exc, "status_code", None) or getattr(getattr(exc, "response", None), "status_code", None)


class FileStore:
    """A local ledger file. The lock is an exclusive lock file next to it; the ETag is the SHA-256 of the bytes."""
    may_create = True

    def __init__(self, path, *, lock_wait: float = 30.0, sleep=time.sleep, clock=time.monotonic):
        self.path = Path(path)
        self.lock_path = self.path.with_name(self.path.name + ".lock")
        self.lock_wait, self.sleep, self.clock = lock_wait, sleep, clock
        self._fd = None

    def describe(self) -> str:
        return str(self.path)

    def lock(self) -> None:
        end = self.clock() + self.lock_wait
        while True:
            try:
                self._fd = os.open(self.lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
                os.write(self._fd, f"{socket.gethostname()} pid {os.getpid()} {utc()}".encode())
                return
            except FileExistsError:
                if self.clock() >= end:
                    raise LedgerError(f"ledger lock {self.lock_path} is held by another execution") from None
                self.sleep(1.0)

    def unlock(self) -> None:
        if self._fd is not None:
            os.close(self._fd)
            self._fd = None
            os.remove(self.lock_path)

    def read(self):
        if not self.path.exists():
            return None, None
        b = self.path.read_bytes()
        return b, sha256(b)

    def write(self, body: bytes, etag) -> str:
        _, current = self.read()
        if current != etag:
            raise LedgerError("ledger changed since it was read (ETag mismatch); nothing written")
        tmp = self.path.with_name(self.path.name + ".tmp")
        tmp.write_bytes(body)
        os.replace(tmp, self.path)
        return sha256(body)

    def read_back(self):
        return self.read()


class BlobStore:
    """The ledger blob. Lock = exclusive blob lease; writes carry the lease id and If-Match on the ETag read."""
    may_create = False

    def __init__(self, blob, *, lease_seconds: int = 30, lock_wait: float = 30.0, sleep=time.sleep, clock=time.monotonic):
        self.blob = blob                 # azure.storage.blob.BlobClient (or a test double with the same calls)
        self.lease_seconds, self.lock_wait, self.sleep, self.clock = lease_seconds, lock_wait, sleep, clock
        self._lease = None

    def describe(self) -> str:
        return getattr(self.blob, "url", "blob")

    def lock(self) -> None:
        end = self.clock() + self.lock_wait
        while True:
            try:
                self._lease = self.blob.acquire_lease(lease_duration=self.lease_seconds)
                return
            except Exception as exc:  # noqa: BLE001
                code = _status(exc)
                if code == 404:
                    raise LedgerError(f"ledger blob not found ({self.describe()}): refusing to start a new ledger") from exc
                if code != 409 or self.clock() >= end:
                    raise LedgerError(f"could not lease the ledger blob: {code or type(exc).__name__}: {exc}") from exc
                self.sleep(2.0)          # 409: another execution holds the lease

    def unlock(self) -> None:
        if self._lease is not None:
            try:
                self._lease.release()
            except Exception:  # noqa: BLE001 - the lease expires on its own within lease_seconds
                pass
            self._lease = None

    def read(self):
        try:
            d = self.blob.download_blob(lease=self._lease)
            return d.readall(), d.properties.etag
        except Exception as exc:  # noqa: BLE001
            raise LedgerError(f"could not read the ledger blob: {_status(exc) or type(exc).__name__}: {exc}") from exc

    def write(self, body: bytes, etag) -> str:
        try:
            from azure.core import MatchConditions
            cond = MatchConditions.IfNotModified
        except ImportError:  # pragma: no cover - the package pins azure-core through azure-storage-blob
            cond = "IfNotModified"
        try:
            r = self.blob.upload_blob(body, overwrite=True, lease=self._lease, etag=etag, match_condition=cond)
            return r["etag"]
        except Exception as exc:  # noqa: BLE001
            code = _status(exc)
            what = "ledger changed since it was read (ETag mismatch)" if code == 412 else "could not write the ledger blob"
            raise LedgerError(f"{what}: {code or type(exc).__name__}: {exc}") from exc

    def read_back(self):
        try:
            d = self.blob.download_blob()
            return d.readall(), d.properties.etag
        except Exception as exc:  # noqa: BLE001
            raise LedgerMismatch(f"ledger written but could not be read back: {exc}") from exc


def open_store(spec: str, **kw):
    """--ledger value: 'blob:https://<account>.blob.core.windows.net/<container>/<name>' or a local path."""
    if spec.startswith("blob:"):
        from azure.identity import ManagedIdentityCredential
        from azure.storage.blob import BlobClient
        cred = ManagedIdentityCredential(client_id=os.environ.get("AZURE_CLIENT_ID") or None)   # None: the VM's identity
        blob = BlobClient.from_blob_url(spec[5:], credential=cred,
                                        connection_timeout=10, read_timeout=15, retry_total=2)
        return BlobStore(blob, **kw)
    return FileStore(spec, **kw)


class Allowance:
    def __init__(self, store, ceiling: int = 15):
        self.store, self.ceiling = store, ceiling
        self.reservation_id: str | None = None
        self.last_sha256 = ""

    @contextmanager
    def _locked(self):
        self.store.lock()
        try:
            yield
        finally:
            self.store.unlock()

    def _load(self):
        body, etag = self.store.read()
        if body is None:
            if not self.store.may_create:
                raise LedgerError(f"ledger not found ({self.store.describe()})")
            return {"ceiling": self.ceiling, "executions": []}, None
        try:
            return json.loads(body.decode("utf-8")), etag
        except ValueError as exc:
            raise LedgerError(f"ledger is not valid JSON: {exc}") from exc

    def _commit(self, data: dict, etag) -> dict:
        body = encode(data)
        written = self.store.write(body, etag)
        back, back_etag = self.store.read_back()
        if back != body:
            raise LedgerMismatch(f"ledger read-back differs: wrote sha256 {sha256(body)}, read "
                                 f"{sha256(back) if back is not None else 'nothing'}; left unresolved")
        if written and back_etag and written != back_etag:
            raise LedgerMismatch(f"ledger ETag after write {written} but read back {back_etag}; left unresolved")
        self.last_sha256 = sha256(body)
        return data

    def reserve(self, folder: str, package: str, wanted: int | None) -> dict:
        """Reconcile unfinished entries, cap this run at what is left, write and verify the reservation."""
        with self._locked():
            data, etag = self._load()
            ceiling = min(int(data.get("ceiling", self.ceiling)), self.ceiling)
            reconciled = []
            for e in data.get("executions") or []:
                if e.get("status") == "reserved" and e.get("oxylabs_jobs_submitted") is None:
                    e["status"] = "unfinished"
                    e["reconciled_utc"] = utc()
                    e["reconciled_note"] = (f"no verified final update; counted at its full reservation of "
                                            f"{e.get('reserved_jobs')} jobs until a person corrects it from evidence")
                    reconciled.append(e.get("folder"))
            used = jobs_counted(data)
            left = ceiling - used
            info = {"ledger": self.store.describe(), "ceiling": ceiling, "counted_before": used, "left_before": left,
                    "reconciled_unfinished": reconciled}
            if left <= 0:
                if reconciled:
                    data["total_jobs_counted"] = used
                    self._commit(data, etag)
                info["cap_this_run"] = 0
                return info
            cap = left if wanted is None else min(int(wanted), left)
            self.reservation_id = uuid.uuid4().hex
            data.setdefault("executions", []).append({
                "folder": folder, "package": package, "host": socket.gethostname(), "reservation_id": self.reservation_id,
                "status": "reserved", "reserved_utc": utc(), "reserved_jobs": cap, "oxylabs_jobs_submitted": None})
            data["ceiling"] = ceiling
            data["total_jobs_counted"] = jobs_counted(data)
            self._commit(data, etag)
            info.update(cap_this_run=cap, reservation_id=self.reservation_id, ledger_sha256_after_reserve=self.last_sha256)
            return info

    def finalize(self, jobs: int, exit_code: int, finished_utc: str) -> dict:
        """Record the jobs actually dispatched on a fresh read, matched by reservation id, and verify it."""
        with self._locked():
            data, etag = self._load()
            entry = next((e for e in data.get("executions") or [] if e.get("reservation_id") == self.reservation_id), None)
            if entry is None:
                raise LedgerError(f"reservation {self.reservation_id} is not in the ledger; left unresolved")
            entry.update(status="final", oxylabs_jobs_submitted=int(jobs), exit_code=exit_code, finished_utc=finished_utc)
            data["total_jobs_counted"] = jobs_counted(data)
            return self._commit(data, etag)