#!/usr/bin/env python3
"""Local admission experiment, Python 3.9+ with SQLite and thread support.

Run without arguments. Only temporary databases are created. Logical test ticks
are not a production clock. No HTTP, SDK, Dataverse or CRM call is made.
"""
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from tempfile import TemporaryDirectory
from threading import Barrier
import json
import sqlite3


MAX_TICK = 2**63 - 1


def integer(value, minimum=0):
    if type(value) is not int or not minimum <= value <= MAX_TICK:
        raise ValueError("invalid integer")
    return value


def label(value):
    if not isinstance(value, str) or not value.strip():
        raise ValueError("empty label")
    return value


@dataclass(frozen=True)
class Scope:
    environment: str
    identity: str

    def __post_init__(self):
        label(self.environment)
        label(self.identity)

    def key(self):
        return self.environment, self.identity


class Admission:
    """One local ledger. An admitted call keeps its slot until explicit finish."""

    def __init__(self, path, timeout=1.0):
        self.path = str(path)
        self.timeout = timeout
        with self.connection() as db:
            db.executescript("""
                CREATE TABLE IF NOT EXISTS policy (
                    environment TEXT NOT NULL,
                    identity TEXT NOT NULL,
                    capacity INTEGER NOT NULL CHECK(capacity > 0),
                    paused_until INTEGER NOT NULL DEFAULT 0,
                    last_tick INTEGER NOT NULL DEFAULT 0,
                    PRIMARY KEY(environment, identity)
                );
                CREATE TABLE IF NOT EXISTS active (
                    environment TEXT NOT NULL,
                    identity TEXT NOT NULL,
                    call_id TEXT NOT NULL,
                    PRIMARY KEY(environment, identity, call_id),
                    FOREIGN KEY(environment, identity)
                        REFERENCES policy(environment, identity)
                );
            """)

    @contextmanager
    def connection(self):
        db = sqlite3.connect(self.path, timeout=self.timeout, isolation_level=None)
        db.execute("PRAGMA foreign_keys=ON")
        try:
            yield db
        finally:
            if db.in_transaction:
                db.rollback()
            db.close()

    @contextmanager
    def transaction(self):
        with self.connection() as db:
            db.execute("BEGIN IMMEDIATE")
            yield db
            db.commit()

    def register(self, scope, capacity):
        integer(capacity, 1)
        with self.transaction() as db:
            db.execute("INSERT INTO policy(environment,identity,capacity) VALUES(?,?,?)",
                       (*scope.key(), capacity))

    def policy_at(self, db, scope, now):
        integer(now)
        row = db.execute(
            "SELECT capacity,paused_until,last_tick FROM policy "
            "WHERE environment=? AND identity=?", scope.key()
        ).fetchone()
        if row is None:
            raise LookupError("unknown scope")
        if now < row[2]:
            raise ValueError("clock moved backwards")
        db.execute("UPDATE policy SET last_tick=? WHERE environment=? AND identity=?",
                   (now, *scope.key()))
        return row

    def admit(self, scope, call_id, now):
        label(call_id)
        with self.transaction() as db:
            capacity, paused_until, _ = self.policy_at(db, scope, now)
            if db.execute("SELECT 1 FROM active WHERE environment=? AND identity=? "
                          "AND call_id=?", (*scope.key(), call_id)).fetchone():
                return "already_active"
            if now < paused_until:
                return "paused"
            count = db.execute(
                "SELECT COUNT(*) FROM active WHERE environment=? AND identity=?",
                scope.key()
            ).fetchone()[0]
            if count >= capacity:
                return "capacity"
            db.execute("INSERT INTO active VALUES(?,?,?)", (*scope.key(), call_id))
            return "admitted"

    def finish(self, scope, call_id, now, cooldown=0):
        """Caller asserts the call has ended. A timeout alone is not that proof."""
        label(call_id)
        integer(now)
        integer(cooldown)
        integer(now + cooldown)
        with self.transaction() as db:
            self.policy_at(db, scope, now)
            deleted = db.execute(
                "DELETE FROM active WHERE environment=? AND identity=? AND call_id=?",
                (*scope.key(), call_id)
            ).rowcount
            if deleted != 1:
                raise LookupError("unknown active call")
            db.execute(
                "UPDATE policy SET paused_until=MAX(paused_until,?) "
                "WHERE environment=? AND identity=?", (now + cooldown, *scope.key())
            )

    def snapshot(self, scope):
        with self.connection() as db:
            row = db.execute(
                "SELECT capacity,paused_until,last_tick FROM policy "
                "WHERE environment=? AND identity=?", scope.key()
            ).fetchone()
            if row is None:
                raise LookupError("unknown scope")
            calls = [r[0] for r in db.execute(
                "SELECT call_id FROM active WHERE environment=? AND identity=? "
                "ORDER BY call_id", scope.key()
            )]
            return {"capacity": row[0], "paused_until": row[1],
                    "last_tick": row[2], "active": calls}


def contenders(ledger, scope, names, now):
    barrier = Barrier(len(names))

    def enter(name):
        barrier.wait(timeout=10)
        return name, ledger.admit(scope, name, now)

    with ThreadPoolExecutor(max_workers=len(names)) as pool:
        return dict(pool.map(enter, names))


def main():
    checks = []

    def check(name, condition):
        if not condition:
            raise AssertionError(name)
        checks.append(name)

    def rejects(name, error, action):
        try:
            action()
        except error:
            check(name, True)
        else:
            raise AssertionError(name)

    with TemporaryDirectory(prefix="dataverse-admission-") as directory:
        root = Path(directory)
        scope = Scope("fictional-environment", "fictional-application-user")
        independent = []
        for worker in range(4):
            ledger = Admission(root / ("worker-" + str(worker) + ".sqlite"))
            ledger.register(scope, 2)
            independent.extend(ledger.admit(scope, str(i), 0) for i in range(2))
        check("independent_budgets_admit_eight", independent.count("admitted") == 8)

        shared = Admission(root / "shared.sqlite")
        shared.register(scope, 2)
        result = contenders(shared, scope, ["call-" + str(i) for i in range(8)], 0)
        check("shared_budget_admits_two", list(result.values()).count("admitted") == 2)
        check("shared_budget_refuses_six", list(result.values()).count("capacity") == 6)
        initial = shared.snapshot(scope)
        check("ledger_contains_two_active_calls", len(initial["active"]) == 2)
        check("reopened_connection_sees_same_state",
              Admission(root / "shared.sqlite").snapshot(scope) == initial)
        first, second = initial["active"]
        check("same_call_is_not_admitted_twice", shared.admit(scope, first, 0) == "already_active")
        check("full_budget_blocks_another_call", shared.admit(scope, "later", 0) == "capacity")

        # A known finished call and its simulated Retry-After are one transaction.
        shared.finish(scope, first, 10, cooldown=8)
        state = shared.snapshot(scope)
        check("finished_call_releases_exactly_one_slot", state["active"] == [second])
        check("pause_is_published_with_release", state["paused_until"] == 18)
        check("sibling_observes_pause", Admission(root / "shared.sqlite").admit(
            scope, "sibling", 11) == "paused")
        check("existing_call_is_not_cancelled", shared.snapshot(scope)["active"] == [second])
        shared.finish(scope, second, 12, cooldown=2)
        check("shorter_delay_does_not_shorten_pause", shared.snapshot(scope)["paused_until"] == 18)
        check("one_tick_before_deadline_is_paused", shared.admit(scope, "later", 17) == "paused")
        check("deadline_allows_admission", shared.admit(scope, "later", 18) == "admitted")
        before = shared.snapshot(scope)
        rejects("unknown_finish_is_refused", LookupError,
                lambda: shared.finish(scope, "unknown", 18, cooldown=100))
        check("unknown_finish_does_not_release_or_pause", shared.snapshot(scope) == before)
        rejects("clock_regression_is_refused", ValueError,
                lambda: shared.admit(scope, "stale", 17))
        check("clock_regression_changes_nothing", shared.snapshot(scope) == before)

        other_environment = Scope("other-environment", scope.identity)
        other_identity = Scope(scope.environment, "other-application-user")
        shared.register(other_environment, 1)
        shared.register(other_identity, 1)
        check("different_environment_has_separate_policy",
              shared.admit(other_environment, "one", 0) == "admitted")
        check("different_identity_has_separate_policy",
              shared.admit(other_identity, "one", 0) == "admitted")

        # No lease expiry is implemented: an unknown outcome retains its slot.
        orphan_scope = Scope("unknown-outcome", "fictional-user")
        shared.register(orphan_scope, 1)
        shared.admit(orphan_scope, "still-unconfirmed", 0)
        check("time_alone_does_not_release_unknown_call",
              shared.admit(orphan_scope, "replacement", 1000) == "capacity")

        racing = Scope("simultaneous-delays", "fictional-user")
        shared.register(racing, 2)
        shared.admit(racing, "short", 0)
        shared.admit(racing, "long", 0)
        barrier = Barrier(2)

        def finish(item):
            name, duration = item
            barrier.wait(timeout=10)
            shared.finish(racing, name, 10, cooldown=duration)

        with ThreadPoolExecutor(max_workers=2) as pool:
            list(pool.map(finish, [("short", 8), ("long", 20)]))
        check("concurrent_delays_keep_latest_deadline",
              shared.snapshot(racing)["paused_until"] == 30)
        check("both_known_finished_calls_are_released", shared.snapshot(racing)["active"] == [])
        check("latest_deadline_still_blocks", shared.admit(racing, "after", 29) == "paused")
        check("latest_deadline_opens", shared.admit(racing, "after", 30) == "admitted")

        last_slot = Scope("last-slot", "fictional-user")
        shared.register(last_slot, 2)
        shared.admit(last_slot, "held", 0)
        race = contenders(shared, last_slot, ["left", "right"], 0)
        check("last_slot_has_one_winner", list(race.values()).count("admitted") == 1)
        check("last_slot_has_one_refusal", list(race.values()).count("capacity") == 1)
        check("last_slot_never_exceeds_capacity", len(shared.snapshot(last_slot)["active"]) == 2)

        before = shared.snapshot(scope)
        rejects("negative_clock_refused", ValueError, lambda: shared.admit(scope, "bad", -1))
        rejects("boolean_clock_refused", ValueError, lambda: shared.admit(scope, "bad", True))
        rejects("fractional_clock_refused", ValueError, lambda: shared.admit(scope, "bad", 18.5))
        rejects("negative_delay_refused", ValueError, lambda: shared.finish(scope, "later", 18, -1))
        rejects("deadline_overflow_refused", ValueError,
                lambda: shared.finish(scope, "later", MAX_TICK, 1))
        rejects("empty_call_refused", ValueError, lambda: shared.admit(scope, "", 18))
        rejects("zero_capacity_refused", ValueError,
                lambda: shared.register(Scope("invalid-capacity", "user"), 0))
        rejects("empty_identity_refused", ValueError, lambda: Scope("environment", " "))
        rejects("unknown_scope_refused", LookupError,
                lambda: shared.admit(Scope("missing", "user"), "call", 0))
        check("invalid_input_does_not_change_known_policy", shared.snapshot(scope) == before)

        # An exclusive writer prevents admission; storage failure is not permission.
        impatient = Admission(root / "shared.sqlite", timeout=0.01)
        with shared.connection() as lock:
            lock.execute("BEGIN IMMEDIATE")
            rejects("busy_coordinator_does_not_admit", sqlite3.OperationalError,
                    lambda: impatient.admit(scope, "bypass", 18))
            lock.rollback()
        check("busy_failure_does_not_insert_a_permit", shared.snapshot(scope) == before)
        check("schema_still_exists_after_failures", shared.snapshot(last_slot)["capacity"] == 2)

        print(json.dumps({
            "checks": len(checks), "passed": checks,
            "independent_admissions": 8, "shared_admissions": 2,
            "shared_pause_until": 18, "concurrent_pause_until": 30,
            "http_calls": 0, "dataverse_requests": 0, "crm_writes": 0,
            "local_sqlite_writes": True,
            "limits": "Local policy and logical clock only; no native throughput or CRM effect proof."
        }, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()