from __future__ import annotations

import sqlite3
import struct
from typing import Any

import aiosqlite

from dot_search.store.protocol import BM25Search, ExactSearch, SubSearch, validate_embedding_column
from dot_search.types import SearchResult


class SQLiteStore:
    def __init__(self, db_url: str) -> None:
        import sqlite_vec
        path = db_url.replace("sqlite+aiosqlite:///", "").replace("sqlite+aiosqlite://", "")
        if not path or path == ":memory:":
            path = ":memory:"
        self._path = path
        self._vec_path = sqlite_vec.loadable_path()
        self._conn: aiosqlite.Connection | None = None

    async def _get_conn(self) -> aiosqlite.Connection:
        if self._conn is None:
            self._conn = await aiosqlite.connect(self._path)
            self._conn.row_factory = sqlite3.Row
            await self._conn.enable_load_extension(True)
            await self._conn.load_extension(self._vec_path)
            await self._conn.enable_load_extension(False)
        return self._conn

    async def add_embedding_column(self, table: str, column: str, dimension: int) -> None:
        validate_embedding_column(column)
        conn = await self._get_conn()
        await conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} BLOB")
        await conn.commit()

    async def add_text_column(self, table: str, column: str) -> None:
        conn = await self._get_conn()
        await conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} TEXT")
        await conn.commit()

    async def write_texts(self, table: str, column: str, pk_column: str, rows: list[tuple[Any, str]]) -> None:
        if not rows:
            return
        conn = await self._get_conn()
        for id_, text in rows:
            await conn.execute(
                f"UPDATE {table} SET {column} = ? WHERE {pk_column} = ?",
                (text, id_),
            )
        await conn.commit()

    async def add_bm25_index(self, table: str, columns: list[str]) -> None:
        pass

    async def write_embeddings(
        self,
        table: str,
        column: str,
        pk_column: str,
        rows: list[tuple[Any, list[float]]],
    ) -> None:
        if not rows:
            return
        conn = await self._get_conn()
        for id_, vec in rows:
            blob = struct.pack(f"{len(vec)}f", *vec)
            await conn.execute(
                f"UPDATE {table} SET {column} = ? WHERE {pk_column} = ?",
                (blob, id_),
            )
        await conn.commit()

    async def read_column(
        self,
        table: str,
        pk_column: str,
        source_column: str,
    ) -> list[tuple[Any, str]]:
        conn = await self._get_conn()
        cursor = await conn.execute(
            f"SELECT {pk_column}, {source_column} FROM {table}"
        )
        rows = await cursor.fetchall()
        return [(row[0], str(row[1]) if row[1] is not None else "") for row in rows]

    async def search(
        self,
        table: str,
        pk_column: str,
        subsearches: list[SubSearch],
        bm25searches: list[BM25Search],
        exact_searches: list[ExactSearch],
        where: str | None,
        limit: int,
    ) -> list[list[SearchResult]]:
        where_clause = f"AND {where}" if where else ""
        conn = await self._get_conn()
        results = []

        # Get payload columns (non-ds_)
        cursor = await conn.execute(f"SELECT * FROM {table} LIMIT 0")
        all_columns = [desc[0] for desc in cursor.description]
        payload_cols = [c for c in all_columns if not c.startswith("ds_")]

        for sub in subsearches:
            dim = len(sub.query_embedding)
            blob = struct.pack(f"{dim}f", *sub.query_embedding)

            select_cols = ", ".join(f"t.{c}" for c in payload_cols)
            sql = f"""
                SELECT {select_cols},
                       vec_distance_cosine(t.{sub.embedding_column}, ?) AS distance
                FROM {table} t
                WHERE t.{sub.embedding_column} IS NOT NULL
                {where_clause}
                ORDER BY distance ASC
                LIMIT ?
            """
            cursor = await conn.execute(sql, (blob, limit))
            rows = await cursor.fetchall()
            col_names = [desc[0] for desc in cursor.description]

            search_results = []
            for row in rows:
                row_dict = dict(zip(col_names, row))
                distance = row_dict.pop("distance")
                score = 1.0 - float(distance)
                id_val = row_dict[pk_column]
                search_results.append(SearchResult(id=id_val, score=score, row=row_dict))
            results.append(search_results)

        for exact in exact_searches:
            select_cols = ", ".join(payload_cols)
            sql = f"""
                SELECT {select_cols}
                FROM {table}
                WHERE {exact.column} LIKE ?
                {where_clause}
                LIMIT ?
            """
            cursor = await conn.execute(sql, (f"%{exact.query}%", limit))
            rows = await cursor.fetchall()
            col_names = [desc[0] for desc in cursor.description]

            search_results = []
            for row in rows:
                row_dict = dict(zip(col_names, row))
                id_val = row_dict[pk_column]
                search_results.append(SearchResult(id=id_val, score=1.0, row=row_dict))
            results.append(search_results)

        return results

    async def get_columns(self, table: str) -> list[str]:
        conn = await self._get_conn()
        cursor = await conn.execute(f"PRAGMA table_info({table})")
        rows = await cursor.fetchall()
        return [row[1] for row in rows]

    async def ensure_config_table(self) -> None:
        conn = await self._get_conn()
        await conn.execute(
            """
            CREATE TABLE IF NOT EXISTS ds_configs (
                table_name TEXT PRIMARY KEY,
                config_json TEXT NOT NULL
            )
            """
        )
        await conn.commit()

    async def save_config(self, index_id: str, config_json: str) -> None:
        conn = await self._get_conn()
        await conn.execute(
            "INSERT INTO ds_configs (table_name, config_json) VALUES (?, ?) ON CONFLICT (table_name) DO UPDATE SET config_json = excluded.config_json",
            (index_id, config_json),
        )
        await conn.commit()

    async def load_config(self, index_id: str) -> str | None:
        conn = await self._get_conn()
        cursor = await conn.execute(
            "SELECT config_json FROM ds_configs WHERE table_name = ?",
            (index_id,),
        )
        row = await cursor.fetchone()
        return row[0] if row else None

    async def close(self) -> None:
        if self._conn is not None:
            await self._conn.close()
            self._conn = None
