"""
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 csv
import io
import json
import sqlite3
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Generator


_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
);
"""


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)

    # ── 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
        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 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
