"""SQLite state store (single writer: the controller). All access under one lock (sqlite + threads gotcha, WP29)."""
from __future__ import annotations
import json, os, sqlite3, threading
from .contracts import utcnow

SCHEMA = """
CREATE TABLE IF NOT EXISTS runs (
  run_id TEXT PRIMARY KEY, task_id TEXT, occupation TEXT, arm TEXT, rep INTEGER, model TEXT, stage TEXT,
  node TEXT, session_id TEXT, status TEXT, assigned_utc TEXT, started_utc TEXT, ended_utc TEXT,
  wall_s REAL, timed_out INTEGER, exit_code INTEGER, result_subtype TEXT, cost_usd REAL, num_turns INTEGER,
  recovery_attempted INTEGER DEFAULT 0, recovery_success INTEGER, recovery_cost_usd REAL, recovery_wall_s REAL,
  selfcheck_pass_first INTEGER, selfcheck_pass_final INTEGER,
  n_events INTEGER, claude_version TEXT, order_index INTEGER, seed INTEGER, notes TEXT
);
CREATE TABLE IF NOT EXISTS events (
  event_id TEXT PRIMARY KEY, run_id TEXT, seq INTEGER, received_utc TEXT, elapsed_s REAL, session_id TEXT,
  event_type TEXT, subtype TEXT, tool_name TEXT, summary TEXT, status TEXT, error_class TEXT, usage TEXT, payload_sha256 TEXT
);
CREATE INDEX IF NOT EXISTS ix_events_run ON events(run_id, seq);
CREATE TABLE IF NOT EXISTS artifacts (run_id TEXT, path TEXT, sha256 TEXT, bytes INTEGER, PRIMARY KEY(run_id, path));
CREATE TABLE IF NOT EXISTS receipts (run_id TEXT PRIMARY KEY, receipt TEXT, evaluated_utc TEXT, det_completion INTEGER, det_score REAL, n_violations INTEGER, abstention_status TEXT, abstention_justified INTEGER);
CREATE TABLE IF NOT EXISTS judgments (run_id TEXT, judge_idx INTEGER, judge_model TEXT, score REAL, scores TEXT, rationale TEXT, cost_usd REAL, judged_utc TEXT, raw TEXT, PRIMARY KEY(run_id, judge_idx));
CREATE TABLE IF NOT EXISTS metrics (run_id TEXT PRIMARY KEY, metrics TEXT);
CREATE TABLE IF NOT EXISTS interventions (id INTEGER PRIMARY KEY AUTOINCREMENT, utc TEXT, run_id TEXT, minutes REAL, description TEXT);
CREATE TABLE IF NOT EXISTS ledger (id INTEGER PRIMARY KEY AUTOINCREMENT, utc TEXT, run_id TEXT, kind TEXT, cost_usd REAL);
"""

class Store:
    def __init__(self, path: str):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        self.path = path; self.lock = threading.RLock()
        self.con = sqlite3.connect(path, check_same_thread=False, timeout=60)
        self.con.row_factory = sqlite3.Row
        with self.lock:
            self.con.executescript(SCHEMA); self.con.commit()
        with self.lock:
            cols = [r[1] for r in self.con.execute("PRAGMA table_info(runs)")]
            if "interrupt_s" not in cols: self.con.execute("ALTER TABLE runs ADD COLUMN interrupt_s INTEGER"); self.con.commit()

    # --- generic helpers ---
    def exec(self, sql, args=()):
        with self.lock:
            cur = self.con.execute(sql, args); self.con.commit(); return cur
    def q(self, sql, args=()):
        with self.lock:
            return [dict(r) for r in self.con.execute(sql, args).fetchall()]
    def one(self, sql, args=()):
        r = self.q(sql, args); return r[0] if r else None

    # --- runs ---
    def add_run(self, **kw):
        kw.setdefault("assigned_utc", utcnow()); kw.setdefault("status", "assigned")
        cols = ",".join(kw); ph = ",".join("?" * len(kw))
        self.exec(f"INSERT OR IGNORE INTO runs ({cols}) VALUES ({ph})", tuple(kw.values()))
    def update_run(self, run_id, **kw):
        sets = ",".join(f"{k}=?" for k in kw)
        self.exec(f"UPDATE runs SET {sets} WHERE run_id=?", (*kw.values(), run_id))
    def get_run(self, run_id):
        return self.one("SELECT * FROM runs WHERE run_id=?", (run_id,))
    def runs(self, where="1=1", args=()):
        return self.q(f"SELECT * FROM runs WHERE {where} ORDER BY order_index", args)

    # --- events ---
    def add_event(self, ev: dict):
        with self.lock:
            self.con.execute("INSERT OR IGNORE INTO events (event_id,run_id,seq,received_utc,elapsed_s,session_id,event_type,subtype,tool_name,summary,status,error_class,usage,payload_sha256) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
                (ev["event_id"], ev["run_id"], ev["seq"], ev["received_utc"], ev["elapsed_s"], ev.get("session_id"), ev.get("event_type"), ev.get("subtype"), ev.get("tool_name"), ev.get("summary"), ev.get("status"), ev.get("error_class"), json.dumps(ev.get("usage")) if ev.get("usage") is not None else None, ev.get("payload_sha256")))
            self.con.commit()
    def events(self, run_id):
        return self.q("SELECT * FROM events WHERE run_id=? ORDER BY seq", (run_id,))

    def add_artifacts(self, run_id, items):
        with self.lock:
            self.con.executemany("INSERT OR REPLACE INTO artifacts VALUES (?,?,?,?)", [(run_id, p, h, b) for p, h, b in items]); self.con.commit()
    def set_receipt(self, run_id, receipt: dict):
        self.exec("INSERT OR REPLACE INTO receipts VALUES (?,?,?,?,?,?,?,?)", (run_id, json.dumps(receipt, default=str), utcnow(), int(bool(receipt.get("det_completion"))), float(receipt.get("det_score") or 0), len(receipt.get("critical_violations") or []), (receipt.get("abstention") or {}).get("status"), None if (receipt.get("abstention") or {}).get("justified") is None else int(bool(receipt["abstention"]["justified"]))))
    def set_judgment(self, run_id, idx, model, score, scores, rationale, cost, raw):
        self.exec("INSERT OR REPLACE INTO judgments VALUES (?,?,?,?,?,?,?,?,?)", (run_id, idx, model, score, json.dumps(scores), rationale, cost, utcnow(), raw[:20000] if raw else None))
    def set_metrics(self, run_id, m: dict):
        self.exec("INSERT OR REPLACE INTO metrics VALUES (?,?)", (run_id, json.dumps(m, default=str)))
    def ledger_add(self, run_id, kind, cost):
        self.exec("INSERT INTO ledger (utc,run_id,kind,cost_usd) VALUES (?,?,?,?)", (utcnow(), run_id, kind, float(cost or 0)))
    def ledger_total(self) -> float:
        r = self.one("SELECT COALESCE(SUM(cost_usd),0) AS t FROM ledger"); return float(r["t"]) if r else 0.0
    def add_intervention(self, run_id, minutes, description):
        self.exec("INSERT INTO interventions (utc,run_id,minutes,description) VALUES (?,?,?,?)", (utcnow(), run_id, minutes, description))