from __future__ import annotations

import json
import sqlite3
from typing import Any

from .models import CheckpointIncompatible


SCHEMA_VERSION = "4"
DEFAULT_RUN = "__default__"
SCHEMA = """
PRAGMA foreign_keys=ON;
CREATE TABLE IF NOT EXISTS schema_meta(key TEXT PRIMARY KEY, value TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS cases(case_id TEXT NOT NULL, run_id TEXT NOT NULL, state TEXT NOT NULL, target_state TEXT NOT NULL, PRIMARY KEY(case_id,run_id));
CREATE TABLE IF NOT EXISTS transitions(id INTEGER PRIMARY KEY, run_id TEXT NOT NULL, case_id TEXT NOT NULL, from_state TEXT, to_state TEXT NOT NULL, reason TEXT NOT NULL, created_at REAL NOT NULL);
CREATE TABLE IF NOT EXISTS journal(id INTEGER PRIMARY KEY, run_id TEXT NOT NULL, case_id TEXT NOT NULL, request_id TEXT, tool TEXT, phase TEXT NOT NULL, detail TEXT NOT NULL, created_at REAL NOT NULL);
CREATE TABLE IF NOT EXISTS effects(id INTEGER PRIMARY KEY, run_id TEXT NOT NULL, case_id TEXT NOT NULL, tool TEXT NOT NULL, idem_key TEXT NOT NULL, args_hash TEXT NOT NULL, result TEXT NOT NULL, UNIQUE(run_id,case_id,idem_key));
CREATE TABLE IF NOT EXISTS approvals(id TEXT PRIMARY KEY, run_id TEXT NOT NULL, case_id TEXT NOT NULL, tool TEXT NOT NULL, args_hash TEXT NOT NULL, status TEXT NOT NULL, actor TEXT, requested_at REAL NOT NULL, decided_at REAL, expires_at REAL, policy_version TEXT NOT NULL, used INTEGER NOT NULL DEFAULT 0);
CREATE TABLE IF NOT EXISTS checkpoints(run_id TEXT NOT NULL, case_id TEXT NOT NULL, payload TEXT NOT NULL, PRIMARY KEY(run_id,case_id));
CREATE TABLE IF NOT EXISTS events(run_id TEXT NOT NULL, case_id TEXT NOT NULL, event_id TEXT NOT NULL, result TEXT NOT NULL, PRIMARY KEY(run_id,case_id,event_id));
CREATE TABLE IF NOT EXISTS decisions(id INTEGER PRIMARY KEY, run_id TEXT NOT NULL, case_id TEXT NOT NULL, payload TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS business_state(case_id TEXT NOT NULL, run_id TEXT NOT NULL, target TEXT NOT NULL, value TEXT NOT NULL, version INTEGER NOT NULL, update_count INTEGER NOT NULL, PRIMARY KEY(case_id,run_id,target));
CREATE TABLE IF NOT EXISTS completion_tokens(token TEXT PRIMARY KEY, run_id TEXT NOT NULL, case_id TEXT NOT NULL, evaluation_digest TEXT NOT NULL, used INTEGER NOT NULL DEFAULT 0);
CREATE TABLE IF NOT EXISTS safety_events(id INTEGER PRIMARY KEY, run_id TEXT NOT NULL, case_id TEXT NOT NULL, category TEXT NOT NULL, detail TEXT NOT NULL, created_at REAL NOT NULL);
CREATE TABLE IF NOT EXISTS run_rejections(id INTEGER PRIMARY KEY, run_id TEXT NOT NULL, case_id TEXT NOT NULL, reason TEXT NOT NULL, detail TEXT NOT NULL, created_at REAL NOT NULL);
"""


class Store:
    def __init__(self, path: str = ":memory:") -> None:
        self.path = path
        self._closed = False
        self.db = sqlite3.connect(path, timeout=10)
        self.db.row_factory = sqlite3.Row
        existing = {row[0] for row in self.db.execute("SELECT name FROM sqlite_master WHERE type='table'")}
        if "cases" in existing and "schema_meta" not in existing:
            self.db.close()
            raise CheckpointIncompatible("legacy store schema requires explicit migration")
        self.db.executescript(SCHEMA)
        row = self.db.execute("SELECT value FROM schema_meta WHERE key='schema_version'").fetchone()
        if row is None:
            self.db.execute("INSERT INTO schema_meta VALUES('schema_version',?)", (SCHEMA_VERSION,))
        elif row[0] != SCHEMA_VERSION:
            self.db.close()
            raise CheckpointIncompatible("store schema version mismatch")
        self.db.commit()

    def __enter__(self) -> "Store":
        return self

    def __exit__(self, exc_type, exc, traceback) -> None:
        self.close()

    def close(self) -> None:
        if self._closed:
            return
        self.db.close()
        self._closed = True

    @property
    def closed(self) -> bool:
        return self._closed

    def __del__(self) -> None:
        try:
            self.close()
        except Exception:
            pass

    def create_case(self, case_id: str, state: str, target_state: str = "COMPLETED", run_id: str | None = None) -> None:
        run = run_id or DEFAULT_RUN
        self.db.execute("INSERT INTO cases VALUES(?,?,?,?)", (case_id, run, state, target_state))
        self.db.commit()

    def _resolve_run(self, case_id: str, run_id: str | None) -> str:
        if run_id:
            return run_id
        rows = self.db.execute("SELECT run_id FROM cases WHERE case_id=? ORDER BY rowid DESC", (case_id,)).fetchall()
        if len(rows) != 1:
            if not rows:
                raise KeyError(case_id)
            raise ValueError("run_id required for ambiguous case")
        return str(rows[0][0])

    def state(self, case_id: str, run_id: str | None = None) -> str:
        run = self._resolve_run(case_id, run_id)
        row = self.db.execute("SELECT state FROM cases WHERE case_id=? AND run_id=?", (case_id, run)).fetchone()
        if not row:
            raise KeyError((case_id, run))
        return str(row[0])

    def set_state(self, case_id: str, old: str, new: str, reason: str, run_id: str | None = None, created_at: float = 0.0) -> None:
        run = self._resolve_run(case_id, run_id)
        def apply() -> None:
            cursor = self.db.execute("UPDATE cases SET state=? WHERE case_id=? AND run_id=? AND state=?", (new, case_id, run, old))
            if cursor.rowcount != 1:
                raise ValueError("concurrent state mismatch")
            self.db.execute("INSERT INTO transitions(run_id,case_id,from_state,to_state,reason,created_at) VALUES(?,?,?,?,?,?)", (run, case_id, old, new, reason, created_at))
        if self.db.in_transaction:
            apply()
        else:
            with self.db:
                apply()

    def log(self, run_id: str, case_id: str, phase: str, detail: dict[str, Any], request_id: str | None = None, tool: str | None = None, created_at: float = 0.0) -> None:
        clean = json.dumps(detail, sort_keys=True, ensure_ascii=True)
        self.db.execute("INSERT INTO journal(run_id,case_id,request_id,tool,phase,detail,created_at) VALUES(?,?,?,?,?,?,?)", (run_id, case_id, request_id, tool, phase, clean, created_at))
        self.db.commit()

    def effect(self, idem_key: str, run_id: str | None = None, case_id: str | None = None) -> sqlite3.Row | None:
        if run_id and case_id:
            return self.db.execute("SELECT * FROM effects WHERE run_id=? AND case_id=? AND idem_key=?", (run_id, case_id, idem_key)).fetchone()
        if run_id:
            raise ValueError("case_id required with run_id for effect lookup")
        return self.db.execute("SELECT * FROM effects WHERE idem_key=?", (idem_key,)).fetchone()

    def add_effect(self, case_id: str, tool: str, idem_key: str, args_hash: str, result: dict[str, Any], run_id: str | None = None) -> None:
        run = run_id or self._resolve_run(case_id, None)
        self.db.execute("INSERT INTO effects(run_id,case_id,tool,idem_key,args_hash,result) VALUES(?,?,?,?,?,?)", (run, case_id, tool, idem_key, args_hash, json.dumps(result, sort_keys=True)))
        self.db.commit()

    def effect_count(self, case_id: str | None = None, run_id: str | None = None) -> int:
        query, params = "SELECT COUNT(*) FROM effects WHERE 1=1", []
        if case_id: query += " AND case_id=?"; params.append(case_id)
        if run_id: query += " AND run_id=?"; params.append(run_id)
        return int(self.db.execute(query, params).fetchone()[0])

    def duplicate_effect_count(self, case_id: str | None = None, run_id: str | None = None) -> int:
        query, params = "SELECT COUNT(*) FROM (SELECT run_id,idem_key,COUNT(*) n FROM effects WHERE 1=1", []
        if case_id: query += " AND case_id=?"; params.append(case_id)
        if run_id: query += " AND run_id=?"; params.append(run_id)
        query += " GROUP BY run_id,idem_key HAVING n>1)"
        return int(self.db.execute(query, params).fetchone()[0])

    def request_approval(self, record: dict[str, Any]) -> None:
        self.db.execute("INSERT INTO approvals(id,run_id,case_id,tool,args_hash,status,actor,requested_at,decided_at,expires_at,policy_version,used) VALUES(?,?,?,?,?,'PENDING',NULL,?,NULL,NULL,?,0)", (record["id"], record["run_id"], record["case_id"], record["tool"], record["args_hash"], record["requested_at"], record["policy_version"]))
        self.db.commit()

    def add_approval(self, record: dict[str, Any]) -> None:
        self.request_approval({**record, "run_id": record.get("run_id", DEFAULT_RUN), "requested_at": record.get("issued_at", 0.0)})
        self.decide_approval(record["id"], "APPROVED", record.get("approver", "external"), record.get("issued_at", 0.0), record.get("expires_at", 10.0))

    def decide_approval(self, approval_id: str, status: str, actor: str, decided_at: float, expires_at: float | None) -> None:
        if status not in {"APPROVED", "REJECTED"}:
            raise ValueError("invalid external decision")
        with self.db:
            cursor = self.db.execute("UPDATE approvals SET status=?,actor=?,decided_at=?,expires_at=? WHERE id=? AND status='PENDING'", (status, actor, decided_at, expires_at, approval_id))
            if cursor.rowcount != 1:
                raise ValueError("approval request is not pending")

    def approval(self, approval_id: str, case_id: str, run_id: str) -> sqlite3.Row | None:
        return self.db.execute("SELECT * FROM approvals WHERE id=? AND case_id=? AND run_id=?", (approval_id, case_id, run_id)).fetchone()

    def mark_approval_used(self, approval_id: str) -> None:
        with self.db:
            cursor = self.db.execute("UPDATE approvals SET used=1 WHERE id=? AND status='APPROVED' AND used=0", (approval_id,))
            if cursor.rowcount != 1:
                raise ValueError("approval is not single-use available")

    def save_checkpoint(self, run_id: str, case_id: str, payload: dict[str, Any]) -> None:
        self.db.execute("INSERT OR REPLACE INTO checkpoints VALUES(?,?,?)", (run_id, case_id, json.dumps(payload, sort_keys=True)))
        self.db.commit()

    def checkpoint(self, run_id: str, case_id: str) -> dict[str, Any] | None:
        row = self.db.execute("SELECT payload FROM checkpoints WHERE run_id=? AND case_id=?", (run_id, case_id)).fetchone()
        return json.loads(row[0]) if row else None

    def claim_event(self, event_id: str, result: dict[str, Any], run_id: str, case_id: str) -> bool:
        try:
            self.db.execute("INSERT INTO events VALUES(?,?,?,?)", (run_id, case_id, event_id, json.dumps(result, sort_keys=True)))
            self.db.commit()
            return True
        except sqlite3.IntegrityError:
            return False

    def add_decision(self, run_id: str, payload: dict[str, Any], case_id: str = "unknown") -> None:
        allowed = {"action", "evidence", "policy", "verification", "next_state", "failure_class"}
        if set(payload) != allowed:
            raise ValueError("decision log schema mismatch")
        self.db.execute("INSERT INTO decisions(run_id,case_id,payload) VALUES(?,?,?)", (run_id, case_id, json.dumps(payload, sort_keys=True)))
        self.db.commit()

    def init_business(self, case_id: str, run_id: str, target: str, value: str = "PENDING") -> None:
        self.db.execute("INSERT INTO business_state VALUES(?,?,?,?,0,0)", (case_id, run_id, target, value))
        self.db.commit()

    def business(self, case_id: str, run_id: str, target: str) -> dict[str, Any] | None:
        row = self.db.execute("SELECT * FROM business_state WHERE case_id=? AND run_id=? AND target=?", (case_id, run_id, target)).fetchone()
        return dict(row) if row else None

    def apply_business_update(self, case_id: str, run_id: str, target: str, value: str, expected_version: int, tool: str, idem_key: str, args_hash: str, result: dict[str, Any], approval_id: str) -> None:
        with self.db:
            cursor = self.db.execute("UPDATE business_state SET value=?,version=version+1,update_count=update_count+1 WHERE case_id=? AND run_id=? AND target=? AND version=?", (value, case_id, run_id, target, expected_version))
            if cursor.rowcount != 1:
                raise ValueError("business state version mismatch")
            self.db.execute("INSERT INTO effects(run_id,case_id,tool,idem_key,args_hash,result) VALUES(?,?,?,?,?,?)", (run_id, case_id, tool, idem_key, args_hash, json.dumps(result, sort_keys=True)))
            used = self.db.execute("UPDATE approvals SET used=1 WHERE id=? AND run_id=? AND case_id=? AND status='APPROVED' AND used=0", (approval_id, run_id, case_id))
            if used.rowcount != 1:
                raise ValueError("approval could not be consumed atomically")

    def issue_completion_token(self, token: str, run_id: str, case_id: str, evaluation_digest: str) -> None:
        self.db.execute("INSERT INTO completion_tokens VALUES(?,?,?,?,0)", (token, run_id, case_id, evaluation_digest))
        self.db.commit()

    def consume_completion_token(self, token: str, run_id: str, case_id: str, evaluation_digest: str) -> None:
        with self.db:
            cursor = self.db.execute("UPDATE completion_tokens SET used=1 WHERE token=? AND run_id=? AND case_id=? AND evaluation_digest=? AND used=0", (token, run_id, case_id, evaluation_digest))
            if cursor.rowcount != 1:
                raise ValueError("completion token invalid or already used")

    def complete_with_token(self, token: str, run_id: str, case_id: str, evaluation_digest: str, state_mutator) -> None:
        """Issue, consume and transition in one SQLite transaction.

        ``state_mutator`` intentionally routes through the runtime's normal CAS
        method so a CAS fault injection exercises the same rollback boundary.
        """
        with self.db:
            self.db.execute("INSERT INTO completion_tokens VALUES(?,?,?,?,0)", (token, run_id, case_id, evaluation_digest))
            state_mutator()
            cursor = self.db.execute("UPDATE completion_tokens SET used=1 WHERE token=? AND run_id=? AND case_id=? AND evaluation_digest=? AND used=0", (token, run_id, case_id, evaluation_digest))
            if cursor.rowcount != 1:
                raise ValueError("completion token invalid or already used")

    def record_run_rejection(self, run_id: str, case_id: str, reason: str, detail: dict[str, Any], created_at: float) -> None:
        self.db.execute("INSERT INTO run_rejections(run_id,case_id,reason,detail,created_at) VALUES(?,?,?,?,?)", (run_id, case_id, reason, json.dumps(detail, sort_keys=True), created_at))
        self.db.commit()

    def safety_block(self, run_id: str, case_id: str, category: str, detail: dict[str, Any], created_at: float) -> None:
        self.db.execute("INSERT INTO safety_events(run_id,case_id,category,detail,created_at) VALUES(?,?,?,?,?)", (run_id, case_id, category, json.dumps(detail, sort_keys=True), created_at))
        self.db.commit()

    def rows(self, table: str) -> list[dict[str, Any]]:
        if table not in {"transitions", "journal", "effects", "approvals", "decisions", "business_state", "completion_tokens", "safety_events", "run_rejections"}:
            raise ValueError("table not allowlisted")
        return [dict(row) for row in self.db.execute(f"SELECT * FROM {table}")]

    def rows_for_case(self, table: str, case_id: str, run_id: str) -> list[dict[str, Any]]:
        if table not in {"transitions", "journal", "effects", "approvals", "decisions", "business_state", "completion_tokens", "safety_events", "run_rejections"}:
            raise ValueError("table not allowlisted")
        return [dict(row) for row in self.db.execute(f"SELECT * FROM {table} WHERE case_id=? AND run_id=?", (case_id, run_id))]
