"""SQLite run state, counters, revocation, and consume-once records."""

from __future__ import annotations

import sqlite3
import threading
from collections.abc import Iterator
from contextlib import closing, contextmanager
from datetime import datetime
from pathlib import Path

from .errors import ReplayError, StateError

_SCHEMA = """
PRAGMA journal_mode=WAL;
PRAGMA foreign_keys=ON;

CREATE TABLE IF NOT EXISTS runs (
    run_id TEXT PRIMARY KEY,
    mandate_id TEXT NOT NULL,
    manifest_sha256 TEXT NOT NULL,
    nonce TEXT NOT NULL UNIQUE,
    status TEXT NOT NULL CHECK (status IN ('active', 'closed', 'revoked')),
    started_at TEXT NOT NULL,
    expires_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS counters (
    run_id TEXT NOT NULL REFERENCES runs(run_id) ON DELETE CASCADE,
    rule_id TEXT NOT NULL,
    use_count INTEGER NOT NULL CHECK (use_count >= 0),
    PRIMARY KEY (run_id, rule_id)
);

CREATE TABLE IF NOT EXISTS consumed_tokens (
    jti TEXT PRIMARY KEY,
    run_id TEXT NOT NULL REFERENCES runs(run_id) ON DELETE CASCADE,
    token_kind TEXT NOT NULL,
    consumed_at TEXT NOT NULL
);
"""


class StateStore:
    def __init__(self, path: str | Path):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.RLock()
        with closing(self._connect()) as connection:
            connection.executescript(_SCHEMA)

    def _connect(self) -> sqlite3.Connection:
        connection = sqlite3.connect(self.path, timeout=10, isolation_level=None)
        connection.row_factory = sqlite3.Row
        connection.execute("PRAGMA foreign_keys=ON")
        connection.execute("PRAGMA busy_timeout=10000")
        return connection

    @contextmanager
    def _transaction(self) -> Iterator[sqlite3.Connection]:
        with self._lock:
            connection = self._connect()
            try:
                connection.execute("BEGIN IMMEDIATE")
                yield connection
                connection.execute("COMMIT")
            except Exception:
                connection.execute("ROLLBACK")
                raise
            finally:
                connection.close()

    def start_run(
        self,
        *,
        run_id: str,
        mandate_id: str,
        manifest_sha256: str,
        nonce: str,
        started_at: datetime,
        expires_at: datetime,
    ) -> None:
        try:
            with self._transaction() as connection:
                connection.execute(
                    """
                    INSERT INTO runs (
                        run_id, mandate_id, manifest_sha256, nonce,
                        status, started_at, expires_at
                    ) VALUES (?, ?, ?, ?, 'active', ?, ?)
                    """,
                    (
                        run_id,
                        mandate_id,
                        manifest_sha256,
                        nonce,
                        started_at.isoformat(),
                        expires_at.isoformat(),
                    ),
                )
        except sqlite3.IntegrityError as exc:
            raise ReplayError("run id or signed manifest nonce has already been used") from exc

    def require_active(self, run_id: str, manifest_sha256: str) -> dict[str, str]:
        with closing(self._connect()) as connection:
            row = connection.execute(
                "SELECT * FROM runs WHERE run_id = ?",
                (run_id,),
            ).fetchone()
        if row is None:
            raise StateError("run does not exist")
        if row["manifest_sha256"] != manifest_sha256:
            raise StateError("run attestation is bound to a different manifest")
        if row["status"] != "active":
            raise StateError(f"run is not active: {row['status']}")
        return dict(row)

    def consume_rule_use(self, run_id: str, rule_id: str, maximum: int | None) -> int:
        with self._transaction() as connection:
            row = connection.execute(
                "SELECT status FROM runs WHERE run_id = ?",
                (run_id,),
            ).fetchone()
            if row is None or row["status"] != "active":
                raise StateError("cannot authorize against an inactive run")
            counter = connection.execute(
                "SELECT use_count FROM counters WHERE run_id = ? AND rule_id = ?",
                (run_id, rule_id),
            ).fetchone()
            current = int(counter["use_count"]) if counter is not None else 0
            if maximum is not None and current >= maximum:
                raise StateError("rule use limit exhausted")
            next_value = current + 1
            connection.execute(
                """
                INSERT INTO counters (run_id, rule_id, use_count)
                VALUES (?, ?, ?)
                ON CONFLICT(run_id, rule_id)
                DO UPDATE SET use_count = excluded.use_count
                """,
                (run_id, rule_id, next_value),
            )
            return next_value

    def consume_token(
        self,
        *,
        jti: str,
        run_id: str,
        token_kind: str,
        consumed_at: datetime,
    ) -> None:
        try:
            with self._transaction() as connection:
                row = connection.execute(
                    "SELECT status FROM runs WHERE run_id = ?",
                    (run_id,),
                ).fetchone()
                if row is None or row["status"] != "active":
                    raise StateError("cannot consume a token for an inactive run")
                connection.execute(
                    """
                    INSERT INTO consumed_tokens (jti, run_id, token_kind, consumed_at)
                    VALUES (?, ?, ?, ?)
                    """,
                    (jti, run_id, token_kind, consumed_at.isoformat()),
                )
        except sqlite3.IntegrityError as exc:
            raise ReplayError(f"{token_kind} token has already been consumed") from exc

    def run_status(self, run_id: str) -> str:
        with closing(self._connect()) as connection:
            row = connection.execute("SELECT status FROM runs WHERE run_id=?", (run_id,)).fetchone()
        if row is None:
            raise StateError("run does not exist")
        return str(row["status"])

    def finish_run(self, run_id: str, status: str, *, allow_same: bool = False) -> None:
        if status not in {"closed", "revoked"}:
            raise ValueError("final run status must be closed or revoked")
        with self._transaction() as connection:
            cursor = connection.execute(
                """
                UPDATE runs SET status = ?
                WHERE run_id = ? AND status = 'active'
                """,
                (status, run_id),
            )
            if cursor.rowcount != 1:
                if allow_same:
                    row = connection.execute(
                        "SELECT status FROM runs WHERE run_id=?", (run_id,)
                    ).fetchone()
                    if row is not None and row["status"] == status:
                        return
                raise StateError("run does not exist or is already final")

    def counters(self, run_id: str) -> dict[str, int]:
        with closing(self._connect()) as connection:
            rows = connection.execute(
                "SELECT rule_id, use_count FROM counters WHERE run_id = ?",
                (run_id,),
            ).fetchall()
        return {row["rule_id"]: int(row["use_count"]) for row in rows}
