import json
import sqlite3
from dataclasses import dataclass
from pathlib import Path

import sqlite_vec

from grimoire.data import entry, meta, schema
from grimoire.data.entry import (
    Entry,
    Filters,
    KeywordHit,
    SemanticHit,
)
from grimoire.embed import Embedder
from grimoire.errors import GrimoireMismatch, GrimoireNotFound


@dataclass(frozen=True, slots=True)
class Peek:
    model: str
    dimension: int
    schema_version: int
    entry_count: int
    group_counts: dict[str | None, int]
    partition_counts: dict[str | None, int]


class Grimoire:
    def __init__(self, conn: sqlite3.Connection, embedder: Embedder) -> None:
        self._conn = conn
        self.embedder = embedder

    def __enter__(self) -> "Grimoire":
        return self

    def __exit__(self, exc_type, *_) -> None:
        if exc_type is None:
            self._conn.commit()
        else:
            self._conn.rollback()

    def add(self, entries: list[Entry]) -> list[Entry]:
        return entry.add(self._conn, entries)

    def update(self, entries: list[Entry]) -> list[Entry]:
        return entry.update(self._conn, entries)

    def remove(self, ids: list[str]) -> list[str]:
        return entry.remove(self._conn, ids)

    def fetch(
        self,
        filters: Filters | None = None,
        limit: int = 100,
        cursor: str | None = None,
    ) -> list[Entry]:
        """Fetch entries with their FTS5 and vec0 index fields populated.

        Entries with no keyword or semantic index row come back with the
        corresponding fields set to None. ULIDs sort lexicographically by
        creation time, so passing `cursor` set to the last entry's id from
        the previous page walks entries chronologically.
        """
        return entry.fetch(self._conn, filters, limit, cursor=cursor)

    def keyword_remove(self, ids: list[str]) -> list[str]:
        """Delete entry_fts rows for the given ids. Returns the ids that had rows.

        Entries themselves are not affected. Empty `ids` is a no-op.
        """
        return entry.keyword_remove(self._conn, ids)

    def embed_remove(self, ids: list[str]) -> list[str]:
        """Delete entry_vec rows for the given ids. Returns the ids that had rows.

        Entries themselves are not affected. Empty `ids` is a no-op.
        """
        return entry.embed_remove(self._conn, ids)

    def keyword(
        self,
        items: list[tuple[str, str]] | None = None,
        *,
        threshold_rank: float | None = None,
    ) -> list[Entry]:
        """Index (or re-index) entries' keyword text from (id, keyword_text) pairs.

        Empty or None `items` is a no-op. Each id must refer to an existing
        entry. Existing fts rows for these ids are replaced. `threshold_rank`
        is stored on every row written in this call.
        """
        if not items:
            return []

        ids = [i for i, _ in items]
        existing = self._conn.execute(
            "SELECT id FROM entry WHERE id IN (SELECT value FROM json_each(?))",
            (json.dumps(ids),),
        ).fetchall()
        known = {r["id"] for r in existing}
        for entry_id in ids:
            if entry_id not in known:
                raise ValueError(f"No entry with id {entry_id!r}")

        return entry.keyword(self._conn, items, threshold_rank=threshold_rank)

    def embed(
        self,
        items: list[tuple[str, str]] | None = None,
        *,
        partition: str | None = None,
        threshold_distance: float | None = None,
    ) -> list[Entry]:
        """Embed (or re-embed) entries from (id, semantic_text) pairs.

        Each pair is written into the given partition. Empty or None `items`
        is a no-op. Each id must refer to an existing entry; the given text
        is embedded and stored on the vec row. Existing vec rows for these
        ids are replaced. `threshold_distance` is stored on every row written
        in this call.
        """
        if not items:
            return []

        ids = [i for i, _ in items]
        texts = [t for _, t in items]

        for text in texts:
            if not text.strip():
                raise ValueError("semantic_text must be non-empty")

        existing = self._conn.execute(
            "SELECT id FROM entry WHERE id IN (SELECT value FROM json_each(?))",
            (json.dumps(ids),),
        ).fetchall()
        known = {r["id"] for r in existing}
        for entry_id in ids:
            if entry_id not in known:
                raise ValueError(f"No entry with id {entry_id!r}")

        embeddings = self.embedder.embed_many(texts)
        return entry.embed(
            self._conn,
            [(i, t, v) for (i, t), v in zip(items, embeddings, strict=True)],
            partition=partition,
            threshold_distance=threshold_distance,
        )

    def keyword_search(
        self,
        query: str,
        filters: Filters | None = None,
        limit: int | None = None,
    ) -> list[KeywordHit]:
        return entry.keyword_search(self._conn, query, filters, limit)

    def semantic_search(
        self,
        query: str,
        partition: str | None = None,
        limit: int = 10,
    ) -> list[SemanticHit]:
        return entry.semantic_search(
            self._conn,
            self.embedder.embed(query),
            partition,
            limit,
        )


def peek(path: str | Path) -> Peek:
    """Inspect a grimoire file without loading an embedder.

    Returns model, dimension, schema version, entry count, per-group counts
    (from `entry`), and per-partition counts (from `entry_vec`). Raises
    `GrimoireNotFound` if the file does not exist or has not been
    initialized.
    """
    p = Path(path)
    if not p.exists():
        raise GrimoireNotFound(f"No grimoire at {p}")

    conn = sqlite3.connect(p)
    conn.row_factory = sqlite3.Row
    conn.enable_load_extension(True)
    sqlite_vec.load(conn)
    conn.enable_load_extension(False)

    try:
        if schema.read_version(conn) == 0:
            raise GrimoireNotFound(f"{p} is not an initialized grimoire")

        schema.validate(conn)
        model = meta.fetch(conn, "model")
        dimension_str = meta.fetch(conn, "dimension")

        if model is None or dimension_str is None:
            raise GrimoireNotFound(f"{p} is missing its embedder lock")

        entry_count = conn.execute("SELECT COUNT(*) FROM entry").fetchone()[0]
        group_rows = conn.execute(
            "SELECT group_key, COUNT(*) AS n FROM entry GROUP BY group_key "
            "ORDER BY group_key IS NULL, group_key"
        ).fetchall()
        partition_rows = conn.execute(
            "SELECT partition, COUNT(*) AS n FROM entry_vec GROUP BY partition "
            "ORDER BY partition IS NULL, partition"
        ).fetchall()

        return Peek(
            model=model,
            dimension=int(dimension_str),
            schema_version=schema.read_version(conn),
            entry_count=entry_count,
            group_counts={r["group_key"]: r["n"] for r in group_rows},
            partition_counts={r["partition"]: r["n"] for r in partition_rows},
        )
    finally:
        conn.close()


def open(path: str | Path, *, embedder: Embedder) -> Grimoire:
    conn = sqlite3.connect(path)
    conn.row_factory = sqlite3.Row

    conn.enable_load_extension(True)
    sqlite_vec.load(conn)
    conn.enable_load_extension(False)

    if schema.read_version(conn) == 0:
        schema.create(conn, model=embedder.model, dimension=embedder.dimension)
    else:
        schema.validate(conn)
        stored_model = meta.fetch(conn, "model")
        stored_dimension = int(meta.fetch(conn, "dimension"))
        if stored_model != embedder.model or stored_dimension != embedder.dimension:
            raise GrimoireMismatch(
                f"Embedder reports model={embedder.model!r}"
                f" dimension={embedder.dimension},"
                f" file locked to model={stored_model!r}"
                f" dimension={stored_dimension}."
            )

    return Grimoire(conn, embedder=embedder)
