from __future__ import annotations

from typing import Any

import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine

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


class PgStore:
    def __init__(self, db_url: str) -> None:
        self._engine = create_async_engine(db_url)
        self._payload_cols_cache: dict[str, list[str]] = {}

    async def _get_payload_cols(self, conn: Any, table: str) -> list[str]:
        if table not in self._payload_cols_cache:
            result = await conn.execute(
                sa.text("SELECT column_name FROM information_schema.columns WHERE table_name = :table ORDER BY ordinal_position"),
                {"table": table},
            )
            self._payload_cols_cache[table] = [r[0] for r in result if not r[0].startswith("ds_")]
        return self._payload_cols_cache[table]

    async def add_embedding_column(self, table: str, column: str, dimension: int) -> None:
        validate_embedding_column(column)
        async with self._engine.begin() as conn:
            await conn.execute(sa.text("CREATE EXTENSION IF NOT EXISTS vector"))
            # Add column if not exists
            await conn.execute(sa.text(
                f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS {column} vector({dimension})"
            ))

    async def add_text_column(self, table: str, column: str) -> None:
        async with self._engine.begin() as conn:
            await conn.execute(sa.text(
                f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS {column} TEXT"
            ))

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

    async def add_bm25_index(self, table: str, columns: list[str]) -> None:
        index_name = f"{table}_bm25_idx"
        async with self._engine.begin() as conn:
            pk_result = await conn.execute(sa.text(
                f"""
                SELECT a.attname
                FROM pg_index i
                JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
                WHERE i.indrelid = '{table}'::regclass AND i.indisprimary
                LIMIT 1
                """
            ))
            pk_row = pk_result.fetchone()
            id_col = pk_row[0] if pk_row else "id"

            all_cols = ", ".join([id_col] + columns)
            await conn.execute(sa.text(f"DROP INDEX IF EXISTS {index_name}"))
            await conn.execute(sa.text(
                f"""
                CREATE INDEX {index_name} ON {table}
                USING bm25 ({all_cols})
                WITH (key_field='{id_col}')
                """
            ))

    async def write_embeddings(
        self,
        table: str,
        column: str,
        pk_column: str,
        rows: list[tuple[Any, list[float]]],
    ) -> None:
        if not rows:
            return
        async with self._engine.begin() as conn:
            await conn.execute(
                sa.text(f"UPDATE {table} SET {column} = :vec WHERE {pk_column} = :id"),
                [{"vec": str(vec), "id": id_} for id_, vec in rows],
            )

    async def read_column(
        self,
        table: str,
        pk_column: str,
        source_column: str,
    ) -> list[tuple[Any, str]]:
        async with self._engine.connect() as conn:
            result = await conn.execute(sa.text(
                f"SELECT {pk_column}, {source_column} FROM {table}"
            ))
            return [(row[0], str(row[1]) if row[1] is not None else "") for row in result]

    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 ""
        results = []

        async with self._engine.connect() as conn:
            payload_cols = await self._get_payload_cols(conn, table)
            select_cols = ", ".join(payload_cols)

            for sub in subsearches:
                vec_str = str(sub.query_embedding)
                sql = sa.text(
                    f"""
                    SELECT {select_cols},
                           1 - ({sub.embedding_column} <=> '{vec_str}'::vector) AS score
                    FROM {table}
                    WHERE {sub.embedding_column} IS NOT NULL
                    {where_clause}
                    ORDER BY {sub.embedding_column} <=> '{vec_str}'::vector
                    LIMIT :limit
                    """
                )
                result = await conn.execute(sql, {"limit": limit})
                rows = result.fetchall()
                col_names = list(result.keys())

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

            for bm25 in bm25searches:
                bm25_query = f'{bm25.column}:"{bm25.query}"'
                sql = sa.text(
                    f"""
                    SELECT {select_cols},
                           paradedb.score({pk_column}) AS score
                    FROM {table}
                    WHERE {table} @@@ '{bm25_query}'
                    {where_clause}
                    ORDER BY score DESC
                    LIMIT :limit
                    """
                )
                result = await conn.execute(sql, {"limit": limit})
                rows = result.fetchall()
                col_names = list(result.keys())

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

            for exact in exact_searches:
                sql = sa.text(
                    f"""
                    SELECT {select_cols}
                    FROM {table}
                    WHERE {exact.column} ILIKE :pattern
                    {where_clause}
                    LIMIT :limit
                    """
                )
                result = await conn.execute(sql, {"pattern": f"%{exact.query}%", "limit": limit})
                rows = result.fetchall()
                col_names = list(result.keys())

                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]:
        async with self._engine.connect() as conn:
            result = await conn.execute(
                sa.text("SELECT column_name FROM information_schema.columns WHERE table_name = :table ORDER BY ordinal_position"),
                {"table": table},
            )
            return [r[0] for r in result]

    async def ensure_config_table(self) -> None:
        async with self._engine.begin() as conn:
            await conn.execute(sa.text(
                """
                CREATE TABLE IF NOT EXISTS ds_configs (
                    table_name TEXT PRIMARY KEY,
                    config_json TEXT NOT NULL
                )
                """
            ))

    async def save_config(self, index_id: str, config_json: str) -> None:
        async with self._engine.begin() as conn:
            await conn.execute(sa.text(
                """
                INSERT INTO ds_configs (table_name, config_json)
                VALUES (:index_id, :config_json)
                ON CONFLICT (table_name) DO UPDATE SET config_json = EXCLUDED.config_json
                """
            ), {"index_id": index_id, "config_json": config_json})

    async def load_config(self, index_id: str) -> str | None:
        async with self._engine.connect() as conn:
            result = await conn.execute(
                sa.text("SELECT config_json FROM ds_configs WHERE table_name = :index_id"),
                {"index_id": index_id},
            )
            row = result.fetchone()
            return row[0] if row else None

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