from __future__ import annotations

import dataclasses
import json

from dot_search.embed import get_embedder
from dot_search.fusion import reciprocal_rank_fusion
from dot_search.store import make_store
from dot_search.store.protocol import BM25Search, ExactSearch, SubSearch
from dot_search.types import SearchConfig, SearchResult, TableConfig


class SearchEngine:
    def __init__(self, db_url: str) -> None:
        self._store = make_store(db_url)

    async def _ensure_config_table(self) -> None:
        await self._store.ensure_config_table()

    async def _save_config(self, config: TableConfig) -> None:
        data = dataclasses.asdict(config)
        await self._store.save_config(config.index_id, json.dumps(data))

    async def _load_config(self, index_id: str) -> TableConfig:
        raw = await self._store.load_config(index_id)
        if raw is None:
            raise ValueError(f"No config found for index_id '{index_id}'. Call index(config=...) first.")
        return TableConfig.from_json(raw)

    async def index(self, config: TableConfig) -> None:
        await self._ensure_config_table()
        await self._save_config(config)
        await self._run_indexing(config)

    async def _run_indexing(self, config: TableConfig) -> None:
        existing_columns = set(await self._store.get_columns(config.table))

        for embedding_config in config.embeddings:
            if embedding_config.source_column not in existing_columns:
                raise ValueError(
                    f"EmbeddingConfig source_column {embedding_config.source_column!r} "
                    f"does not exist in table {config.table!r}. "
                    f"Available columns: {sorted(existing_columns)}"
                )

        for bm25cfg in config.bm25:
            if bm25cfg.source_column not in existing_columns:
                raise ValueError(
                    f"BM25Config source_column {bm25cfg.source_column!r} "
                    f"does not exist in table {config.table!r}. "
                    f"Available columns: {sorted(existing_columns)}"
                )

        for embedding_config in config.embeddings:
            embedder = get_embedder(embedding_config.provider, embedding_config.model)
            target_col = embedding_config.target_column
            rows = await self._store.read_column(config.table, config.pk_column, embedding_config.source_column)

            if not rows:
                continue

            ids = [r[0] for r in rows]
            texts = [r[1] for r in rows]

            await self._store.add_embedding_column(config.table, target_col, embedder.dimension)

            checked = False
            for batch_start in range(0, len(texts), config.batch_size):
                batch_texts = texts[batch_start : batch_start + config.batch_size]
                batch_ids = ids[batch_start : batch_start + config.batch_size]
                batch_vecs = await embedder.fn(batch_texts)
                if not checked and batch_vecs and len(batch_vecs[0]) != embedder.dimension:
                    raise ValueError(
                        f"Embedder dimension mismatch: configured {embedder.dimension}, "
                        f"but fn returned vectors of size {len(batch_vecs[0])}."
                    )
                checked = True
                await self._store.write_embeddings(
                    config.table,
                    target_col,
                    config.pk_column,
                    list(zip(batch_ids, batch_vecs)),
                )

        if config.bm25:
            await self._store.add_bm25_index(config.table, [t.source_column for t in config.bm25])

    async def search(
        self,
        query: str,
        index_id: str,
        config: SearchConfig | None = None,
    ) -> list[SearchResult]:
        if config is None:
            config = SearchConfig()

        await self._ensure_config_table()
        table_config = await self._load_config(index_id)

        strategy = config.strategy
        limit = config.limit
        weights_map = config.weights

        subsearches: list[SubSearch] = []
        weights: list[float] = []

        if strategy in ("vector", "hybrid"):
            for embedding_config in table_config.embeddings:
                embedder = get_embedder(embedding_config.provider, embedding_config.model)
                target_col = embedding_config.target_column
                w = weights_map.get(embedding_config.key, 1.0)
                query_embedding = (await embedder.fn([query]))[0]
                subsearches.append(SubSearch(
                    embedding_column=target_col,
                    query_embedding=query_embedding,
                    weight=w,
                ))
                weights.append(w)

        bm25searches = []
        if strategy in ("bm25", "hybrid"):
            if strategy == "bm25" and not table_config.bm25:
                raise ValueError(f"strategy='bm25' but index '{index_id}' has no bm25 configured.")
            bm25searches = [
                BM25Search(
                    column=t.source_column,
                    query=query,
                    weight=weights_map.get(t.key, 1.0),
                )
                for t in table_config.bm25
            ]
            weights.extend(s.weight for s in bm25searches)

        exact_searches = []
        if strategy in ("exact", "hybrid"):
            if strategy == "exact" and not table_config.exact:
                raise ValueError(f"strategy='exact' but index '{index_id}' has no exact config.")
            exact_searches = [
                ExactSearch(
                    column=t.source_column,
                    query=query,
                    weight=weights_map.get(t.key, 1.0),
                )
                for t in table_config.exact
            ]
            weights.extend(s.weight for s in exact_searches)

        ranked_lists = await self._store.search(
            table_config.table,
            table_config.pk_column,
            subsearches,
            bm25searches,
            exact_searches,
            config.where,
            limit,
        )

        return reciprocal_rank_fusion(ranked_lists, weights)[:limit]

    async def close(self) -> None:
        await self._store.close()
