backend/database.py

import asyncio
import re
from dataclasses import dataclass
import json
import os
import time
from typing import Any, Dict, List, Optional

from sql_security import strip_trailing_semicolon

try:
    import asyncpg
except Exception:  # pragma: no cover - optional until requirements are installed
    asyncpg = None


def _get_database_url() -> Optional[str]:
    return (
        os.environ.get("READONLY_DATABASE_URL")
        or os.environ.get("SUPABASE_DB_URL")
        or os.environ.get("DATABASE_URL")
        or os.environ.get("POSTGRES_URL")
    )


@dataclass
class QueryExecution:
    success: bool
    rows: List[Dict[str, Any]]
    columns: List[str]
    row_count: int
    execution_ms: float
    execution_plan: Any = None
    indexes_used: Optional[List[str]] = None
    error: Optional[str] = None


class PostgresService:
    def __init__(self) -> None:
        self.database_url = _get_database_url()
        self.pool = None

    @property
    def configured(self) -> bool:
        return bool(self.database_url and asyncpg is not None)

    async def connect(self) -> None:
        if not self.configured or self.pool is not None:
            return
        self.pool = await asyncpg.create_pool(
            dsn=self.database_url,
            min_size=1,
            max_size=int(os.environ.get("POSTGRES_POOL_SIZE", "5")),
            command_timeout=12,
        )

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

    async def execute_select(self, sql: str, limit: int = 100, timeout: float = 10.0) -> QueryExecution:
        if self.pool is None:
            raise RuntimeError("PostgreSQL is not configured")

        safe_sql = strip_trailing_semicolon(sql)
        limited_sql = f"SELECT * FROM ({safe_sql}) AS neuralvault_query LIMIT {int(limit)}"
        explain_sql = f"EXPLAIN (FORMAT JSON, ANALYZE false) {safe_sql}"

        start = time.perf_counter()
        try:
            async with self.pool.acquire() as conn:
                # ELEGANT Spec 5: Sandboxing statement timeouts at transaction session level
                async with conn.transaction():
                    await conn.execute("SET statement_timeout = 2000; SET lock_timeout = 1000;")
                    rows = await conn.fetch(limited_sql, timeout=timeout)
                    plan = await conn.fetchval(explain_sql, timeout=timeout)
        except Exception as exc:
            return QueryExecution(
                success=False,
                rows=[],
                columns=[],
                row_count=0,
                execution_ms=round((time.perf_counter() - start) * 1000, 2),
                error=str(exc),
            )

        dict_rows = [dict(row) for row in rows]
        columns = list(dict_rows[0].keys()) if dict_rows else []
        return QueryExecution(
            success=True,
            rows=dict_rows,
            columns=columns,
            row_count=len(dict_rows),
            execution_ms=round((time.perf_counter() - start) * 1000, 2),
            execution_plan=plan,
            indexes_used=_extract_indexes(plan),
        )

    async def health(self) -> dict:
        if asyncpg is None:
            return {"configured": False, "ok": False, "reason": "asyncpg is not installed"}
        if not self.database_url:
            return {"configured": False, "ok": False, "reason": "No PostgreSQL URL configured"}
        if self.pool is None:
            return {"configured": True, "ok": False, "reason": "Pool not initialized"}

        try:
            async with self.pool.acquire() as conn:
                await conn.fetchval("SELECT 1")
                vector_count = await _safe_fetchval(conn, "SELECT COUNT(*) FROM products WHERE embedding IS NOT NULL")
                hnsw_indexes = await conn.fetch(
                    """
                    SELECT indexname
                    FROM pg_indexes
                    WHERE schemaname = 'public'
                      AND indexdef ILIKE '%hnsw%'
                    """
                )
            return {
                "configured": True,
                "ok": True,
                "vector_count": vector_count,
                "hnsw_indexes": [row["indexname"] for row in hnsw_indexes],
            }
        except Exception as exc:
            return {"configured": True, "ok": False, "reason": str(exc)}

    async def run_startup_validation(self) -> dict:
        if self.pool is None:
            return {"ok": False, "error": "Database not connected"}
        
        try:
            async with self.pool.acquire() as conn:
                # 1. Check pgvector extension
                has_vector = await conn.fetchval(
                    "SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'vector')"
                )
                
                # 2. Check products table
                has_products = await conn.fetchval(
                    "SELECT EXISTS(SELECT 1 FROM information_schema.tables WHERE table_schema = 'public' AND table_name = 'products')"
                )
                
                # 3. Check HNSW index
                hnsw_rows = await conn.fetch(
                    "SELECT indexname FROM pg_indexes WHERE schemaname = 'public' AND tablename = 'products' AND indexdef ILIKE '%hnsw%'"
                )
                has_hnsw = len(hnsw_rows) > 0
                
            return {
                "ok": True,
                "pgvector_extension": has_vector,
                "products_table": has_products,
                "hnsw_index": has_hnsw
            }
        except Exception as exc:
            return {"ok": False, "error": str(exc)}


    async def hybrid_search(
        self,
        query: str,
        embedding: List[float],
        mode: str = "hybrid",
        limit: int = 10,
        ef_search: int = 40,
        price_limit: Optional[float] = None,
    ) -> dict:
        if self.pool is None:
            raise RuntimeError("PostgreSQL is not configured")

        vector_literal = "[" + ",".join(f"{value:.8f}" for value in embedding) + "]"
        timings = {"embedding_ms": 0.0, "vector_ms": 0.0, "fts_ms": 0.0, "fusion_ms": 0.0}
        vector_rows: List[Dict[str, Any]] = []
        text_rows: List[Dict[str, Any]] = []

        async with self.pool.acquire() as conn:
            await conn.execute(f"SET hnsw.ef_search = {int(ef_search)}")

            if mode in ("vector", "hybrid"):
                start = time.perf_counter()
                if price_limit is not None:
                    # Spec 7: Vector Metadata Pre-Filtering using HNSW indexes with WHERE constraints
                    vector_rows = [
                        dict(row)
                        for row in await conn.fetch(
                            """
                            SELECT id::text, title AS name, category, price::text, rating,
                                   1 - (embedding <=> $1::vector) AS vector_score
                            FROM products
                            WHERE embedding IS NOT NULL AND price <= $2
                            ORDER BY embedding <=> $1::vector
                            LIMIT 50
                            """,
                            vector_literal,
                            float(price_limit),
                        )
                    ]
                else:
                    vector_rows = [
                        dict(row)
                        for row in await conn.fetch(
                            """
                            SELECT id::text, title AS name, category, price::text, rating,
                                   1 - (embedding <=> $1::vector) AS vector_score
                            FROM products
                            WHERE embedding IS NOT NULL
                            ORDER BY embedding <=> $1::vector
                            LIMIT 50
                            """,
                            vector_literal,
                        )
                    ]
                timings["vector_ms"] = round((time.perf_counter() - start) * 1000, 2)

            if mode in ("fulltext", "hybrid"):
                start = time.perf_counter()
                if price_limit is not None:
                    text_rows = [
                        dict(row)
                        for row in await conn.fetch(
                            """
                            SELECT id::text, title AS name, category, price::text, rating,
                                   ts_rank(
                                     to_tsvector('english', title || ' ' || COALESCE(description, '')),
                                     plainto_tsquery('english', $1)
                                   ) AS text_score
                            FROM products
                            WHERE to_tsvector('english', title || ' ' || COALESCE(description, ''))
                                  @@ plainto_tsquery('english', $1) AND price <= $2
                            ORDER BY text_score DESC
                            LIMIT 50
                            """,
                            query,
                            float(price_limit),
                        )
                    ]
                else:
                    text_rows = [
                        dict(row)
                        for row in await conn.fetch(
                            """
                            SELECT id::text, title AS name, category, price::text, rating,
                                   ts_rank(
                                     to_tsvector('english', title || ' ' || COALESCE(description, '')),
                                     plainto_tsquery('english', $1)
                                   ) AS text_score
                            FROM products
                            WHERE to_tsvector('english', title || ' ' || COALESCE(description, ''))
                                  @@ plainto_tsquery('english', $1)
                            ORDER BY text_score DESC
                            LIMIT 50
                            """,
                            query,
                        )
                    ]
                timings["fts_ms"] = round((time.perf_counter() - start) * 1000, 2)

        start = time.perf_counter()
        # Stage 1: Blended candidate results via RRF Fusion
        candidates = _merge_search_results(vector_rows, text_rows, mode, limit=50)

        # Spec 8 Stage 2: Cross-Encoder Reranking
        if mode == "hybrid" and candidates:
            try:
                from embedding_service import rerank_candidates
                results = await rerank_candidates(query, candidates, limit=limit)
            except Exception as exc:
                print(f"Failed to run Cross-Encoder reranker: {exc}")
                results = candidates[:limit]
        else:
            results = candidates[:limit]

        timings["fusion_ms"] = round((time.perf_counter() - start) * 1000, 2)
        timings["total_ms"] = round(sum(timings.values()), 2)

        return {
            "results": results,
            "timings": timings,
            "ef_search": ef_search,
            "mode": mode,
            "price_limit": price_limit,
        }


async def _safe_fetchval(conn, sql: str) -> Any:
    try:
        return await conn.fetchval(sql)
    except Exception:
        return None


def _merge_search_results(vector_rows: List[dict], text_rows: List[dict], mode: str, limit: int) -> List[dict]:
    if mode == "vector":
        return [_format_search_row(row, vector_rank=i + 1) for i, row in enumerate(vector_rows[:limit])]
    if mode == "fulltext":
        return [_format_search_row(row, text_rank=i + 1) for i, row in enumerate(text_rows[:limit])]

    by_id: Dict[str, dict] = {}
    for rank, row in enumerate(vector_rows, start=1):
        item = by_id.setdefault(row["id"], row.copy())
        item["vector_rank"] = rank
        item["vector_score"] = float(row.get("vector_score") or 0)
    for rank, row in enumerate(text_rows, start=1):
        item = by_id.setdefault(row["id"], row.copy())
        item["text_rank"] = rank
        item["text_score"] = float(row.get("text_score") or 0)

    fused = []
    for item in by_id.values():
        vector_rank = item.get("vector_rank")
        text_rank = item.get("text_rank")
        rrf_score = 0.0
        if vector_rank:
            rrf_score += 1 / (60 + vector_rank)
        if text_rank:
            rrf_score += 1 / (60 + text_rank)
        formatted = _format_search_row(item, vector_rank=vector_rank, text_rank=text_rank)
        formatted["rrf_score"] = round(rrf_score, 4)
        fused.append(formatted)

    return sorted(fused, key=lambda row: row["rrf_score"], reverse=True)[:limit]


def _format_search_row(row: dict, vector_rank: Optional[int] = None, text_rank: Optional[int] = None) -> dict:
    vector_score = float(row.get("vector_score") or 0)
    text_score = float(row.get("text_score") or 0)
    rrf_score = row.get("rrf_score")
    if rrf_score is None:
        ranks = [rank for rank in (vector_rank, text_rank) if rank]
        rrf_score = sum(1 / (60 + rank) for rank in ranks) if ranks else 0

    reasons = []
    if vector_rank:
        reasons.append(f"semantic rank #{vector_rank}")
    if text_rank:
        reasons.append(f"full-text rank #{text_rank}")

    return {
        "id": row.get("id"),
        "name": row.get("name"),
        "category": row.get("category"),
        "price": row.get("price") or "n/a",
        "rating": row.get("rating"),
        "vector_score": round(vector_score, 4),
        "text_score": round(text_score, 4),
        "rrf_score": round(float(rrf_score), 4),
        "why_matched": "Matched by " + " and ".join(reasons) if reasons else "Matched product search query",
        "in_stock": True,
    }


def _extract_indexes(plan: Any) -> List[str]:
    if not plan:
        return []
    if isinstance(plan, str):
        try:
            plan = json.loads(plan)
        except Exception:
            return sorted(set(re.findall(r"Index Scan using ([a-zA-Z0-9_]+)", plan)))

    indexes = set()

    def walk(node: Any) -> None:
        if isinstance(node, list):
            for item in node:
                walk(item)
        elif isinstance(node, dict):
            if "Index Name" in node:
                indexes.add(node["Index Name"])
            for value in node.values():
                walk(value)

    walk(plan)
    return sorted(indexes)


postgres = PostgresService()