"""
ChainStore — SQLite persistence layer for llm-guard-kit.

Stores every tracked chain with risk scores, supports auto-calibration
pool retrieval, audit log export, and drift window queries.

Usage:
    from qppg_service.store import ChainStore

    store = ChainStore()                           # ~/.qppg/chains.db
    store = ChainStore("/var/data/qppg.db")        # custom path

    row_id = store.add_chain("prod", chain_dict, risk_score=0.72,
                              alert=True, failure_mode="EXCESSIVE_SEARCH")
    chains = store.get_calibration_pool("prod", n=200)
    risks   = store.get_recent_risk("prod", days=7)
    csv_str = store.export_audit("prod", start=0, end=time.time())
"""

from __future__ import annotations

import base64
import csv
import hashlib
import io
import json
import os
import sqlite3
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Generator


# ── Fernet encryption helpers for secrets (LLM API keys) ──────────────────────

def _get_cipher():
    """Return a Fernet cipher if QPPG_SECRET_KEY is set, else None (no encryption)."""
    try:
        from cryptography.fernet import Fernet
        secret = os.environ.get("QPPG_SECRET_KEY", "")
        if not secret:
            return None
        key = base64.urlsafe_b64encode(hashlib.sha256(secret.encode()).digest())
        return Fernet(key)
    except ImportError:
        return None


def _encrypt(v: str) -> str:
    cipher = _get_cipher()
    if cipher and v:
        return "enc:" + cipher.encrypt(v.encode()).decode()
    return v


def _decrypt(v: str) -> str:
    if not v or not v.startswith("enc:"):
        return v
    cipher = _get_cipher()
    if cipher:
        try:
            return cipher.decrypt(v[4:].encode()).decode()
        except Exception:
            return v
    return v


_DEFAULT_DB = Path.home() / ".qppg" / "chains.db"

_SCHEMA = """
CREATE TABLE IF NOT EXISTS chains (
    id               INTEGER PRIMARY KEY AUTOINCREMENT,
    domain           TEXT    NOT NULL,
    question         TEXT    DEFAULT '',
    steps_json       TEXT    DEFAULT '[]',
    final_answer     TEXT    DEFAULT '',
    finished         INTEGER DEFAULT 1,
    risk_score       REAL    DEFAULT 0.0,
    alert_triggered  INTEGER DEFAULT 0,
    failure_mode     TEXT    DEFAULT '',
    n_steps          INTEGER DEFAULT 0,
    model_name       TEXT    DEFAULT '',
    timestamp        REAL    NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_domain_ts ON chains(domain, timestamp);

CREATE TABLE IF NOT EXISTS calibration_log (
    id             INTEGER PRIMARY KEY AUTOINCREMENT,
    domain         TEXT  NOT NULL,
    n_chains       INTEGER,
    auroc_estimate REAL,
    timestamp      REAL  NOT NULL
);

CREATE TABLE IF NOT EXISTS api_keys (
    key_hash      TEXT    PRIMARY KEY,
    customer_id   TEXT    NOT NULL,
    domain_prefix TEXT    NOT NULL DEFAULT '',
    created_at    REAL    NOT NULL,
    is_active     INTEGER DEFAULT 1,
    request_count INTEGER DEFAULT 0
);

CREATE TABLE IF NOT EXISTS users (
    id                     INTEGER PRIMARY KEY AUTOINCREMENT,
    email                  TEXT    UNIQUE NOT NULL,
    password_hash          TEXT    DEFAULT '',
    google_id              TEXT    UNIQUE,
    name                   TEXT    DEFAULT '',
    plan                   TEXT    DEFAULT 'free',
    stripe_customer_id     TEXT    DEFAULT '',
    stripe_subscription_id TEXT    DEFAULT '',
    chains_this_month      INTEGER DEFAULT 0,
    month_reset_at         REAL    NOT NULL,
    created_at             REAL    NOT NULL,
    is_active              INTEGER DEFAULT 1
);

CREATE TABLE IF NOT EXISTS usage_events (
    id         INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id    INTEGER NOT NULL,
    domain     TEXT    NOT NULL,
    chain_id   INTEGER,
    event_type TEXT    NOT NULL,
    timestamp  REAL    NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_usage_user ON usage_events(user_id, timestamp);

CREATE TABLE IF NOT EXISTS workflows (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id     INTEGER NOT NULL,
    name        TEXT    NOT NULL DEFAULT 'Untitled workflow',
    nodes_json  TEXT    NOT NULL DEFAULT '[]',
    edges_json  TEXT    NOT NULL DEFAULT '[]',
    enabled     INTEGER DEFAULT 1,
    created_at  REAL    NOT NULL,
    updated_at  REAL    NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_workflows_user ON workflows(user_id);

CREATE TABLE IF NOT EXISTS guardrails (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id         INTEGER NOT NULL,
    domain          TEXT    NOT NULL,
    threshold       REAL    NOT NULL DEFAULT 0.65,
    action          TEXT    NOT NULL DEFAULT 'alert',
    webhook_url     TEXT    DEFAULT '',
    slack_webhook   TEXT    DEFAULT '',
    enabled         INTEGER DEFAULT 1,
    created_at      REAL    NOT NULL,
    updated_at      REAL    NOT NULL,
    UNIQUE(user_id, domain)
);

CREATE TABLE IF NOT EXISTS pipelines (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id     INTEGER NOT NULL,
    name        TEXT    NOT NULL DEFAULT 'Untitled pipeline',
    nodes_json  TEXT    NOT NULL DEFAULT '[]',
    edges_json  TEXT    NOT NULL DEFAULT '[]',
    enabled     INTEGER DEFAULT 1,
    created_at  REAL    NOT NULL,
    updated_at  REAL    NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_pipelines_user ON pipelines(user_id);

CREATE TABLE IF NOT EXISTS pipeline_runs (
    id               INTEGER PRIMARY KEY AUTOINCREMENT,
    pipeline_id      INTEGER NOT NULL,
    user_id          INTEGER NOT NULL,
    query            TEXT,
    result           TEXT,
    trace_json       TEXT    NOT NULL DEFAULT '[]',
    risk_score       REAL,
    guardrail_action TEXT,
    status           TEXT    NOT NULL DEFAULT 'pending',
    started_at       REAL    NOT NULL,
    finished_at      REAL
);
CREATE INDEX IF NOT EXISTS idx_pipeline_runs ON pipeline_runs(pipeline_id, user_id);

CREATE TABLE IF NOT EXISTS user_configs (
    user_id     INTEGER PRIMARY KEY,
    config_json TEXT    NOT NULL DEFAULT '{}'
);

CREATE TABLE IF NOT EXISTS pipeline_versions (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    pipeline_id INTEGER NOT NULL,
    user_id     INTEGER NOT NULL,
    version     INTEGER NOT NULL DEFAULT 1,
    name        TEXT    NOT NULL DEFAULT '',
    nodes_json  TEXT    NOT NULL DEFAULT '[]',
    edges_json  TEXT    NOT NULL DEFAULT '[]',
    note        TEXT    DEFAULT '',
    created_at  REAL    NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_pipeline_versions ON pipeline_versions(pipeline_id, version DESC);

CREATE TABLE IF NOT EXISTS ab_comparisons (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id         INTEGER NOT NULL,
    pipeline_a_id   INTEGER NOT NULL,
    pipeline_b_id   INTEGER NOT NULL,
    query           TEXT    NOT NULL,
    result_a        TEXT    DEFAULT '',
    result_b        TEXT    DEFAULT '',
    trace_a_json    TEXT    NOT NULL DEFAULT '[]',
    trace_b_json    TEXT    NOT NULL DEFAULT '[]',
    risk_score_a    REAL,
    risk_score_b    REAL,
    status          TEXT    NOT NULL DEFAULT 'pending',
    created_at      REAL    NOT NULL,
    finished_at     REAL
);
CREATE INDEX IF NOT EXISTS idx_ab_user ON ab_comparisons(user_id, created_at DESC);

CREATE TABLE IF NOT EXISTS orgs (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    name        TEXT    NOT NULL,
    owner_id    INTEGER NOT NULL,
    created_at  REAL    NOT NULL
);
CREATE TABLE IF NOT EXISTS org_members (
    org_id      INTEGER NOT NULL,
    user_id     INTEGER NOT NULL,
    role        TEXT    NOT NULL DEFAULT 'member',
    joined_at   REAL    NOT NULL,
    PRIMARY KEY (org_id, user_id)
);
CREATE TABLE IF NOT EXISTS org_invites (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    org_id      INTEGER NOT NULL,
    email       TEXT    NOT NULL,
    token       TEXT    UNIQUE NOT NULL,
    created_at  REAL    NOT NULL,
    expires_at  REAL    NOT NULL
);

CREATE TABLE IF NOT EXISTS test_sets (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id     INTEGER NOT NULL,
    name        TEXT    NOT NULL DEFAULT 'Untitled test set',
    cases_json  TEXT    NOT NULL DEFAULT '[]',
    pipeline_id INTEGER,
    created_at  REAL    NOT NULL,
    updated_at  REAL    NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_test_sets_user ON test_sets(user_id);

CREATE TABLE IF NOT EXISTS test_runs (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    test_set_id     INTEGER NOT NULL,
    user_id         INTEGER NOT NULL,
    pipeline_id     INTEGER NOT NULL,
    results_json    TEXT    NOT NULL DEFAULT '[]',
    avg_risk_score  REAL,
    pass_rate       REAL,
    status          TEXT    NOT NULL DEFAULT 'pending',
    started_at      REAL    NOT NULL,
    finished_at     REAL
);
CREATE INDEX IF NOT EXISTS idx_test_runs_set ON test_runs(test_set_id, user_id);

CREATE TABLE IF NOT EXISTS proxy_requests (
    id                  INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id             INTEGER NOT NULL,
    provider            TEXT    NOT NULL,
    model               TEXT    DEFAULT '',
    prompt_tokens       INTEGER DEFAULT 0,
    completion_tokens   INTEGER DEFAULT 0,
    risk_score          REAL,
    latency_ms          INTEGER DEFAULT 0,
    status              TEXT    NOT NULL DEFAULT 'ok',
    created_at          REAL    NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_proxy_user ON proxy_requests(user_id, created_at DESC);

CREATE TABLE IF NOT EXISTS automation_log (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id         INTEGER NOT NULL,
    trigger_type    TEXT    NOT NULL DEFAULT 'chain_scored',
    domain          TEXT    NOT NULL DEFAULT '',
    risk_score      REAL,
    chain_id        INTEGER,
    workflow_id     INTEGER,
    guardrail_id    INTEGER,
    pipeline_run_id INTEGER,
    action          TEXT    NOT NULL DEFAULT '',
    result          TEXT    NOT NULL DEFAULT '',
    error           TEXT    DEFAULT '',
    fired_at        REAL    NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_auto_log_user ON automation_log(user_id, fired_at DESC);
"""


class ChainStore:
    """
    SQLite-backed persistence layer for agent chain history.

    Thread-safe: each call acquires its own connection via check_same_thread=False.
    """

    def __init__(self, db_path: str | Path | None = None):
        self.db_path = Path(db_path or _DEFAULT_DB).expanduser()
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        # Initialise schema once
        with self._conn() as conn:
            conn.executescript(_SCHEMA)
            # Migrate: add review columns added in v0.3.1 (no-op if already present)
            for _col_sql in [
                "ALTER TABLE chains ADD COLUMN reviewed INTEGER DEFAULT 0",
                "ALTER TABLE chains ADD COLUMN label    TEXT    DEFAULT ''",
                "ALTER TABLE api_keys ADD COLUMN user_id INTEGER DEFAULT NULL",
                "ALTER TABLE pipelines ADD COLUMN org_id INTEGER DEFAULT NULL",
                "ALTER TABLE workflows ADD COLUMN org_id INTEGER DEFAULT NULL",
                "ALTER TABLE guardrails ADD COLUMN org_id INTEGER DEFAULT NULL",
                "ALTER TABLE guardrails ADD COLUMN email_to TEXT DEFAULT ''",
                "ALTER TABLE pipeline_runs ADD COLUMN version_id INTEGER DEFAULT NULL",
            ]:
                try:
                    conn.execute(_col_sql)
                except Exception:
                    pass  # column already exists

    # ── Context manager ────────────────────────────────────────────────────────

    @contextmanager
    def _conn(self) -> Generator[sqlite3.Connection, None, None]:
        conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")       # concurrent reads + one writer
        conn.execute("PRAGMA synchronous=NORMAL")     # safe with WAL, faster than FULL
        try:
            yield conn
            conn.commit()
        finally:
            conn.close()

    # ── Write ──────────────────────────────────────────────────────────────────

    def add_chain(
        self,
        domain: str,
        chain: dict,
        risk_score: float = 0.0,
        alert: bool = False,
        failure_mode: str = "",
        model_name: str = "",
    ) -> int:
        """
        Persist a tracked chain. Returns the new row id.

        Parameters
        ----------
        domain       : monitoring domain name
        chain        : chain dict with keys question, steps, final_answer, finished
        risk_score   : LabelFreeResult.risk_score
        alert        : True if a QppgAlert was fired
        failure_mode : FailureTaxonomist primary_mode (empty if not run)
        model_name   : LLM model identifier for cross-model tracking
        """
        steps = chain.get("steps", [])
        n_steps = sum(1 for s in steps if s.get("action_type") == "Search")
        with self._conn() as conn:
            cur = conn.execute(
                """INSERT INTO chains
                   (domain, question, steps_json, final_answer, finished,
                    risk_score, alert_triggered, failure_mode, n_steps,
                    model_name, timestamp)
                   VALUES (?,?,?,?,?,?,?,?,?,?,?)""",
                (
                    domain,
                    chain.get("question", ""),
                    json.dumps(steps),
                    chain.get("final_answer", ""),
                    int(chain.get("finished", True)),
                    float(risk_score),
                    int(alert),
                    failure_mode,
                    n_steps,
                    model_name,
                    time.time(),
                ),
            )
            return cur.lastrowid

    def log_calibration(self, domain: str, n_chains: int, auroc: float) -> None:
        """Record a calibration event."""
        with self._conn() as conn:
            conn.execute(
                "INSERT INTO calibration_log (domain, n_chains, auroc_estimate, timestamp) "
                "VALUES (?,?,?,?)",
                (domain, n_chains, auroc, time.time()),
            )

    # ── Read ───────────────────────────────────────────────────────────────────

    def get_chains(
        self,
        domain: str,
        since: float | None = None,
        limit: int = 500,
    ) -> list[dict]:
        """Return recent chains for a domain, newest first."""
        since = since or 0.0
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT * FROM chains WHERE domain=? AND timestamp>=? "
                "ORDER BY timestamp DESC LIMIT ?",
                (domain, since, limit),
            ).fetchall()
        return [dict(r) for r in rows]

    def get_calibration_pool(self, domain: str, n: int = 200) -> list[dict]:
        """
        Return the most recent N chains as calibration candidates.
        Reconstructs chain dicts (question + steps) from stored JSON.
        """
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT question, steps_json, final_answer, finished "
                "FROM chains WHERE domain=? ORDER BY timestamp DESC LIMIT ?",
                (domain, n),
            ).fetchall()
        pool = []
        for r in rows:
            pool.append({
                "question":     r["question"],
                "steps":        json.loads(r["steps_json"] or "[]"),
                "final_answer": r["final_answer"],
                "finished":     bool(r["finished"]),
            })
        return pool

    def get_recent_risk(
        self,
        domain: str,
        days: int = 7,
    ) -> list[tuple[float, float]]:
        """
        Return (timestamp, risk_score) pairs for the last N days.
        Used by DriftDetector and the dashboard timeseries endpoint.
        """
        since = time.time() - days * 86400
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT timestamp, risk_score FROM chains "
                "WHERE domain=? AND timestamp>=? ORDER BY timestamp ASC",
                (domain, since),
            ).fetchall()
        return [(r["timestamp"], r["risk_score"]) for r in rows]

    def get_risk_window(
        self,
        domain: str,
        start: float,
        end: float,
    ) -> list[float]:
        """Return risk scores in [start, end] time range."""
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT risk_score FROM chains "
                "WHERE domain=? AND timestamp>=? AND timestamp<=?",
                (domain, start, end),
            ).fetchall()
        return [r["risk_score"] for r in rows]

    def get_domains(self) -> list[str]:
        """Return all distinct domain names in the store."""
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT DISTINCT domain FROM chains ORDER BY domain"
            ).fetchall()
        return [r["domain"] for r in rows]

    def get_domain_stats(self, domain: str) -> dict:
        """Summary stats for a domain (used by CLI status command)."""
        with self._conn() as conn:
            row = conn.execute(
                "SELECT COUNT(*) as n, AVG(risk_score) as avg_risk, "
                "SUM(alert_triggered) as n_alerts, MAX(timestamp) as last_ts "
                "FROM chains WHERE domain=?",
                (domain,),
            ).fetchone()
            cal_row = conn.execute(
                "SELECT auroc_estimate, timestamp FROM calibration_log "
                "WHERE domain=? ORDER BY timestamp DESC LIMIT 1",
                (domain,),
            ).fetchone()
        return {
            "domain":      domain,
            "n_chains":    row["n"] or 0,
            "avg_risk":    round(row["avg_risk"] or 0.0, 3),
            "n_alerts":    row["n_alerts"] or 0,
            "last_active": row["last_ts"],
            "last_auroc":  cal_row["auroc_estimate"] if cal_row else None,
            "last_cal_ts": cal_row["timestamp"] if cal_row else None,
        }

    def get_failure_mode_counts(self, domain: str, since: float = 0.0) -> dict:
        """Return {failure_mode: count} for alerted chains."""
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT failure_mode, COUNT(*) as cnt FROM chains "
                "WHERE domain=? AND alert_triggered=1 AND timestamp>=? "
                "AND failure_mode!='' GROUP BY failure_mode",
                (domain, since),
            ).fetchall()
        return {r["failure_mode"]: r["cnt"] for r in rows}

    # ── Delete ─────────────────────────────────────────────────────────────────

    def clear_domain(self, domain: str) -> int:
        """Delete all chains for a domain. Returns number of rows deleted."""
        with self._conn() as conn:
            cur = conn.execute("DELETE FROM chains WHERE domain=?", (domain,))
            conn.execute("DELETE FROM calibration_log WHERE domain=?", (domain,))
            return cur.rowcount

    # ── Audit export ───────────────────────────────────────────────────────────

    def export_audit(
        self,
        domain: str,
        start: float = 0.0,
        end: float | None = None,
        fmt: str = "csv",
    ) -> str:
        """
        Export an audit log for a domain over [start, end].

        Parameters
        ----------
        domain : domain name
        start  : Unix timestamp (default 0 = all time)
        end    : Unix timestamp (default now)
        fmt    : "csv" or "json"

        Returns
        -------
        String containing the audit log in the requested format.
        """
        end = end or time.time()
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT id, domain, question, final_answer, finished, "
                "risk_score, alert_triggered, failure_mode, n_steps, "
                "model_name, timestamp "
                "FROM chains WHERE domain=? AND timestamp>=? AND timestamp<=? "
                "ORDER BY timestamp ASC",
                (domain, start, end),
            ).fetchall()

        records = []
        for r in rows:
            records.append({
                "chain_id":       r["id"],
                "domain":         r["domain"],
                "timestamp":      r["timestamp"],
                "question_trunc": (r["question"] or "")[:120],
                "final_answer":   (r["final_answer"] or "")[:80],
                "finished":       bool(r["finished"]),
                "risk_score":     r["risk_score"],
                "alert_triggered": bool(r["alert_triggered"]),
                "failure_mode":   r["failure_mode"] or "",
                "n_search_steps": r["n_steps"],
                "model_name":     r["model_name"] or "",
            })

        if fmt == "json":
            return json.dumps(records, indent=2)

        # CSV
        if not records:
            return ""
        buf = io.StringIO()
        writer = csv.DictWriter(buf, fieldnames=list(records[0].keys()))
        writer.writeheader()
        writer.writerows(records)
        return buf.getvalue()

    # ── API key management (SaaS) ──────────────────────────────────────────────

    def create_api_key(
        self,
        customer_id: str,
        domain_prefix: str = "",
    ) -> str:
        """
        Generate and store a new API key.
        Returns the raw key (only time it is returned in plaintext).
        """
        import hashlib
        import secrets
        raw_key = secrets.token_urlsafe(32)
        key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
        with self._conn() as conn:
            conn.execute(
                "INSERT INTO api_keys (key_hash, customer_id, domain_prefix, created_at) "
                "VALUES (?,?,?,?)",
                (key_hash, customer_id, domain_prefix, time.time()),
            )
        return raw_key

    def verify_api_key(self, raw_key: str) -> dict | None:
        """
        Verify a raw API key.
        Returns the key record dict on success, None on failure.
        Increments request_count if valid.
        """
        import hashlib
        key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM api_keys WHERE key_hash=? AND is_active=1",
                (key_hash,),
            ).fetchone()
            if row:
                conn.execute(
                    "UPDATE api_keys SET request_count=request_count+1 WHERE key_hash=?",
                    (key_hash,),
                )
                return dict(row)
        return None

    def get_key_stats(self, key_hash: str) -> dict | None:
        """Return usage stats for an API key (by full hash)."""
        with self._conn() as conn:
            row = conn.execute(
                "SELECT customer_id, domain_prefix, created_at, "
                "request_count, is_active FROM api_keys WHERE key_hash=?",
                (key_hash,),
            ).fetchone()
        return dict(row) if row else None

    def get_key_by_prefix(self, key_prefix: str) -> dict | None:
        """
        Look up an API key by its 12-char hash prefix (as shown in the admin UI).

        Returns the full key record (without hash) on match, None if not found.
        Raises ValueError if the prefix matches more than one key (collision).
        """
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT key_hash, customer_id, domain_prefix, created_at, "
                "is_active, request_count FROM api_keys WHERE key_hash LIKE ?",
                (key_prefix + "%",),
            ).fetchall()
        if not rows:
            return None
        if len(rows) > 1:
            raise ValueError(f"Ambiguous prefix '{key_prefix}' matches {len(rows)} keys")
        r = dict(rows[0])
        r.pop("key_hash", None)   # never expose hash
        return r

    def revoke_key_by_prefix(self, key_prefix: str) -> bool:
        """
        Deactivate an API key by its 12-char hash prefix.

        Returns True if a key was revoked, False if prefix not found.
        Raises ValueError if the prefix is ambiguous (matches > 1 key).
        """
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT key_hash FROM api_keys WHERE key_hash LIKE ? AND is_active=1",
                (key_prefix + "%",),
            ).fetchall()
            if not rows:
                return False
            if len(rows) > 1:
                raise ValueError(f"Ambiguous prefix '{key_prefix}' matches {len(rows)} active keys")
            conn.execute(
                "UPDATE api_keys SET is_active=0 WHERE key_hash=?",
                (rows[0]["key_hash"],),
            )
        return True

    # ── Active learning review queue ────────────────────────────────────────────

    def get_review_queue(self, domain: str, n: int = 10) -> list[dict]:
        """
        Return the N most recent alerted, unreviewed chains for the review queue.

        Returns list of dicts with keys: chain_id, question, steps, final_answer,
        risk_score, failure_mode, n_steps, timestamp.
        """
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT id, question, steps_json, final_answer, risk_score, "
                "failure_mode, n_steps, timestamp "
                "FROM chains WHERE domain=? AND alert_triggered=1 AND reviewed=0 "
                "ORDER BY timestamp DESC LIMIT ?",
                (domain, n),
            ).fetchall()
        result = []
        for r in rows:
            result.append({
                "chain_id":     r["id"],
                "question":     (r["question"] or "")[:300],
                "steps":        json.loads(r["steps_json"] or "[]"),
                "final_answer": (r["final_answer"] or "")[:200],
                "risk_score":   r["risk_score"],
                "failure_mode": r["failure_mode"] or "",
                "n_steps":      r["n_steps"],
                "timestamp":    r["timestamp"],
            })
        return result

    def label_chain(self, chain_id: int, label: str) -> bool:
        """
        Mark a chain as reviewed and record its label ('correct' or 'incorrect').

        Parameters
        ----------
        chain_id : primary key of the chain row
        label    : 'correct' — agent was right despite the alert; or
                   'incorrect' — alert was justified (true positive)

        Returns True if a row was updated, False if chain_id not found.
        """
        with self._conn() as conn:
            cur = conn.execute(
                "UPDATE chains SET reviewed=1, label=? WHERE id=?",
                (label, chain_id),
            )
            return cur.rowcount > 0

    # ── User management (SaaS) ─────────────────────────────────────────────────

    def create_user(
        self,
        email: str,
        password_hash: str = "",
        google_id: str | None = None,
        name: str = "",
    ) -> dict:
        """Create a new user. Returns the created user dict."""
        now = time.time()
        # month_reset_at = start of next calendar month
        import calendar
        import datetime
        dt = datetime.datetime.utcfromtimestamp(now)
        if dt.month == 12:
            next_month = datetime.datetime(dt.year + 1, 1, 1)
        else:
            next_month = datetime.datetime(dt.year, dt.month + 1, 1)
        month_reset_at = calendar.timegm(next_month.timetuple())
        with self._conn() as conn:
            cur = conn.execute(
                "INSERT INTO users (email, password_hash, google_id, name, "
                "month_reset_at, created_at) VALUES (?,?,?,?,?,?)",
                (email, password_hash, google_id, name, month_reset_at, now),
            )
            user_id = cur.lastrowid
        return self.get_user_by_id(user_id)

    def get_user_by_email(self, email: str) -> dict | None:
        """Return user dict by email, or None if not found."""
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM users WHERE email=?", (email,)
            ).fetchone()
        return dict(row) if row else None

    def get_user_by_google_id(self, google_id: str) -> dict | None:
        """Return user dict by Google OAuth sub, or None if not found."""
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM users WHERE google_id=?", (google_id,)
            ).fetchone()
        return dict(row) if row else None

    def get_user_by_id(self, user_id: int) -> dict | None:
        """Return user dict by primary key, or None if not found."""
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM users WHERE id=?", (user_id,)
            ).fetchone()
        return dict(row) if row else None

    def update_user_plan(
        self,
        user_id: int,
        plan: str,
        stripe_customer_id: str = "",
        stripe_subscription_id: str = "",
    ) -> None:
        """Update a user's plan and Stripe IDs."""
        with self._conn() as conn:
            conn.execute(
                "UPDATE users SET plan=?, stripe_customer_id=?, "
                "stripe_subscription_id=? WHERE id=?",
                (plan, stripe_customer_id, stripe_subscription_id, user_id),
            )

    def increment_usage(self, user_id: int) -> int:
        """
        Increment chains_this_month for a user.
        Resets the counter if the month has rolled over.
        Returns the new count.
        """
        now = time.time()
        with self._conn() as conn:
            row = conn.execute(
                "SELECT chains_this_month, month_reset_at FROM users WHERE id=?",
                (user_id,),
            ).fetchone()
            if not row:
                return 0
            if now >= row["month_reset_at"]:
                # New month — compute next reset
                import calendar
                import datetime
                dt = datetime.datetime.utcfromtimestamp(now)
                if dt.month == 12:
                    next_month = datetime.datetime(dt.year + 1, 1, 1)
                else:
                    next_month = datetime.datetime(dt.year, dt.month + 1, 1)
                month_reset_at = calendar.timegm(next_month.timetuple())
                conn.execute(
                    "UPDATE users SET chains_this_month=1, month_reset_at=? WHERE id=?",
                    (month_reset_at, user_id),
                )
                return 1
            else:
                new_count = row["chains_this_month"] + 1
                conn.execute(
                    "UPDATE users SET chains_this_month=? WHERE id=?",
                    (new_count, user_id),
                )
                return new_count

    def get_monthly_usage(self, user_id: int) -> int:
        """Return current month's chain count for a user."""
        now = time.time()
        with self._conn() as conn:
            row = conn.execute(
                "SELECT chains_this_month, month_reset_at FROM users WHERE id=?",
                (user_id,),
            ).fetchone()
        if not row:
            return 0
        if now >= row["month_reset_at"]:
            return 0  # month rolled over, counter would be reset on next increment
        return row["chains_this_month"]

    def link_api_key_to_user(self, key_hash: str, user_id: int) -> None:
        """Associate an API key (by full hash) with a user."""
        with self._conn() as conn:
            conn.execute(
                "UPDATE api_keys SET user_id=? WHERE key_hash=?",
                (user_id, key_hash),
            )

    def get_keys_for_user(self, user_id: int) -> list[dict]:
        """Return all API keys belonging to a user (hashes excluded)."""
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT key_hash, customer_id, domain_prefix, created_at, "
                "is_active, request_count FROM api_keys WHERE user_id=? ORDER BY created_at DESC",
                (user_id,),
            ).fetchall()
        result = []
        for r in rows:
            d = dict(r)
            key_hash = d.pop("key_hash")
            d["key_prefix"] = key_hash[:12]
            result.append(d)
        return result

    # ── Domain stats for user dashboard ────────────────────────────────────────

    def get_domains_for_user(self, user_id: int) -> list[dict]:
        """Return per-domain stats for all domains belonging to this user's keys."""
        customer_prefix = f"user_{user_id}:"
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT domain, COUNT(*) as n_chains, AVG(risk_score) as avg_risk, "
                "SUM(alert_triggered) as n_alerts, MAX(timestamp) as last_ts "
                "FROM chains WHERE domain LIKE ? GROUP BY domain ORDER BY last_ts DESC",
                (customer_prefix + "%",),
            ).fetchall()
        return [
            {
                "domain":      r["domain"].replace(customer_prefix, ""),
                "chain_count": r["n_chains"],
                "avg_risk":    round(r["avg_risk"] or 0.0, 3),
                "n_alerts":    r["n_alerts"] or 0,
                "last_seen":   r["last_ts"],
            }
            for r in rows
        ]

    def get_recent_chains_for_user(self, user_id: int, limit: int = 20) -> list[dict]:
        """Return the most recent scored chains across all domains for this user."""
        customer_prefix = f"user_{user_id}:"
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT id, domain, question, final_answer, risk_score, "
                "alert_triggered, failure_mode, n_steps, timestamp "
                "FROM chains WHERE domain LIKE ? ORDER BY timestamp DESC LIMIT ?",
                (customer_prefix + "%", limit),
            ).fetchall()
        prefix = customer_prefix
        return [
            {
                "id":            r["id"],
                "domain":        r["domain"].replace(prefix, ""),
                "question":      (r["question"] or "")[:120],
                "final_answer":  (r["final_answer"] or "")[:80],
                "risk_score":    round(r["risk_score"], 3),
                "alert_triggered": bool(r["alert_triggered"]),
                "failure_mode":  r["failure_mode"] or "",
                "n_steps":       r["n_steps"],
                "timestamp":     r["timestamp"],
            }
            for r in rows
        ]

    # ── Workflow management ─────────────────────────────────────────────────────

    def save_workflow(self, user_id: int, name: str, nodes: list, edges: list, workflow_id: int | None = None) -> dict:
        now = time.time()
        with self._conn() as conn:
            if workflow_id:
                conn.execute(
                    "UPDATE workflows SET name=?, nodes_json=?, edges_json=?, updated_at=? WHERE id=? AND user_id=?",
                    (name, json.dumps(nodes), json.dumps(edges), now, workflow_id, user_id),
                )
                row_id = workflow_id
            else:
                cur = conn.execute(
                    "INSERT INTO workflows (user_id, name, nodes_json, edges_json, created_at, updated_at) VALUES (?,?,?,?,?,?)",
                    (user_id, name, json.dumps(nodes), json.dumps(edges), now, now),
                )
                row_id = cur.lastrowid
        return self.get_workflow(row_id, user_id)

    def get_workflows(self, user_id: int) -> list[dict]:
        with self._conn() as conn:
            rows = conn.execute(
                """SELECT DISTINCT w.* FROM workflows w WHERE w.user_id=?
                   UNION
                   SELECT DISTINCT w.* FROM workflows w
                   JOIN org_members m ON w.org_id = m.org_id
                   WHERE m.user_id=? AND w.org_id IS NOT NULL
                   ORDER BY updated_at DESC""",
                (user_id, user_id),
            ).fetchall()
        return [self._parse_workflow(dict(r)) for r in rows]

    def get_workflow(self, workflow_id: int, user_id: int) -> dict | None:
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM workflows WHERE id=? AND user_id=?", (workflow_id, user_id)
            ).fetchone()
        return self._parse_workflow(dict(row)) if row else None

    def delete_workflow(self, workflow_id: int, user_id: int) -> bool:
        with self._conn() as conn:
            cur = conn.execute("DELETE FROM workflows WHERE id=? AND user_id=?", (workflow_id, user_id))
        return cur.rowcount > 0

    def toggle_workflow(self, workflow_id: int, user_id: int, enabled: bool) -> dict | None:
        with self._conn() as conn:
            conn.execute(
                "UPDATE workflows SET enabled=?, updated_at=? WHERE id=? AND user_id=?",
                (int(enabled), time.time(), workflow_id, user_id),
            )
        return self.get_workflow(workflow_id, user_id)

    @staticmethod
    def _parse_workflow(r: dict) -> dict:
        r["nodes"] = json.loads(r.pop("nodes_json", "[]"))
        r["edges"] = json.loads(r.pop("edges_json", "[]"))
        return r

    # ── Guardrail management ────────────────────────────────────────────────────

    def save_guardrail(self, user_id: int, domain: str, threshold: float,
                       action: str, webhook_url: str = "", slack_webhook: str = "",
                       enabled: bool = True) -> dict:
        now = time.time()
        with self._conn() as conn:
            conn.execute(
                "INSERT INTO guardrails (user_id, domain, threshold, action, webhook_url, slack_webhook, enabled, created_at, updated_at) "
                "VALUES (?,?,?,?,?,?,?,?,?) ON CONFLICT(user_id, domain) DO UPDATE SET "
                "threshold=excluded.threshold, action=excluded.action, webhook_url=excluded.webhook_url, "
                "slack_webhook=excluded.slack_webhook, enabled=excluded.enabled, updated_at=excluded.updated_at",
                (user_id, domain, threshold, action, webhook_url, slack_webhook, int(enabled), now, now),
            )
        return self.get_guardrail(user_id, domain)

    def get_guardrails(self, user_id: int) -> list[dict]:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT * FROM guardrails WHERE user_id=? ORDER BY domain", (user_id,)
            ).fetchall()
        return [dict(r) for r in rows]

    def get_guardrail(self, user_id: int, domain: str) -> dict | None:
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM guardrails WHERE user_id=? AND domain=?", (user_id, domain)
            ).fetchone()
        return dict(row) if row else None

    def delete_guardrail(self, user_id: int, domain: str) -> bool:
        with self._conn() as conn:
            cur = conn.execute("DELETE FROM guardrails WHERE user_id=? AND domain=?", (user_id, domain))
        return cur.rowcount > 0

    # ── Pipeline management ─────────────────────────────────────────────────────

    def save_pipeline(self, user_id: int, name: str, nodes: list, edges: list,
                      pipeline_id: int | None = None) -> dict:
        now = time.time()
        with self._conn() as conn:
            if pipeline_id:
                conn.execute(
                    "UPDATE pipelines SET name=?, nodes_json=?, edges_json=?, updated_at=? WHERE id=? AND user_id=?",
                    (name, json.dumps(nodes), json.dumps(edges), now, pipeline_id, user_id),
                )
                row_id = pipeline_id
            else:
                cur = conn.execute(
                    "INSERT INTO pipelines (user_id, name, nodes_json, edges_json, created_at, updated_at) VALUES (?,?,?,?,?,?)",
                    (user_id, name, json.dumps(nodes), json.dumps(edges), now, now),
                )
                row_id = cur.lastrowid
        return self.get_pipeline(row_id, user_id)

    def get_pipelines(self, user_id: int) -> list[dict]:
        with self._conn() as conn:
            rows = conn.execute(
                """SELECT DISTINCT p.* FROM pipelines p WHERE p.user_id=?
                   UNION
                   SELECT DISTINCT p.* FROM pipelines p
                   JOIN org_members m ON p.org_id = m.org_id
                   WHERE m.user_id=? AND p.org_id IS NOT NULL
                   ORDER BY updated_at DESC""",
                (user_id, user_id),
            ).fetchall()
        return [self._parse_pipeline(dict(r)) for r in rows]

    def get_pipeline(self, pipeline_id: int, user_id: int) -> dict | None:
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM pipelines WHERE id=? AND user_id=?", (pipeline_id, user_id)
            ).fetchone()
        return self._parse_pipeline(dict(row)) if row else None

    def delete_pipeline(self, pipeline_id: int, user_id: int) -> bool:
        with self._conn() as conn:
            cur = conn.execute("DELETE FROM pipelines WHERE id=? AND user_id=?", (pipeline_id, user_id))
        return cur.rowcount > 0

    def toggle_pipeline(self, pipeline_id: int, user_id: int, enabled: bool) -> dict | None:
        with self._conn() as conn:
            conn.execute(
                "UPDATE pipelines SET enabled=?, updated_at=? WHERE id=? AND user_id=?",
                (int(enabled), time.time(), pipeline_id, user_id),
            )
        return self.get_pipeline(pipeline_id, user_id)

    @staticmethod
    def _parse_pipeline(r: dict) -> dict:
        r["nodes"] = json.loads(r.pop("nodes_json", "[]"))
        r["edges"] = json.loads(r.pop("edges_json", "[]"))
        return r

    # ── Pipeline run history ────────────────────────────────────────────────────

    def create_run(self, pipeline_id: int, user_id: int, query: str) -> dict:
        now = time.time()
        with self._conn() as conn:
            cur = conn.execute(
                "INSERT INTO pipeline_runs (pipeline_id, user_id, query, status, started_at) VALUES (?,?,?,?,?)",
                (pipeline_id, user_id, query, "running", now),
            )
        return {"id": cur.lastrowid, "pipeline_id": pipeline_id, "query": query, "status": "running", "started_at": now}

    def update_run(self, run_id: int, result: str = "", trace: list | None = None,
                   risk_score: float | None = None, guardrail_action: str | None = None,
                   status: str = "done") -> None:
        now = time.time()
        with self._conn() as conn:
            conn.execute(
                "UPDATE pipeline_runs SET result=?, trace_json=?, risk_score=?, guardrail_action=?, status=?, finished_at=? WHERE id=?",
                (result, json.dumps(trace or []), risk_score, guardrail_action, status, now, run_id),
            )

    def get_runs(self, pipeline_id: int, user_id: int, limit: int = 10) -> list[dict]:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT id, pipeline_id, query, result, risk_score, guardrail_action, status, started_at, finished_at "
                "FROM pipeline_runs WHERE pipeline_id=? AND user_id=? ORDER BY started_at DESC LIMIT ?",
                (pipeline_id, user_id, limit),
            ).fetchall()
        return [dict(r) for r in rows]

    def get_run_by_id(self, run_id: int, user_id: int) -> dict | None:
        """Return a single pipeline run including trace JSON (for SSE resume)."""
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM pipeline_runs WHERE id=? AND user_id=?",
                (run_id, user_id),
            ).fetchone()
        if not row:
            return None
        d = dict(row)
        d["trace"] = json.loads(d.pop("trace_json", "[]"))
        return d

    # ── User configs (LLM API keys) ─────────────────────────────────────────────

    def get_user_config(self, user_id: int) -> dict:
        with self._conn() as conn:
            row = conn.execute("SELECT config_json FROM user_configs WHERE user_id=?", (user_id,)).fetchone()
        if not row:
            return {}
        raw = json.loads(row["config_json"])
        # Decrypt secret values transparently
        return {k: _decrypt(v) if isinstance(v, str) else v for k, v in raw.items()}

    def set_user_config(self, user_id: int, config: dict) -> None:
        # Encrypt secret values before storing
        encrypted = {k: _encrypt(v) if isinstance(v, str) and v else v for k, v in config.items()}
        with self._conn() as conn:
            conn.execute(
                "INSERT INTO user_configs (user_id, config_json) VALUES (?,?) "
                "ON CONFLICT(user_id) DO UPDATE SET config_json=excluded.config_json",
                (user_id, json.dumps(encrypted)),
            )

    # ── Prompt versioning ───────────────────────────────────────────────────────

    def snapshot_pipeline_version(self, pipeline_id: int, user_id: int, note: str = "") -> dict:
        """Snapshot current pipeline state as a new version. Returns version record."""
        p = self.get_pipeline(pipeline_id, user_id)
        if not p:
            raise ValueError(f"Pipeline {pipeline_id} not found")
        now = time.time()
        with self._conn() as conn:
            row = conn.execute(
                "SELECT MAX(version) as max_v FROM pipeline_versions WHERE pipeline_id=?",
                (pipeline_id,),
            ).fetchone()
            next_v = (row["max_v"] or 0) + 1
            cur = conn.execute(
                "INSERT INTO pipeline_versions (pipeline_id, user_id, version, name, nodes_json, edges_json, note, created_at) "
                "VALUES (?,?,?,?,?,?,?,?)",
                (pipeline_id, user_id, next_v, p["name"],
                 json.dumps(p["nodes"]), json.dumps(p["edges"]), note, now),
            )
            vid = cur.lastrowid
        return {"id": vid, "pipeline_id": pipeline_id, "version": next_v,
                "name": p["name"], "note": note, "created_at": now}

    def get_pipeline_versions(self, pipeline_id: int, user_id: int) -> list[dict]:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT id, pipeline_id, version, name, note, created_at "
                "FROM pipeline_versions WHERE pipeline_id=? AND user_id=? ORDER BY version DESC",
                (pipeline_id, user_id),
            ).fetchall()
        return [dict(r) for r in rows]

    def get_pipeline_version(self, version_id: int, user_id: int) -> dict | None:
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM pipeline_versions WHERE id=? AND user_id=?",
                (version_id, user_id),
            ).fetchone()
        if not row:
            return None
        d = dict(row)
        d["nodes"] = json.loads(d.pop("nodes_json", "[]"))
        d["edges"] = json.loads(d.pop("edges_json", "[]"))
        return d

    def restore_pipeline_version(self, version_id: int, pipeline_id: int, user_id: int) -> dict | None:
        """Restore a pipeline to a saved version (also snapshots current as new version first)."""
        v = self.get_pipeline_version(version_id, user_id)
        if not v or v["pipeline_id"] != pipeline_id:
            return None
        self.snapshot_pipeline_version(pipeline_id, user_id, note="auto-snapshot before restore")
        now = time.time()
        with self._conn() as conn:
            conn.execute(
                "UPDATE pipelines SET nodes_json=?, edges_json=?, updated_at=? WHERE id=? AND user_id=?",
                (json.dumps(v["nodes"]), json.dumps(v["edges"]), now, pipeline_id, user_id),
            )
        return self.get_pipeline(pipeline_id, user_id)

    # ── A/B comparison ──────────────────────────────────────────────────────────

    def create_comparison(self, user_id: int, pipeline_a_id: int, pipeline_b_id: int, query: str) -> dict:
        now = time.time()
        with self._conn() as conn:
            cur = conn.execute(
                "INSERT INTO ab_comparisons (user_id, pipeline_a_id, pipeline_b_id, query, status, created_at) "
                "VALUES (?,?,?,?,?,?)",
                (user_id, pipeline_a_id, pipeline_b_id, query, "running", now),
            )
        return {"id": cur.lastrowid, "user_id": user_id, "pipeline_a_id": pipeline_a_id,
                "pipeline_b_id": pipeline_b_id, "query": query, "status": "running", "created_at": now}

    def update_comparison(self, cmp_id: int, result_a: str, result_b: str,
                          trace_a: list, trace_b: list,
                          risk_a: float | None, risk_b: float | None, status: str = "done") -> None:
        with self._conn() as conn:
            conn.execute(
                "UPDATE ab_comparisons SET result_a=?, result_b=?, trace_a_json=?, trace_b_json=?, "
                "risk_score_a=?, risk_score_b=?, status=?, finished_at=? WHERE id=?",
                (result_a, result_b, json.dumps(trace_a), json.dumps(trace_b),
                 risk_a, risk_b, status, time.time(), cmp_id),
            )

    def get_comparisons(self, user_id: int, limit: int = 20) -> list[dict]:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT id, pipeline_a_id, pipeline_b_id, query, result_a, result_b, "
                "risk_score_a, risk_score_b, status, created_at, finished_at "
                "FROM ab_comparisons WHERE user_id=? ORDER BY created_at DESC LIMIT ?",
                (user_id, limit),
            ).fetchall()
        return [dict(r) for r in rows]

    def get_comparison(self, cmp_id: int, user_id: int) -> dict | None:
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM ab_comparisons WHERE id=? AND user_id=?",
                (cmp_id, user_id),
            ).fetchone()
        if not row:
            return None
        d = dict(row)
        d["trace_a"] = json.loads(d.pop("trace_a_json", "[]"))
        d["trace_b"] = json.loads(d.pop("trace_b_json", "[]"))
        return d

    # ── Org / team management ───────────────────────────────────────────────────

    def create_org(self, owner_id: int, name: str) -> dict:
        now = time.time()
        with self._conn() as conn:
            cur = conn.execute(
                "INSERT INTO orgs (name, owner_id, created_at) VALUES (?,?,?)",
                (name, owner_id, now),
            )
            org_id = cur.lastrowid
            conn.execute(
                "INSERT INTO org_members (org_id, user_id, role, joined_at) VALUES (?,?,?,?)",
                (org_id, owner_id, "owner", now),
            )
        return {"id": org_id, "name": name, "owner_id": owner_id, "created_at": now, "role": "owner"}

    def get_orgs_for_user(self, user_id: int) -> list[dict]:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT o.id, o.name, o.owner_id, o.created_at, m.role "
                "FROM orgs o JOIN org_members m ON o.id=m.org_id WHERE m.user_id=?",
                (user_id,),
            ).fetchall()
        return [dict(r) for r in rows]

    def get_org_members(self, org_id: int) -> list[dict]:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT m.user_id, m.role, m.joined_at, u.email, u.name "
                "FROM org_members m JOIN users u ON m.user_id=u.id WHERE m.org_id=?",
                (org_id,),
            ).fetchall()
        return [dict(r) for r in rows]

    def create_invite(self, org_id: int, email: str) -> dict:
        import secrets
        token = secrets.token_urlsafe(32)
        now = time.time()
        with self._conn() as conn:
            cur = conn.execute(
                "INSERT INTO org_invites (org_id, email, token, created_at, expires_at) VALUES (?,?,?,?,?)",
                (org_id, email, token, now, now + 7 * 86400),  # 7 days
            )
        return {"id": cur.lastrowid, "org_id": org_id, "email": email,
                "token": token, "expires_at": now + 7 * 86400}

    def accept_invite(self, token: str, user_id: int) -> dict | None:
        now = time.time()
        with self._conn() as conn:
            inv = conn.execute(
                "SELECT * FROM org_invites WHERE token=? AND expires_at>?", (token, now)
            ).fetchone()
            if not inv:
                return None
            inv = dict(inv)
            existing = conn.execute(
                "SELECT 1 FROM org_members WHERE org_id=? AND user_id=?",
                (inv["org_id"], user_id),
            ).fetchone()
            if not existing:
                conn.execute(
                    "INSERT INTO org_members (org_id, user_id, role, joined_at) VALUES (?,?,?,?)",
                    (inv["org_id"], user_id, "member", now),
                )
            conn.execute("DELETE FROM org_invites WHERE token=?", (token,))
        return inv

    def remove_member(self, org_id: int, user_id: int, owner_id: int) -> bool:
        if user_id == owner_id:
            return False  # can't remove owner
        with self._conn() as conn:
            cur = conn.execute(
                "DELETE FROM org_members WHERE org_id=? AND user_id=?", (org_id, user_id)
            )
        return cur.rowcount > 0

    def is_org_member(self, org_id: int, user_id: int) -> bool:
        with self._conn() as conn:
            row = conn.execute(
                "SELECT 1 FROM org_members WHERE org_id=? AND user_id=?", (org_id, user_id)
            ).fetchone()
        return row is not None

    # ── Test sets & regression runs ─────────────────────────────────────────────

    def save_test_set(self, user_id: int, name: str, cases: list,
                     pipeline_id: int | None = None, test_set_id: int | None = None) -> dict:
        now = time.time()
        with self._conn() as conn:
            if test_set_id:
                conn.execute(
                    "UPDATE test_sets SET name=?, cases_json=?, pipeline_id=?, updated_at=? "
                    "WHERE id=? AND user_id=?",
                    (name, json.dumps(cases), pipeline_id, now, test_set_id, user_id),
                )
                row_id = test_set_id
            else:
                cur = conn.execute(
                    "INSERT INTO test_sets (user_id, name, cases_json, pipeline_id, created_at, updated_at) "
                    "VALUES (?,?,?,?,?,?)",
                    (user_id, name, json.dumps(cases), pipeline_id, now, now),
                )
                row_id = cur.lastrowid
        return self.get_test_set(row_id, user_id)

    def get_test_sets(self, user_id: int) -> list[dict]:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT id, user_id, name, pipeline_id, created_at, updated_at, "
                "json_array_length(cases_json) as case_count "
                "FROM test_sets WHERE user_id=? ORDER BY updated_at DESC",
                (user_id,),
            ).fetchall()
        return [dict(r) for r in rows]

    def get_test_set(self, test_set_id: int, user_id: int) -> dict | None:
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM test_sets WHERE id=? AND user_id=?", (test_set_id, user_id)
            ).fetchone()
        if not row:
            return None
        d = dict(row)
        d["cases"] = json.loads(d.pop("cases_json", "[]"))
        return d

    def delete_test_set(self, test_set_id: int, user_id: int) -> bool:
        with self._conn() as conn:
            cur = conn.execute("DELETE FROM test_sets WHERE id=? AND user_id=?", (test_set_id, user_id))
        return cur.rowcount > 0

    def create_test_run(self, test_set_id: int, user_id: int, pipeline_id: int) -> dict:
        now = time.time()
        with self._conn() as conn:
            cur = conn.execute(
                "INSERT INTO test_runs (test_set_id, user_id, pipeline_id, status, started_at) "
                "VALUES (?,?,?,?,?)",
                (test_set_id, user_id, pipeline_id, "running", now),
            )
        return {"id": cur.lastrowid, "test_set_id": test_set_id, "pipeline_id": pipeline_id,
                "status": "running", "started_at": now}

    def update_test_run(self, run_id: int, results: list, avg_risk: float,
                        pass_rate: float, status: str = "done") -> None:
        with self._conn() as conn:
            conn.execute(
                "UPDATE test_runs SET results_json=?, avg_risk_score=?, pass_rate=?, "
                "status=?, finished_at=? WHERE id=?",
                (json.dumps(results), avg_risk, pass_rate, status, time.time(), run_id),
            )

    def get_test_runs(self, test_set_id: int, user_id: int) -> list[dict]:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT id, test_set_id, pipeline_id, avg_risk_score, pass_rate, "
                "status, started_at, finished_at FROM test_runs "
                "WHERE test_set_id=? AND user_id=? ORDER BY started_at DESC LIMIT 20",
                (test_set_id, user_id),
            ).fetchall()
        return [dict(r) for r in rows]

    def get_test_run(self, run_id: int, user_id: int) -> dict | None:
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM test_runs WHERE id=? AND user_id=?", (run_id, user_id)
            ).fetchone()
        if not row:
            return None
        d = dict(row)
        d["results"] = json.loads(d.pop("results_json", "[]"))
        return d

    # ── Proxy request log ───────────────────────────────────────────────────────

    def log_proxy_request(self, user_id: int, provider: str, model: str,
                          prompt_tokens: int, completion_tokens: int,
                          risk_score: float | None, latency_ms: int, status: str = "ok") -> int:
        with self._conn() as conn:
            cur = conn.execute(
                "INSERT INTO proxy_requests (user_id, provider, model, prompt_tokens, "
                "completion_tokens, risk_score, latency_ms, status, created_at) "
                "VALUES (?,?,?,?,?,?,?,?,?)",
                (user_id, provider, model, prompt_tokens, completion_tokens,
                 risk_score, latency_ms, status, time.time()),
            )
        return cur.lastrowid

    def get_proxy_stats(self, user_id: int, days: int = 30) -> dict:
        since = time.time() - days * 86400
        with self._conn() as conn:
            row = conn.execute(
                "SELECT COUNT(*) as n_calls, SUM(prompt_tokens) as total_prompt, "
                "SUM(completion_tokens) as total_completion, AVG(latency_ms) as avg_latency, "
                "AVG(risk_score) as avg_risk, SUM(CASE WHEN status!='ok' THEN 1 ELSE 0 END) as n_errors "
                "FROM proxy_requests WHERE user_id=? AND created_at>=?",
                (user_id, since),
            ).fetchone()
        return dict(row) if row else {}

    def get_proxy_requests(self, user_id: int, limit: int = 50) -> list[dict]:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT id, provider, model, prompt_tokens, completion_tokens, "
                "risk_score, latency_ms, status, created_at "
                "FROM proxy_requests WHERE user_id=? ORDER BY created_at DESC LIMIT ?",
                (user_id, limit),
            ).fetchall()
        return [dict(r) for r in rows]

    # ── Time-series analytics ───────────────────────────────────────────────────

    def get_timeseries(self, user_id: int, domain: str | None = None,
                       days: int = 30, bucket: str = "day") -> list[dict]:
        """
        Return bucketed (timestamp, avg_risk, n_chains, n_alerts) for the dashboard.
        bucket: 'hour' | 'day' | 'week'
        Covers pipeline_runs (from pipeline guard nodes) and chains (from API scoring).
        """
        since = time.time() - days * 86400
        secs = {"hour": 3600, "day": 86400, "week": 604800}.get(bucket, 86400)

        prefix = f"user_{user_id}:"
        domain_filter = f"{prefix}{domain}" if domain else None

        with self._conn() as conn:
            if domain_filter:
                rows = conn.execute(
                    "SELECT timestamp, risk_score, alert_triggered FROM chains "
                    "WHERE domain=? AND timestamp>=? ORDER BY timestamp ASC",
                    (domain_filter, since),
                ).fetchall()
            else:
                rows = conn.execute(
                    "SELECT timestamp, risk_score, alert_triggered FROM chains "
                    "WHERE domain LIKE ? AND timestamp>=? ORDER BY timestamp ASC",
                    (prefix + "%", since),
                ).fetchall()

            # Also add pipeline run data
            pr_rows = conn.execute(
                "SELECT started_at as timestamp, risk_score, 0 as alert_triggered "
                "FROM pipeline_runs WHERE user_id=? AND started_at>=? AND risk_score IS NOT NULL",
                (user_id, since),
            ).fetchall()

        all_rows = [(r["timestamp"], r["risk_score"], r["alert_triggered"])
                    for r in list(rows) + list(pr_rows)]

        if not all_rows:
            return []

        # Bucket by time
        buckets: dict[int, list] = {}
        for ts, risk, alert in all_rows:
            bucket_key = int(ts // secs) * secs
            buckets.setdefault(bucket_key, []).append((risk, alert))

        result = []
        for bk in sorted(buckets.keys()):
            items = buckets[bk]
            risks = [r for r, _ in items]
            alerts = sum(a for _, a in items)
            result.append({
                "timestamp": bk,
                "avg_risk": round(sum(risks) / len(risks), 4),
                "n_chains": len(items),
                "n_alerts": alerts,
            })
        return result

    # ── Automation / guardrail execution log ────────────────────────────────────

    def log_automation(
        self,
        user_id: int,
        trigger_type: str,
        domain: str,
        risk_score: float | None,
        chain_id: int | None = None,
        workflow_id: int | None = None,
        guardrail_id: int | None = None,
        pipeline_run_id: int | None = None,
        action: str = "",
        result: str = "",
        error: str = "",
    ) -> int:
        """Persist a single automation/guardrail execution event. Returns row id."""
        with self._conn() as conn:
            cur = conn.execute(
                "INSERT INTO automation_log (user_id, trigger_type, domain, risk_score, "
                "chain_id, workflow_id, guardrail_id, pipeline_run_id, action, result, error, fired_at) "
                "VALUES (?,?,?,?,?,?,?,?,?,?,?,?)",
                (user_id, trigger_type, domain, risk_score, chain_id, workflow_id,
                 guardrail_id, pipeline_run_id, action, result, error, time.time()),
            )
        return cur.lastrowid

    def get_automation_logs(self, user_id: int, limit: int = 100, domain: str = "") -> list[dict]:
        """Return recent automation execution log entries, newest first."""
        with self._conn() as conn:
            if domain:
                rows = conn.execute(
                    "SELECT * FROM automation_log WHERE user_id=? AND domain=? "
                    "ORDER BY fired_at DESC LIMIT ?",
                    (user_id, domain, limit),
                ).fetchall()
            else:
                rows = conn.execute(
                    "SELECT * FROM automation_log WHERE user_id=? ORDER BY fired_at DESC LIMIT ?",
                    (user_id, limit),
                ).fetchall()
        return [dict(r) for r in rows]

    # ── Version diff ────────────────────────────────────────────────────────────

    @staticmethod
    def diff_pipeline_versions(v_old: dict, v_new: dict) -> dict:
        """
        Compare two pipeline versions by node id.
        Returns {nodes_added, nodes_removed, nodes_changed, edges_added, edges_removed}.
        """
        old_nodes = {n["id"]: n for n in v_old.get("nodes", [])}
        new_nodes = {n["id"]: n for n in v_new.get("nodes", [])}
        added   = [n for nid, n in new_nodes.items() if nid not in old_nodes]
        removed = [n for nid, n in old_nodes.items() if nid not in new_nodes]
        changed = []
        for nid in set(old_nodes) & set(new_nodes):
            o, n = old_nodes[nid], new_nodes[nid]
            if json.dumps(o, sort_keys=True) != json.dumps(n, sort_keys=True):
                changed.append({"id": nid, "before": o, "after": n})
        old_edges = {(e["source"], e["target"]) for e in v_old.get("edges", [])}
        new_edges = {(e["source"], e["target"]) for e in v_new.get("edges", [])}
        return {
            "nodes_added":   added,
            "nodes_removed": removed,
            "nodes_changed": changed,
            "edges_added":   [{"source": s, "target": t} for s, t in (new_edges - old_edges)],
            "edges_removed": [{"source": s, "target": t} for s, t in (old_edges - new_edges)],
        }
