"""SQLite storage layer for MVP.

Stores baselines, tests, and results in a local SQLite database.
Default location: ~/.config/kredo-drift/drift.db

This is intentionally simple for Phase 0. The schema mirrors the
PostgreSQL design in ARCHITECTURE.md but without append-only enforcement,
hash chains, or encryption. Those come in Phase 2.
"""

from __future__ import annotations

import json
import os
import sqlite3
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from kredo_drift.models import (
    DimensionScore,
    DriftBaseline,
    DriftClassification,
    DriftResult,
    DriftTest,
    Dimension,
)


DEFAULT_DB_PATH = Path.home() / ".config" / "kredo-drift" / "drift.db"


class DriftStore:
    """SQLite-backed storage for drift baselines, tests, and results."""

    def __init__(self, db_path: Optional[str | Path] = None):
        self.db_path = Path(db_path) if db_path else DEFAULT_DB_PATH
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._conn: Optional[sqlite3.Connection] = None
        self._init_db()

    def _get_conn(self) -> sqlite3.Connection:
        if self._conn is None:
            self._conn = sqlite3.connect(str(self.db_path))
            self._conn.row_factory = sqlite3.Row
            self._conn.execute("PRAGMA journal_mode=WAL")
            self._conn.execute("PRAGMA foreign_keys=ON")
        return self._conn

    def _init_db(self) -> None:
        """Create tables if they don't exist."""
        conn = self._get_conn()
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS baselines (
                id TEXT PRIMARY KEY,
                agent_name TEXT NOT NULL,
                created_at TEXT NOT NULL,
                model_name TEXT NOT NULL,
                model_version TEXT NOT NULL,
                vector_json TEXT NOT NULL,
                dimension_vectors_json TEXT NOT NULL,
                response_hashes_json TEXT NOT NULL,
                prompt_ids_json TEXT NOT NULL,
                prompt_count INTEGER NOT NULL,
                content_hash TEXT NOT NULL,
                pair_similarities_json TEXT NOT NULL DEFAULT '{}',
                anchor_values_json TEXT NOT NULL DEFAULT '{}',
                fidelity_vector_json TEXT NOT NULL DEFAULT '[]',
                dimension_coherence_json TEXT NOT NULL DEFAULT '{}'
            );

            CREATE INDEX IF NOT EXISTS idx_baselines_agent
                ON baselines(agent_name);

            CREATE TABLE IF NOT EXISTS drift_tests (
                id TEXT PRIMARY KEY,
                agent_name TEXT NOT NULL,
                baseline_id TEXT NOT NULL,
                created_at TEXT NOT NULL,
                model_name TEXT NOT NULL,
                model_version TEXT NOT NULL,
                vector_json TEXT NOT NULL,
                dimension_vectors_json TEXT NOT NULL,
                response_hashes_json TEXT NOT NULL,
                prompt_ids_json TEXT NOT NULL,
                prompt_count INTEGER NOT NULL,
                content_hash TEXT NOT NULL,
                FOREIGN KEY (baseline_id) REFERENCES baselines(id)
            );

            CREATE INDEX IF NOT EXISTS idx_tests_agent
                ON drift_tests(agent_name);

            CREATE TABLE IF NOT EXISTS drift_results (
                id TEXT PRIMARY KEY,
                agent_name TEXT NOT NULL,
                test_id TEXT NOT NULL,
                baseline_id TEXT NOT NULL,
                created_at TEXT NOT NULL,
                score REAL NOT NULL,
                dimension_scores_json TEXT NOT NULL,
                classification TEXT NOT NULL,
                baseline_hash TEXT NOT NULL,
                test_hash TEXT NOT NULL,
                FOREIGN KEY (test_id) REFERENCES drift_tests(id),
                FOREIGN KEY (baseline_id) REFERENCES baselines(id)
            );

            CREATE INDEX IF NOT EXISTS idx_results_agent
                ON drift_results(agent_name);
            CREATE INDEX IF NOT EXISTS idx_results_created
                ON drift_results(agent_name, created_at);
        """)
        conn.commit()
        self._migrate(conn)

    def _migrate(self, conn: sqlite3.Connection) -> None:
        """Add columns introduced after the initial schema."""
        # WHY ALTER TABLE instead of recreate: Existing databases have data
        # we don't want to lose. SQLite ADD COLUMN is safe and additive.
        # The DEFAULT ensures old rows get sensible empty values.
        existing = {
            row[1]
            for row in conn.execute("PRAGMA table_info(baselines)").fetchall()
        }
        migrations = [
            ("pair_similarities_json", "TEXT NOT NULL DEFAULT '{}'"),
            ("anchor_values_json", "TEXT NOT NULL DEFAULT '{}'"),
            ("fidelity_vector_json", "TEXT NOT NULL DEFAULT '[]'"),
        ]
        for col_name, col_def in migrations:
            if col_name not in existing:
                conn.execute(
                    f"ALTER TABLE baselines ADD COLUMN {col_name} {col_def}"
                )
        conn.commit()

    def save_baseline(self, baseline: DriftBaseline) -> None:
        """Save a baseline to the store."""
        conn = self._get_conn()
        conn.execute(
            """INSERT OR REPLACE INTO baselines
               (id, agent_name, created_at, model_name, model_version,
                vector_json, dimension_vectors_json, response_hashes_json,
                prompt_ids_json, prompt_count, content_hash,
                pair_similarities_json, anchor_values_json, fidelity_vector_json,
                dimension_coherence_json)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                baseline.id,
                baseline.agent_name,
                baseline.created_at.isoformat(),
                baseline.model_name,
                baseline.model_version,
                json.dumps(baseline.vector),
                json.dumps(baseline.dimension_vectors),
                json.dumps(baseline.response_hashes),
                json.dumps(baseline.prompt_ids),
                baseline.prompt_count,
                baseline.content_hash(),
                json.dumps(baseline.pair_similarities),
                json.dumps(baseline.anchor_values),
                json.dumps(baseline.fidelity_vector),
                json.dumps(baseline.dimension_coherence),
            ),
        )
        conn.commit()

    def load_baseline(self, agent_name: str) -> Optional[DriftBaseline]:
        """Load the latest baseline for an agent."""
        conn = self._get_conn()
        row = conn.execute(
            """SELECT * FROM baselines
               WHERE agent_name = ?
               ORDER BY created_at DESC
               LIMIT 1""",
            (agent_name,),
        ).fetchone()

        if row is None:
            return None

        return self._row_to_baseline(row)

    def load_baseline_by_id(self, baseline_id: str) -> Optional[DriftBaseline]:
        """Load a specific baseline by ID."""
        conn = self._get_conn()
        row = conn.execute(
            "SELECT * FROM baselines WHERE id = ?",
            (baseline_id,),
        ).fetchone()

        if row is None:
            return None

        return self._row_to_baseline(row)

    def _row_to_baseline(self, row: sqlite3.Row) -> DriftBaseline:
        return DriftBaseline(
            id=row["id"],
            agent_name=row["agent_name"],
            created_at=datetime.fromisoformat(row["created_at"]),
            model_name=row["model_name"],
            model_version=row["model_version"],
            vector=json.loads(row["vector_json"]),
            dimension_vectors=json.loads(row["dimension_vectors_json"]),
            response_hashes=json.loads(row["response_hashes_json"]),
            prompt_ids=json.loads(row["prompt_ids_json"]),
            prompt_count=row["prompt_count"],
            pair_similarities=json.loads(row["pair_similarities_json"]),
            anchor_values=json.loads(row["anchor_values_json"]),
            fidelity_vector=json.loads(row["fidelity_vector_json"]),
            dimension_coherence=json.loads(row["dimension_coherence_json"] if "dimension_coherence_json" in row.keys() else "{}"),
        )

    def save_test(self, test: DriftTest) -> None:
        """Save a drift test to the store."""
        conn = self._get_conn()
        conn.execute(
            """INSERT OR REPLACE INTO drift_tests
               (id, agent_name, baseline_id, created_at, model_name, model_version,
                vector_json, dimension_vectors_json, response_hashes_json,
                prompt_ids_json, prompt_count, content_hash)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                test.id,
                test.agent_name,
                test.baseline_id,
                test.created_at.isoformat(),
                test.model_name,
                test.model_version,
                json.dumps(test.vector),
                json.dumps(test.dimension_vectors),
                json.dumps(test.response_hashes),
                json.dumps(test.prompt_ids),
                test.prompt_count,
                test.content_hash(),
            ),
        )
        conn.commit()

    def save_result(self, result: DriftResult) -> None:
        """Save a drift result to the store."""
        conn = self._get_conn()

        dim_scores_data = [
            {
                "dimension": ds.dimension.value,
                "score": ds.score,
                "similarity": ds.similarity,
                "prompt_count": ds.prompt_count,
            }
            for ds in result.dimension_scores
        ]

        conn.execute(
            """INSERT OR REPLACE INTO drift_results
               (id, agent_name, test_id, baseline_id, created_at,
                score, dimension_scores_json, classification,
                baseline_hash, test_hash)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                result.id,
                result.agent_name,
                result.test_id,
                result.baseline_id,
                result.created_at.isoformat(),
                result.score,
                json.dumps(dim_scores_data),
                result.classification.value,
                result.baseline_hash,
                result.test_hash,
            ),
        )
        conn.commit()

    def load_latest_result(self, agent_name: str) -> Optional[DriftResult]:
        """Load the most recent drift result for an agent."""
        conn = self._get_conn()
        row = conn.execute(
            """SELECT * FROM drift_results
               WHERE agent_name = ?
               ORDER BY created_at DESC
               LIMIT 1""",
            (agent_name,),
        ).fetchone()

        if row is None:
            return None

        return self._row_to_result(row)

    def load_results(
        self,
        agent_name: str,
        limit: int = 20,
    ) -> list[DriftResult]:
        """Load drift result history for an agent (most recent first)."""
        conn = self._get_conn()
        rows = conn.execute(
            """SELECT * FROM drift_results
               WHERE agent_name = ?
               ORDER BY created_at DESC
               LIMIT ?""",
            (agent_name, limit),
        ).fetchall()

        return [self._row_to_result(row) for row in rows]

    def _row_to_result(self, row: sqlite3.Row) -> DriftResult:
        dim_scores_data = json.loads(row["dimension_scores_json"])
        dimension_scores = [
            DimensionScore(
                dimension=Dimension(d["dimension"]),
                score=d["score"],
                similarity=d["similarity"],
                prompt_count=d["prompt_count"],
            )
            for d in dim_scores_data
        ]

        return DriftResult(
            id=row["id"],
            agent_name=row["agent_name"],
            test_id=row["test_id"],
            baseline_id=row["baseline_id"],
            created_at=datetime.fromisoformat(row["created_at"]),
            score=row["score"],
            dimension_scores=dimension_scores,
            classification=DriftClassification(row["classification"]),
            baseline_hash=row["baseline_hash"],
            test_hash=row["test_hash"],
        )

    def list_agents(self) -> list[str]:
        """List all agent names that have baselines."""
        conn = self._get_conn()
        rows = conn.execute(
            "SELECT DISTINCT agent_name FROM baselines ORDER BY agent_name"
        ).fetchall()
        return [row["agent_name"] for row in rows]

    # WHY hardcoded allowlist: table names can't be parameterized in SQL.
    # This frozenset ensures only known tables are used in the DELETE below.
    _AGENT_TABLES = frozenset(["drift_results", "drift_tests", "baselines"])

    def delete_agent(self, agent_name: str) -> int:
        """Delete all data for an agent. Returns the number of records deleted."""
        conn = self._get_conn()
        counts = 0
        for table in sorted(self._AGENT_TABLES):
            assert table.isidentifier(), f"Invalid table name: {table}"
            cursor = conn.execute(
                f"DELETE FROM {table} WHERE agent_name = ?",  # nosec B608 — table from allowlist
                (agent_name,),
            )
            counts += cursor.rowcount
        conn.commit()
        return counts

    def close(self) -> None:
        """Close the database connection."""
        if self._conn:
            self._conn.close()
            self._conn = None
