import json
import re
import sqlite3
import struct

import sqlite_vec

from pydstl.types import Evidence, Skill

_SAFE_KEY = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")

EMBEDDING_DIM = 384


def _serialize_f32(vec: list[float]) -> bytes:
    return struct.pack(f"{len(vec)}f", *vec)


class Store:
    def __init__(self, db_path: str):
        self._conn = sqlite3.connect(db_path)
        self._conn.enable_load_extension(True)
        sqlite_vec.load(self._conn)
        self._conn.enable_load_extension(False)
        self._create_tables()

    def _create_tables(self) -> None:
        self._conn.executescript(f"""
            CREATE TABLE IF NOT EXISTS evidence (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                content TEXT NOT NULL,
                source_json TEXT,
                created_at TEXT DEFAULT (datetime('now'))
            );

            CREATE VIRTUAL TABLE IF NOT EXISTS evidence_vec USING vec0(
                id INTEGER PRIMARY KEY,
                embedding FLOAT[{EMBEDDING_DIM}]
            );

            CREATE TABLE IF NOT EXISTS skills (
                id TEXT PRIMARY KEY,
                title TEXT,
                file_path TEXT NOT NULL,
                created_at TEXT DEFAULT (datetime('now')),
                updated_at TEXT
            );

            CREATE TABLE IF NOT EXISTS skill_evidence (
                skill_id TEXT REFERENCES skills(id),
                evidence_id INTEGER REFERENCES evidence(id),
                PRIMARY KEY (skill_id, evidence_id)
            );

            CREATE TABLE IF NOT EXISTS outcomes (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                skill_id TEXT REFERENCES skills(id),
                success BOOLEAN NOT NULL,
                notes TEXT,
                created_at TEXT DEFAULT (datetime('now'))
            );
        """)

    def add_evidence(self, content: str, source: dict | None, embedding: list[float]) -> int:
        cur = self._conn.execute(
            "INSERT INTO evidence (content, source_json) VALUES (?, ?)",
            (content, json.dumps(source) if source else None),
        )
        eid = cur.lastrowid
        assert eid is not None
        self._conn.execute(
            "INSERT INTO evidence_vec (id, embedding) VALUES (?, ?)",
            (eid, _serialize_f32(embedding)),
        )
        self._conn.commit()
        return eid

    def search_evidence(self, embedding: list[float], top_k: int = 5) -> list[Evidence]:
        rows = self._conn.execute(
            """
            SELECT e.id, e.content, e.source_json, e.created_at, v.distance
            FROM evidence_vec v
            JOIN evidence e ON e.id = v.id
            WHERE v.embedding MATCH ? AND k = ?
            ORDER BY v.distance
            """,
            (_serialize_f32(embedding), top_k),
        ).fetchall()
        return [
            Evidence(
                id=r[0],
                content=r[1],
                source=json.loads(r[2]) if r[2] else None,
                created_at=r[3],
                score=1.0 - r[4],  # convert distance to similarity
            )
            for r in rows
        ]

    def get_all_evidence(self) -> list[Evidence]:
        rows = self._conn.execute(
            "SELECT id, content, source_json, created_at FROM evidence"
        ).fetchall()
        return [
            Evidence(
                id=r[0],
                content=r[1],
                source=json.loads(r[2]) if r[2] else None,
                created_at=r[3],
            )
            for r in rows
        ]

    def list_evidence(self, source_filter: dict | None = None) -> list[Evidence]:
        if not source_filter:
            return self.get_all_evidence()

        # Build WHERE clauses using json_extract
        conditions = []
        params: list = []
        for key, value in source_filter.items():
            if not _SAFE_KEY.match(key):
                raise ValueError(f"Invalid source filter key: {key!r}")
            conditions.append(f"json_extract(source_json, '$.{key}') = ?")
            params.append(value)

        where = " AND ".join(conditions)
        rows = self._conn.execute(
            f"SELECT id, content, source_json, created_at FROM evidence WHERE {where}",
            params,
        ).fetchall()
        return [
            Evidence(
                id=r[0],
                content=r[1],
                source=json.loads(r[2]) if r[2] else None,
                created_at=r[3],
            )
            for r in rows
        ]

    def save_skill(self, skill: Skill, evidence_ids: list[int] | None = None) -> None:
        self._conn.execute(
            """
            INSERT INTO skills (id, title, file_path, created_at, updated_at)
            VALUES (?, ?, ?, ?, ?)
            ON CONFLICT(id) DO UPDATE SET
                title = excluded.title,
                file_path = excluded.file_path,
                updated_at = datetime('now')
            """,
            (skill.id, skill.title, skill.file_path, skill.created_at, skill.updated_at),
        )
        if evidence_ids:
            for eid in evidence_ids:
                self._conn.execute(
                    "INSERT OR IGNORE INTO skill_evidence (skill_id, evidence_id) VALUES (?, ?)",
                    (skill.id, eid),
                )
        self._conn.commit()

    def get_skill(self, skill_id: str) -> Skill | None:
        row = self._conn.execute(
            "SELECT id, title, file_path, created_at, updated_at FROM skills WHERE id = ?",
            (skill_id,),
        ).fetchone()
        if not row:
            return None
        return Skill(
            id=row[0], title=row[1], file_path=row[2], created_at=row[3], updated_at=row[4]
        )

    def add_outcome(self, skill_id: str, success: bool, notes: str) -> int:
        cur = self._conn.execute(
            "INSERT INTO outcomes (skill_id, success, notes) VALUES (?, ?, ?)",
            (skill_id, success, notes),
        )
        self._conn.commit()
        oid = cur.lastrowid
        assert oid is not None
        return oid

    def close(self) -> None:
        self._conn.close()
