"""PostgreSQL full-text search backend using asyncpg.

Leverages ``tsvector`` / ``tsquery`` for full-text search, ``pg_trgm`` for
fuzzy matching, GIN indexes for performance, and ``ts_headline`` for
highlighted snippets.
"""

from __future__ import annotations

import json
import time
from typing import TYPE_CHECKING, Any

from search.backends.base import SearchBackend
from search.models import (
    IndexConfig,
    SearchDocument,
    SearchQuery,
    SearchResponse,
    SearchResult,
)

if TYPE_CHECKING:
    import asyncpg


def _table_name(index_name: str) -> str:
    """Sanitize index name into a valid table name."""
    safe = "".join(c if c.isalnum() else "_" for c in index_name)
    return f"search_{safe}"


class PostgresBackend(SearchBackend):
    """PostgreSQL backend using tsvector, pg_trgm, and GIN indexes.

    Requires the ``pg_trgm`` extension to be available. Each search index maps
    to a dedicated table with a ``tsvector`` column and GIN index.
    """

    def __init__(self, pool: asyncpg.Pool, schema: str = "public") -> None:
        self._pool = pool
        self._schema = schema
        self._configs: dict[str, IndexConfig] = {}

    async def create_index(self, config: IndexConfig) -> None:
        table = _table_name(config.name)
        fq_table = f"{self._schema}.{table}"

        weight_cases = []
        for i, field in enumerate(config.searchable_fields):
            weight = "A" if i == 0 else ("B" if i == 1 else "C")
            weight_cases.append(
                f"setweight(to_tsvector('english', coalesce("
                f"content->>'{field}', '')), '{weight}')"
            )
        tsvector_expr = " || ".join(weight_cases) if weight_cases else "to_tsvector('')"

        async with self._pool.acquire() as conn:
            await conn.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm")
            await conn.execute(f"""
                CREATE TABLE IF NOT EXISTS {fq_table} (
                    id TEXT PRIMARY KEY,
                    content JSONB NOT NULL,
                    tsv tsvector GENERATED ALWAYS AS ({tsvector_expr}) STORED,
                    created_at TIMESTAMPTZ DEFAULT now(),
                    updated_at TIMESTAMPTZ DEFAULT now()
                )
            """)
            await conn.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{table}_tsv ON {fq_table} USING GIN (tsv)"
            )
            await conn.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{table}_content ON {fq_table} USING GIN (content)"
            )
            for field in config.filterable_fields:
                idx_name = f"idx_{table}_{field}"
                await conn.execute(
                    f"CREATE INDEX IF NOT EXISTS {idx_name} "
                    f"ON {fq_table} USING BTREE ((content->>'{field}'))"
                )
        self._configs[config.name] = config

    async def delete_index(self, index_name: str) -> None:
        table = _table_name(index_name)
        async with self._pool.acquire() as conn:
            await conn.execute(f"DROP TABLE IF EXISTS {self._schema}.{table} CASCADE")
        self._configs.pop(index_name, None)

    async def index(self, document: SearchDocument) -> None:
        table = _table_name(document.index_name)
        fq_table = f"{self._schema}.{table}"
        content_json = json.dumps(document.content)
        async with self._pool.acquire() as conn:
            await conn.execute(
                f"""
                INSERT INTO {fq_table} (id, content, updated_at)
                VALUES ($1, $2::jsonb, now())
                ON CONFLICT (id) DO UPDATE SET
                    content = EXCLUDED.content,
                    updated_at = now()
                """,
                document.id,
                content_json,
            )

    async def bulk_index(self, documents: list[SearchDocument]) -> int:
        if not documents:
            return 0
        by_index: dict[str, list[SearchDocument]] = {}
        for doc in documents:
            by_index.setdefault(doc.index_name, []).append(doc)

        count = 0
        for index_name, docs in by_index.items():
            table = _table_name(index_name)
            fq_table = f"{self._schema}.{table}"
            records = [(d.id, json.dumps(d.content)) for d in docs]
            async with self._pool.acquire() as conn:
                await conn.executemany(
                    f"""
                    INSERT INTO {fq_table} (id, content, updated_at)
                    VALUES ($1, $2::jsonb, now())
                    ON CONFLICT (id) DO UPDATE SET
                        content = EXCLUDED.content,
                        updated_at = now()
                    """,
                    records,
                )
            count += len(docs)
        return count

    async def delete(self, index_name: str, document_id: str) -> bool:
        table = _table_name(index_name)
        fq_table = f"{self._schema}.{table}"
        async with self._pool.acquire() as conn:
            result = await conn.execute(
                f"DELETE FROM {fq_table} WHERE id = $1", document_id
            )
            return result == "DELETE 1"

    async def search(self, index_name: str, query: SearchQuery) -> SearchResponse:
        start = time.monotonic()
        table = _table_name(index_name)
        fq_table = f"{self._schema}.{table}"
        config = self._configs.get(index_name)

        conditions: list[str] = []
        params: list[Any] = []
        param_idx = 1

        if query.query:
            conditions.append(f"tsv @@ plainto_tsquery('english', ${param_idx})")
            params.append(query.query)
            param_idx += 1

        for key, value in query.filters.items():
            if isinstance(value, list):
                placeholders = ", ".join(f"${param_idx + i}" for i in range(len(value)))
                conditions.append(f"content->>'{key}' IN ({placeholders})")
                params.extend(str(v) for v in value)
                param_idx += len(value)
            else:
                conditions.append(f"content->>'{key}' = ${param_idx}")
                params.append(str(value))
                param_idx += 1

        where = " AND ".join(conditions) if conditions else "TRUE"

        score_expr = (
            f"ts_rank_cd(tsv, plainto_tsquery('english', $1))"
            if query.query
            else "1.0"
        )

        order = f"{score_expr} DESC"
        if query.sort_by:
            desc = query.sort_by.startswith("-")
            field = query.sort_by.lstrip("-")
            direction = "DESC" if desc else "ASC"
            order = f"content->>'{field}' {direction}"

        offset = (query.page - 1) * query.per_page

        async with self._pool.acquire() as conn:
            count_row = await conn.fetchrow(
                f"SELECT COUNT(*) as cnt FROM {fq_table} WHERE {where}", *params
            )
            total = count_row["cnt"] if count_row else 0

            rows = await conn.fetch(
                f"""
                SELECT id, content, {score_expr} as score
                FROM {fq_table}
                WHERE {where}
                ORDER BY {order}
                LIMIT {query.per_page} OFFSET {offset}
                """,
                *params,
            )

            results: list[SearchResult] = []
            for row in rows:
                content = json.loads(row["content"]) if isinstance(row["content"], str) else row["content"]
                highlights: dict[str, str] = {}
                if query.highlight and query.query:
                    highlights = await self._get_highlights(
                        conn, fq_table, row["id"], query.query, config
                    )
                results.append(
                    SearchResult(
                        id=row["id"],
                        score=float(row["score"]),
                        highlights=highlights,
                        document=content,
                    )
                )

            facets = await self._compute_facets(conn, fq_table, where, params, query.facets)

        elapsed = (time.monotonic() - start) * 1000
        return SearchResponse(
            results=results,
            total=total,
            facets=facets,
            query_time_ms=round(elapsed, 2),
            page=query.page,
            per_page=query.per_page,
        )

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

    # -- private helpers -------------------------------------------------------

    @staticmethod
    async def _get_highlights(
        conn: Any,
        fq_table: str,
        doc_id: str,
        query: str,
        config: IndexConfig | None,
    ) -> dict[str, str]:
        fields = config.searchable_fields if config else ["title", "body"]
        highlights: dict[str, str] = {}
        for field in fields:
            row = await conn.fetchrow(
                f"""
                SELECT ts_headline(
                    'english',
                    coalesce(content->>'{field}', ''),
                    plainto_tsquery('english', $1),
                    'StartSel=<em>, StopSel=</em>, MaxFragments=2, FragmentDelimiter=...'
                ) as snippet
                FROM {fq_table}
                WHERE id = $2
                """,
                query,
                doc_id,
            )
            if row and row["snippet"] and "<em>" in row["snippet"]:
                highlights[field] = row["snippet"]
        return highlights

    @staticmethod
    async def _compute_facets(
        conn: Any,
        fq_table: str,
        where: str,
        params: list[Any],
        facet_fields: list[str],
    ) -> dict[str, dict[str, int]]:
        if not facet_fields:
            return {}
        facets: dict[str, dict[str, int]] = {}
        for field in facet_fields:
            rows = await conn.fetch(
                f"""
                SELECT content->>'{field}' as val, COUNT(*) as cnt
                FROM {fq_table}
                WHERE {where} AND content->>'{field}' IS NOT NULL
                GROUP BY content->>'{field}'
                ORDER BY cnt DESC
                LIMIT 50
                """,
                *params,
            )
            facets[field] = {row["val"]: row["cnt"] for row in rows}
        return facets
