"""
PostgreSQL pg_vector-based context-aware Genie cache implementation.

This module provides a context-aware cache that uses PostgreSQL with pg_vector
for semantic similarity search. It supports both standard PostgreSQL and
Databricks Lakebase connections via the DatabaseModel abstraction.

Features:
- Dual embedding matching (question + conversation context)
- pg_vector cosine similarity search
- Prompt history tracking for conversation context
- TTL-based expiration with refresh-on-hit
- Space-partitioned cache entries
"""

from __future__ import annotations

import json
import math
from datetime import datetime, timedelta
from typing import Any, Self

import mlflow
from databricks.sdk import WorkspaceClient
from loguru import logger
from psycopg import sql

from dao_ai.config import (
    DatabaseModel,
    GenieContextAwareCacheParametersModel,
    WarehouseModel,
)
from dao_ai.genie.cache.base import (
    CacheResult,
    GenieServiceBase,
    SQLCacheEntry,
)
from dao_ai.genie.cache.context_aware.persistent import (
    DbRow,
    PersistentContextAwareGenieCacheService,
)
from dao_ai.genie.core import GenieResponse


def _serialize_embedding(embedding: list[float]) -> str:
    """Serialize embedding vector to pg_vector string format using C-optimized json.dumps."""
    return json.dumps(embedding, separators=(",", ":"))


class PostgresContextAwareGenieService(PersistentContextAwareGenieCacheService):
    """
    PostgreSQL pg_vector-based context-aware caching decorator.

    This service caches the SQL query generated by Genie along with dual embeddings
    (question + conversation context) for high-precision semantic matching. On
    subsequent queries, it performs similarity search using pg_vector to find
    cached queries that match both the question intent AND conversation context.

    Supports both standard PostgreSQL and Databricks Lakebase via DatabaseModel.

    Cache entries are partitioned by genie_space_id to ensure queries from different
    Genie spaces don't return incorrect cache hits.

    On cache hit, it re-executes the cached SQL using the provided warehouse
    to return fresh data while avoiding the Genie NL-to-SQL translation cost.

    Example:
        from dao_ai.config import GenieContextAwareCacheParametersModel, DatabaseModel
        from dao_ai.genie.cache.context_aware import PostgresContextAwareGenieService

        cache_params = GenieContextAwareCacheParametersModel(
            database=database_model,
            warehouse=warehouse_model,
            embedding_model="databricks-gte-large-en",
            time_to_live_seconds=86400,  # 24 hours
            similarity_threshold=0.85
        )
        genie = PostgresContextAwareGenieService(
            impl=GenieService(Genie(space_id="my-space")),
            parameters=cache_params
        )

    Thread-safe: Uses connection pooling from psycopg_pool.
    """

    impl: GenieServiceBase
    parameters: GenieContextAwareCacheParametersModel
    _workspace_client: WorkspaceClient | None
    name: str
    _embeddings: Any  # DatabricksEmbeddings
    _pool: Any  # ConnectionPool
    _embedding_dims: int | None
    _setup_complete: bool

    def __init__(
        self,
        impl: GenieServiceBase,
        parameters: GenieContextAwareCacheParametersModel,
        workspace_client: WorkspaceClient | None = None,
        name: str | None = None,
    ) -> None:
        """
        Initialize the PostgreSQL context-aware cache service.

        Args:
            impl: The underlying GenieServiceBase to delegate to on cache miss.
                The space_id will be obtained from impl.space_id.
            parameters: Cache configuration including database, warehouse, embedding model
            workspace_client: Optional WorkspaceClient for retrieving conversation history.
                If None, conversation context will not be used.
            name: Name for this cache layer (for logging). Defaults to class name.
        """
        self.impl = impl
        self.parameters = parameters
        self._workspace_client = workspace_client
        self.name = name if name is not None else self.__class__.__name__
        self._embeddings = None
        self._pool = None
        self._embedding_dims = None
        self._setup_complete = False
        self._prompt_stored_for_current_request = False
        self._resolved_ivfflat_lists: int = (
            100  # Set during _create_table_if_not_exists
        )

    def _setup(self) -> None:
        """Initialize embeddings and database connection pool lazily."""
        if self._setup_complete:
            return

        from dao_ai.memory.postgres import PostgresPoolManager

        # Initialize embeddings using base class helper
        self._initialize_embeddings(
            self.parameters.embedding_model,
            self.parameters.embedding_dims,
        )

        # Get connection pool
        self._pool = PostgresPoolManager.get_pool(self.parameters.database)

        # Ensure table exists
        self._create_table_if_not_exists()

        self._setup_complete = True
        logger.debug(
            "PostgreSQL context-aware cache initialized",
            layer=self.name,
            space_id=self.space_id,
            table_name=self.table_name,
            dims=self._embedding_dims,
        )

    # Property implementations
    @property
    def database(self) -> DatabaseModel:
        """The database used for storing cache entries."""
        return self.parameters.database

    @property
    def warehouse(self) -> WarehouseModel:
        """The warehouse used for executing cached SQL queries."""
        return self.parameters.warehouse

    @property
    def time_to_live(self) -> timedelta | None:
        """Time-to-live for cache entries. None means never expires."""
        ttl = self.parameters.time_to_live_seconds
        if ttl is None or ttl < 0:
            return None
        return timedelta(seconds=ttl)

    @property
    def time_to_live_seconds(self) -> int | None:
        """TTL in seconds (None or negative = never expires)."""
        return self.parameters.time_to_live_seconds

    @property
    def similarity_threshold(self) -> float:
        """Minimum similarity for cache hit (using L2 distance converted to similarity)."""
        return self.parameters.similarity_threshold

    @property
    def context_similarity_threshold(self) -> float:
        """Minimum similarity for context matching."""
        return self.parameters.context_similarity_threshold

    @property
    def question_weight(self) -> float:
        """Weight for question similarity in combined score."""
        return self.parameters.question_weight

    @property
    def context_weight(self) -> float:
        """Weight for context similarity in combined score."""
        return self.parameters.context_weight

    @property
    def table_name(self) -> str:
        """Name of the cache table."""
        return self.parameters.table_name

    @property
    def prompt_history_table(self) -> str:
        """Name of the prompt history table."""
        return self.parameters.prompt_history_table

    @property
    def context_window_size(self) -> int:
        """Number of previous prompts to include in context."""
        return self.parameters.context_window_size

    @property
    def max_context_tokens(self) -> int:
        """Maximum tokens for context string."""
        return self.parameters.max_context_tokens

    @property
    def max_prompt_history_length(self) -> int:
        """Maximum number of prompts to keep per conversation."""
        return self.parameters.max_prompt_history_length

    def _create_table_if_not_exists(self) -> None:
        """Create the cache table and prompt history table with pg_vector extension."""
        create_extension_sql: str = "CREATE EXTENSION IF NOT EXISTS vector"

        # Check if table exists and get current embedding dimensions
        check_dims_sql: str = """
            SELECT atttypmod 
            FROM pg_attribute 
            WHERE attrelid = %s::regclass 
              AND attname = 'question_embedding'
        """

        table_id = sql.Identifier(self.table_name)
        dims = self.embedding_dims

        create_table_sql = sql.SQL("""
            CREATE TABLE IF NOT EXISTS {} (
                id SERIAL PRIMARY KEY,
                genie_space_id TEXT NOT NULL,
                question TEXT NOT NULL,
                conversation_context TEXT,
                context_string TEXT,
                question_embedding vector({}),
                context_embedding vector({}),
                sql_query TEXT NOT NULL,
                description TEXT,
                conversation_id TEXT,
                message_id TEXT,
                created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
            )
        """).format(table_id, sql.Literal(dims), sql.Literal(dims))

        # Migration: Add message_id column if it doesn't exist
        add_message_id_sql = sql.SQL("""
            ALTER TABLE {} 
            ADD COLUMN IF NOT EXISTS message_id TEXT
        """).format(table_id)
        # Compute ivfflat lists parameter.
        # Explicit config value takes priority; otherwise auto-compute from table size.
        ivfflat_lists = self.parameters.ivfflat_lists
        # ivfflat_lists will be resolved after table creation (needs row count)

        # Index for efficient similarity search partitioned by genie_space_id
        # The lists parameter is set dynamically below after table creation.
        def _make_ivfflat_index_sql(
            idx_name: str, column: str, lists_val: int
        ) -> sql.Composed:
            return sql.SQL("""
                CREATE INDEX IF NOT EXISTS {} 
                ON {} 
                USING ivfflat ({} vector_cosine_ops)
                WITH (lists = {})
            """).format(
                sql.Identifier(idx_name),
                table_id,
                sql.Identifier(column),
                sql.Literal(lists_val),
            )

        create_space_index_sql = sql.SQL("""
            CREATE INDEX IF NOT EXISTS {} 
            ON {} (genie_space_id)
        """).format(
            sql.Identifier(f"{self.table_name}_space_idx"),
            table_id,
        )
        create_unique_question_index_sql = sql.SQL("""
            CREATE UNIQUE INDEX IF NOT EXISTS {}
            ON {} (genie_space_id, question)
        """).format(
            sql.Identifier(f"{self.table_name}_unique_question_idx"),
            table_id,
        )

        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                cur.execute(create_extension_sql)

                # Check if table exists and verify embedding dimensions
                try:
                    cur.execute(check_dims_sql, (self.table_name,))
                    row: DbRow | None = cur.fetchone()
                    if row is not None:
                        current_dims = row.get("atttypmod", 0)
                        if current_dims != self.embedding_dims:
                            logger.warning(
                                "Embedding dimension mismatch, dropping and recreating table",
                                layer=self.name,
                                table_dims=current_dims,
                                expected_dims=self.embedding_dims,
                                table_name=self.table_name,
                            )
                            cur.execute(sql.SQL("DROP TABLE {}").format(table_id))
                except Exception:
                    pass

                try:
                    cur.execute(create_table_sql)
                except Exception as e:
                    logger.debug(
                        f"Table creation skipped (may already exist): {e}",
                        layer=self.name,
                    )

                # Migration: Add message_id column if it doesn't exist (for existing tables)
                try:
                    cur.execute(add_message_id_sql)
                    logger.debug(
                        "Added message_id column (or already exists)",
                        layer=self.name,
                        table_name=self.table_name,
                    )
                except Exception as e:
                    # Column might already exist or other error
                    logger.debug(
                        f"message_id column migration skipped: {e}",
                        layer=self.name,
                    )

                # Auto-compute ivfflat lists if not explicitly configured.
                if ivfflat_lists is None:
                    try:
                        cur.execute(
                            sql.SQL("SELECT COUNT(*) as cnt FROM {}").format(table_id)
                        )
                        count_row = cur.fetchone()
                        row_count = count_row.get("cnt", 0) if count_row else 0
                        ivfflat_lists = max(100, int(math.sqrt(row_count)))
                    except Exception:
                        ivfflat_lists = 100

                # Store resolved lists for probes auto-computation in _find_similar
                self._resolved_ivfflat_lists = ivfflat_lists

                logger.debug(
                    "Using ivfflat lists",
                    layer=self.name,
                    ivfflat_lists=ivfflat_lists,
                )

                create_question_embedding_index_sql = _make_ivfflat_index_sql(
                    f"{self.table_name}_question_embedding_idx",
                    "question_embedding",
                    ivfflat_lists,
                )
                create_context_embedding_index_sql = _make_ivfflat_index_sql(
                    f"{self.table_name}_context_embedding_idx",
                    "context_embedding",
                    ivfflat_lists,
                )

                # Create indexes
                for idx_name, idx_sql in [
                    (f"{self.table_name}_space_idx", create_space_index_sql),
                    (
                        f"{self.table_name}_question_embedding_idx",
                        create_question_embedding_index_sql,
                    ),
                    (
                        f"{self.table_name}_context_embedding_idx",
                        create_context_embedding_index_sql,
                    ),
                    (
                        f"{self.table_name}_unique_question_idx",
                        create_unique_question_index_sql,
                    ),
                ]:
                    if self._index_exists(cur, idx_name):
                        logger.debug(
                            f"Index {idx_name} already exists", layer=self.name
                        )
                        continue
                    try:
                        cur.execute(idx_sql)
                    except Exception as e:
                        logger.warning(
                            f"Could not create {idx_name}: {e}", layer=self.name
                        )

                # Create prompt history table
                try:
                    self._create_prompt_history_table(cur)
                except Exception as e:
                    logger.error(
                        f"Failed to create prompt history table: {e}",
                        layer=self.name,
                        exc_info=True,
                    )

    def _create_prompt_history_table(self, cur: Any) -> None:
        """Create the prompt history table for tracking user prompts."""
        prompt_table_name = self.prompt_history_table
        prompt_table_id = sql.Identifier(prompt_table_name)

        create_prompt_table_sql = sql.SQL("""
            CREATE TABLE IF NOT EXISTS {} (
                id SERIAL PRIMARY KEY,
                genie_space_id TEXT NOT NULL,
                conversation_id TEXT NOT NULL,
                prompt TEXT NOT NULL,
                cache_hit BOOLEAN DEFAULT FALSE,
                cache_entry_id INTEGER,
                created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
            )
        """).format(prompt_table_id)

        # Migration: Add cache_entry_id column if it doesn't exist
        add_cache_entry_id_sql = sql.SQL("""
            ALTER TABLE {} 
            ADD COLUMN IF NOT EXISTS cache_entry_id INTEGER
        """).format(prompt_table_id)

        create_conversation_index_sql = sql.SQL("""
            CREATE INDEX IF NOT EXISTS {} 
            ON {} (genie_space_id, conversation_id, created_at DESC)
        """).format(
            sql.Identifier(f"{prompt_table_name}_conversation_idx"),
            prompt_table_id,
        )
        create_space_index_sql = sql.SQL("""
            CREATE INDEX IF NOT EXISTS {} 
            ON {} (genie_space_id, created_at DESC)
        """).format(
            sql.Identifier(f"{prompt_table_name}_space_idx"),
            prompt_table_id,
        )
        create_unique_prompt_index_sql = sql.SQL("""
            CREATE UNIQUE INDEX IF NOT EXISTS {}
            ON {} (genie_space_id, conversation_id, prompt)
        """).format(
            sql.Identifier(f"{prompt_table_name}_unique_prompt_idx"),
            prompt_table_id,
        )
        create_cache_entry_index_sql = sql.SQL("""
            CREATE INDEX IF NOT EXISTS {}
            ON {} (cache_entry_id)
            WHERE cache_entry_id IS NOT NULL
        """).format(
            sql.Identifier(f"{prompt_table_name}_cache_entry_idx"),
            prompt_table_id,
        )

        try:
            cur.execute(create_prompt_table_sql)
        except Exception as e:
            if "duplicate key" in str(e) or "already exists" in str(e):
                logger.debug("Prompt history table already exists", layer=self.name)
            else:
                raise

        # Migration: Add cache_entry_id column if it doesn't exist (for existing tables)
        try:
            cur.execute(add_cache_entry_id_sql)
            logger.debug(
                "Added cache_entry_id column (or already exists)",
                layer=self.name,
                table_name=prompt_table_name,
            )
        except Exception as e:
            logger.debug(
                f"cache_entry_id column migration skipped: {e}",
                layer=self.name,
            )

        for idx_name, idx_sql in [
            (f"{prompt_table_name}_conversation_idx", create_conversation_index_sql),
            (f"{prompt_table_name}_space_idx", create_space_index_sql),
            (f"{prompt_table_name}_unique_prompt_idx", create_unique_prompt_index_sql),
            (f"{prompt_table_name}_cache_entry_idx", create_cache_entry_index_sql),
        ]:
            if self._index_exists(cur, idx_name):
                continue
            try:
                cur.execute(idx_sql)
            except Exception as e:
                if "duplicate key" not in str(e) and "already exists" not in str(e):
                    logger.warning(
                        f"Could not create index {idx_name}: {e}", layer=self.name
                    )

        logger.info(
            "Prompt history table ready", layer=self.name, table=prompt_table_name
        )

    @mlflow.trace(name="semantic_search_postgres")
    def _find_similar(
        self,
        question: str,
        conversation_context: str,
        question_embedding: list[float],
        context_embedding: list[float],
        conversation_id: str | None = None,
    ) -> tuple[SQLCacheEntry, float] | None:
        """Find a semantically similar cached entry using pg_vector.

        Uses a two-phase approach for scalability:
        Phase 1 (SQL): Use pg_vector's native ``ORDER BY embedding <=> query LIMIT K``
            pattern which is optimized to use the ivfflat index. This retrieves the
            top-K candidates by question distance. ``SET LOCAL ivfflat.probes`` controls
            recall quality.
        Phase 2 (Python): Compute combined similarity on the K candidates, apply
            threshold checks (question, context, TTL), and pick the best match.

        This avoids the previous CTE approach which materialized all filtered rows
        before ranking, defeating the ivfflat index optimization.
        """
        ttl_seconds = self.time_to_live_seconds
        ttl_disabled = ttl_seconds is None or ttl_seconds < 0

        question_weight = self.question_weight
        context_weight = self.context_weight
        candidates_k = self.parameters.ivfflat_candidates

        # Resolve ivfflat probes: explicit config or auto-compute as sqrt(lists).
        # sqrt(lists) is the standard recommendation for good recall at any scale.
        configured_probes = self.parameters.ivfflat_probes
        if configured_probes is not None:
            ivfflat_probes = configured_probes
        else:
            ivfflat_probes = max(10, int(math.sqrt(self._resolved_ivfflat_lists)))

        # Phase 1: Retrieve top-K candidates using pg_vector index scan.
        # The ORDER BY ... LIMIT K pattern is specifically optimized by pg_vector
        # to use the ivfflat index for approximate nearest neighbor search,
        # unlike WHERE-based distance filters or CTEs which cause full scans.
        search_sql = sql.SQL("""
            SELECT id, question, conversation_context, sql_query,
                   description, conversation_id, message_id, created_at,
                   (question_embedding <=> %s::vector) as question_distance,
                   COALESCE(NULLIF((context_embedding <=> %s::vector), 'NaN'::float8), 0.0) as context_distance
            FROM {}
            WHERE genie_space_id = %s
            ORDER BY question_embedding <=> %s::vector
            LIMIT %s
        """).format(sql.Identifier(self.table_name))

        question_emb_str = _serialize_embedding(question_embedding)
        context_emb_str = _serialize_embedding(context_embedding)

        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                # Set ivfflat probes for this transaction only.
                # SET LOCAL resets when the transaction ends, so it is
                # safe for connection pools.
                cur.execute(
                    sql.SQL("SET LOCAL ivfflat.probes = {}").format(
                        sql.Literal(ivfflat_probes)
                    )
                )

                cur.execute(
                    search_sql,
                    (
                        question_emb_str,
                        context_emb_str,
                        self.space_id,
                        question_emb_str,
                        candidates_k,
                    ),
                )
                rows: list[DbRow] = cur.fetchall()

                if not rows:
                    logger.info(
                        "Cache MISS (no entries)",
                        layer=self.name,
                        question=question[:50],
                        space=self.space_id,
                    )
                    return None

                # Phase 2: Python-side reranking of K candidates.
                # Convert distances to similarities and apply threshold checks.
                query_has_context = bool(
                    conversation_context and conversation_context.strip()
                )

                best_row: DbRow | None = None
                best_question_sim: float = 0.0
                best_context_sim: float = 0.0
                best_combined_sim: float = 0.0

                for row in rows:
                    question_similarity = 1.0 - row["question_distance"]
                    context_similarity = 1.0 - row["context_distance"]

                    # Question threshold gate
                    if question_similarity < self.similarity_threshold:
                        continue

                    # Bidirectional context-skip: skip context threshold when
                    # EITHER side has no context.
                    cached_context = row.get("conversation_context", "")
                    has_stored_context = bool(cached_context and cached_context.strip())

                    if has_stored_context and query_has_context:
                        if context_similarity < self.context_similarity_threshold:
                            continue

                    combined = (
                        question_weight * question_similarity
                        + context_weight * context_similarity
                    )

                    if combined > best_combined_sim:
                        best_row = row
                        best_question_sim = question_similarity
                        best_context_sim = context_similarity
                        best_combined_sim = combined

                if best_row is None:
                    logger.info(
                        "Cache MISS (no candidates pass thresholds)",
                        layer=self.name,
                        question=question[:50],
                        space=self.space_id,
                        candidates_checked=len(rows),
                    )
                    return None

                # TTL check on the best candidate
                entry_id = best_row.get("id")
                created_at = best_row["created_at"]

                if not ttl_disabled and created_at:
                    age_seconds = (
                        datetime.now(created_at.tzinfo) - created_at
                    ).total_seconds()
                    if age_seconds > ttl_seconds:
                        cur.execute(
                            sql.SQL("DELETE FROM {} WHERE id = %s").format(
                                sql.Identifier(self.table_name)
                            ),
                            (entry_id,),
                        )
                        logger.info("Cache MISS (expired, deleted)", layer=self.name)
                        return None

                cached_question = best_row.get("question", "")
                cache_age_seconds = None
                if created_at:
                    cache_age_seconds = (
                        datetime.now(created_at.tzinfo) - created_at
                    ).total_seconds()

                logger.info(
                    "Cache HIT",
                    layer=self.name,
                    question=question[:80],
                    matched_question=cached_question[:80],
                    cache_age_seconds=round(cache_age_seconds, 1)
                    if cache_age_seconds
                    else None,
                    question_similarity=f"{best_question_sim:.4f}",
                    context_similarity=f"{best_context_sim:.4f}",
                    combined_similarity=f"{best_combined_sim:.4f}",
                    candidates_checked=len(rows),
                )

                entry = SQLCacheEntry(
                    query=best_row["sql_query"],
                    description=best_row.get("description", ""),
                    conversation_id=best_row.get("conversation_id", ""),
                    created_at=created_at,
                    message_id=best_row.get("message_id"),
                    cache_entry_id=entry_id,
                )
                return entry, best_combined_sim

    def _store_entry(
        self,
        question: str,
        conversation_context: str,
        question_embedding: list[float],
        context_embedding: list[float],
        response: GenieResponse,
        message_id: str | None = None,
    ) -> None:
        """Store a new cache entry with dual embeddings and message_id."""
        insert_sql = sql.SQL("""
            INSERT INTO {} 
            (genie_space_id, question, conversation_context, context_string, 
             question_embedding, context_embedding, sql_query, description, 
             conversation_id, message_id)
            VALUES (%s, %s, %s, %s, %s::vector, %s::vector, %s, %s, %s, %s)
            ON CONFLICT (genie_space_id, question)
            DO UPDATE SET
                conversation_context = EXCLUDED.conversation_context,
                context_string = EXCLUDED.context_string,
                question_embedding = EXCLUDED.question_embedding,
                context_embedding = EXCLUDED.context_embedding,
                sql_query = EXCLUDED.sql_query,
                description = EXCLUDED.description,
                conversation_id = EXCLUDED.conversation_id,
                message_id = EXCLUDED.message_id,
                created_at = CURRENT_TIMESTAMP
        """).format(sql.Identifier(self.table_name))
        question_emb_str = _serialize_embedding(question_embedding)
        context_emb_str = _serialize_embedding(context_embedding)

        if conversation_context:
            full_context_string = f"{conversation_context}\nCurrent: {question}"
        else:
            full_context_string = question

        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    insert_sql,
                    (
                        self.space_id,
                        question,
                        conversation_context,
                        full_context_string,
                        question_emb_str,
                        context_emb_str,
                        response.query,
                        response.description,
                        response.conversation_id,
                        message_id,
                    ),
                )
                logger.debug(
                    "Stored cache entry",
                    layer=self.name,
                    question=question[:50],
                    space=self.space_id,
                    message_id=message_id,
                )

    def _on_stale_cache_entry(self, question: str) -> None:
        """Delete stale cache entry from database."""
        delete_sql = sql.SQL(
            "DELETE FROM {} WHERE genie_space_id = %s AND question = %s"
        ).format(sql.Identifier(self.table_name))
        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                cur.execute(delete_sql, (self.space_id, question))
                deleted_rows = cur.rowcount
                logger.info(
                    "Deleted stale cache entry",
                    layer=self.name,
                    deleted_rows=deleted_rows,
                    space_id=self.space_id,
                )

    def _invalidate_by_question(self, question: str) -> bool:
        """
        Invalidate cache entries matching a specific question.

        This method is called when negative feedback is received to remove
        the corresponding cache entry from the PostgreSQL database.

        Args:
            question: The question text to match and invalidate

        Returns:
            True if an entry was found and invalidated, False otherwise
        """
        delete_sql = sql.SQL(
            "DELETE FROM {} WHERE genie_space_id = %s AND question = %s"
        ).format(sql.Identifier(self.table_name))
        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                cur.execute(delete_sql, (self.space_id, question))
                deleted_rows = cur.rowcount if isinstance(cur.rowcount, int) else 0
                if deleted_rows > 0:
                    logger.info(
                        "Invalidated cache entry by question",
                        layer=self.name,
                        question=question[:50],
                        deleted_rows=deleted_rows,
                        space_id=self.space_id,
                    )
                    return True
                return False

    # Template Method hook implementations

    def _before_cache_lookup(self, question: str, conversation_id: str | None) -> None:
        """Store prompt before cache lookup."""
        if conversation_id:
            self._store_user_prompt(
                prompt=question,
                conversation_id=conversation_id,
                cache_hit=False,
            )
            # Track that we stored the prompt
            self._prompt_stored_for_current_request = True
        else:
            self._prompt_stored_for_current_request = False

    def _after_cache_hit(
        self,
        question: str,
        conversation_id: str | None,
        result: CacheResult,
    ) -> None:
        """Update cache_hit flag and cache_entry_id after a cache hit."""
        if result.cache_hit and self._prompt_stored_for_current_request:
            actual_conv_id = result.response.conversation_id or conversation_id
            if actual_conv_id:
                self._update_prompt_cache_hit(
                    conversation_id=actual_conv_id,
                    prompt=question,
                    cache_hit=True,
                    cache_entry_id=result.cache_entry_id,
                )

    def _after_cache_miss(
        self,
        question: str,
        conversation_id: str | None,
        result: CacheResult,
    ) -> None:
        """Store prompt if not done earlier (when conversation_id comes from response)."""
        if (
            not self._prompt_stored_for_current_request
            and result.response.conversation_id
        ):
            self._store_user_prompt(
                prompt=question,
                conversation_id=result.response.conversation_id,
                cache_hit=False,
            )

    # Template Method implementations for invalidate_expired() and clear()

    def _get_empty_expiration_result(self) -> dict[str, int]:
        """Return empty dict for PostgresContextAwareGenieService."""
        return {"cache": 0, "prompt_history": 0}

    def _delete_expired_entries(self, ttl_seconds: int) -> dict[str, int]:
        """Delete expired entries from cache and prompt history."""
        prompt_ttl_seconds = self.parameters.prompt_history_ttl_seconds
        if prompt_ttl_seconds is None:
            prompt_ttl_seconds = ttl_seconds

        result: dict[str, int] = {"cache": 0, "prompt_history": 0}

        # Delete expired cache entries
        delete_cache_sql = sql.SQL("""
            DELETE FROM {}
            WHERE genie_space_id = %s
              AND created_at < NOW() - INTERVAL '%s seconds'
        """).format(sql.Identifier(self.table_name))

        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                cur.execute(delete_cache_sql, (self.space_id, ttl_seconds))
                deleted = cur.rowcount if isinstance(cur.rowcount, int) else 0
                result["cache"] = deleted
                logger.debug(
                    "Deleted expired cache entries",
                    layer=self.name,
                    deleted_count=deleted,
                )

        # Delete expired prompt history
        if prompt_ttl_seconds is not None and prompt_ttl_seconds >= 0:
            try:
                delete_prompt_sql = sql.SQL("""
                    DELETE FROM {}
                    WHERE genie_space_id = %s
                      AND created_at < NOW() - INTERVAL '%s seconds'
                """).format(sql.Identifier(self.prompt_history_table))
                with self._pool.connection() as conn:
                    with conn.cursor() as cur:
                        cur.execute(
                            delete_prompt_sql, (self.space_id, prompt_ttl_seconds)
                        )
                        deleted = cur.rowcount if isinstance(cur.rowcount, int) else 0
                        result["prompt_history"] = deleted
            except Exception as e:
                logger.warning(
                    f"Failed to clean up prompt history: {e}", layer=self.name
                )

        return result

    def _delete_all_entries(self) -> int:
        """Delete all cache entries for this Genie space."""
        delete_sql = sql.SQL("DELETE FROM {} WHERE genie_space_id = %s").format(
            sql.Identifier(self.table_name)
        )

        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                cur.execute(delete_sql, (self.space_id,))
                deleted: int = cur.rowcount
                logger.debug(
                    "Cleared cache entries", layer=self.name, deleted_count=deleted
                )
                return deleted

    # Template Method implementations for stats()

    def _count_all_entries(self) -> int:
        """Count all cache entries for this Genie space."""
        count_sql = sql.SQL(
            "SELECT COUNT(*) as total FROM {} WHERE genie_space_id = %s"
        ).format(sql.Identifier(self.table_name))
        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                cur.execute(count_sql, (self.space_id,))
                row = cur.fetchone()
                return row.get("total", 0) if row else 0

    def _count_entries_with_ttl(self, ttl_seconds: int) -> tuple[int, int]:
        """Count total and expired entries for this Genie space."""
        stats_sql = sql.SQL("""
            SELECT
                COUNT(*) as total,
                COUNT(*) FILTER (WHERE created_at <= NOW() - INTERVAL '%s seconds') as expired
            FROM {}
            WHERE genie_space_id = %s
        """).format(sql.Identifier(self.table_name))
        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                cur.execute(stats_sql, (ttl_seconds, self.space_id))
                row = cur.fetchone()
                if row:
                    return row.get("total", 0), row.get("expired", 0)
                return 0, 0

    def _get_additional_stats(self) -> dict[str, Any]:
        """Add prompt history stats."""
        prompt_stats_sql = sql.SQL("""
            SELECT
                COUNT(*) as total_prompts,
                COUNT(*) FILTER (WHERE cache_hit = true) as cache_hit_prompts,
                COUNT(*) FILTER (WHERE cache_hit = false) as cache_miss_prompts,
                COUNT(DISTINCT conversation_id) as total_conversations
            FROM {}
            WHERE genie_space_id = %s
        """).format(sql.Identifier(self.prompt_history_table))
        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                cur.execute(prompt_stats_sql, (self.space_id,))
                row = cur.fetchone()
                if row:
                    total_prompts = row.get("total_prompts", 0)
                    return {
                        "prompt_history": {
                            "total_prompts": total_prompts,
                            "cache_hit_prompts": row.get("cache_hit_prompts", 0),
                            "cache_miss_prompts": row.get("cache_miss_prompts", 0),
                            "total_conversations": row.get("total_conversations", 0),
                            "cache_hit_rate": (
                                row.get("cache_hit_prompts", 0) / total_prompts
                                if total_prompts > 0
                                else 0.0
                            ),
                        }
                    }
        return {}

    def get_entries(
        self,
        limit: int | None = None,
        offset: int | None = None,
        include_embeddings: bool = False,
        conversation_id: str | None = None,
        created_after: datetime | None = None,
        created_before: datetime | None = None,
        question_contains: str | None = None,
    ) -> list[dict[str, Any]]:
        """
        Get cache entries with optional filtering.

        This method retrieves cache entries for inspection, debugging, or
        generating evaluation datasets for threshold optimization.

        Args:
            limit: Maximum number of entries to return (None = no limit)
            offset: Number of entries to skip for pagination (None = 0)
            include_embeddings: Whether to include embedding vectors in results.
                Embeddings are large, so set False for general inspection.
            conversation_id: Filter by conversation ID (None = all conversations)
            created_after: Only entries created after this time (None = no filter)
            created_before: Only entries created before this time (None = no filter)
            question_contains: Case-insensitive text search on question field

        Returns:
            List of cache entry dicts. See base class for full key documentation.

        Example:
            # Get entries with embeddings for evaluation dataset generation
            entries = cache.get_entries(include_embeddings=True, limit=100)
            eval_dataset = generate_eval_dataset_from_cache(entries)
        """
        self._setup()

        # Build column list
        base_columns = [
            "id",
            "question",
            "conversation_context",
            "sql_query",
            "description",
            "conversation_id",
            "created_at",
        ]

        if include_embeddings:
            columns = base_columns + ["question_embedding", "context_embedding"]
        else:
            columns = base_columns

        columns_str = ", ".join(columns)

        # Build WHERE clause with parameters
        where_clauses = ["genie_space_id = %s"]
        params: list[Any] = [self.space_id]

        if conversation_id is not None:
            where_clauses.append("conversation_id = %s")
            params.append(conversation_id)

        if created_after is not None:
            where_clauses.append("created_at > %s")
            params.append(created_after)

        if created_before is not None:
            where_clauses.append("created_at < %s")
            params.append(created_before)

        if question_contains is not None:
            where_clauses.append("question ILIKE %s")
            params.append(f"%{question_contains}%")

        where_str = " AND ".join(where_clauses)

        # Build full query using safe SQL composition
        query = sql.SQL("""
            SELECT {}
            FROM {}
            WHERE {}
            ORDER BY created_at DESC
        """).format(
            sql.SQL(columns_str),
            sql.Identifier(self.table_name),
            sql.SQL(where_str),
        )

        if limit is not None:
            query += sql.SQL(" LIMIT {}").format(sql.Literal(int(limit)))

        if offset is not None:
            query += sql.SQL(" OFFSET {}").format(sql.Literal(int(offset)))

        # Execute query
        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                cur.execute(query, params)
                rows = cur.fetchall()

                entries: list[dict[str, Any]] = []
                for row in rows:
                    entry: dict[str, Any] = {
                        "id": row.get("id"),
                        "question": row.get("question"),
                        "conversation_context": row.get("conversation_context"),
                        "sql_query": row.get("sql_query"),
                        "description": row.get("description"),
                        "conversation_id": row.get("conversation_id"),
                        "created_at": row.get("created_at"),
                    }

                    if include_embeddings:
                        # Convert pgvector to list
                        q_emb = row.get("question_embedding")
                        c_emb = row.get("context_embedding")
                        entry["question_embedding"] = (
                            list(q_emb) if q_emb is not None else []
                        )
                        entry["context_embedding"] = (
                            list(c_emb) if c_emb is not None else []
                        )

                    entries.append(entry)

                logger.debug(
                    "Retrieved cache entries",
                    layer=self.name,
                    count=len(entries),
                    include_embeddings=include_embeddings,
                    filters={
                        "conversation_id": conversation_id,
                        "created_after": str(created_after) if created_after else None,
                        "created_before": (
                            str(created_before) if created_before else None
                        ),
                        "question_contains": question_contains,
                    },
                )

                return entries

    def from_space(
        self,
        space_id: str | None = None,
        *,
        include_all_messages: bool = True,
        from_datetime: datetime | None = None,
        to_datetime: datetime | None = None,
        max_messages: int | None = None,
        max_conversations: int | None = None,
    ) -> Self:
        """Populate cache from existing Genie space conversations.

        Fetches all conversations from a Genie space and populates:
        1. Prompt history table - all user messages
        2. Cache embeddings table - messages with SQL query attachments

        Uses ON CONFLICT DO NOTHING to avoid duplicate entries.

        Args:
            space_id: Genie space ID to import from (defaults to self.space_id)
            include_all_messages: If True, fetch all users' conversations
            from_datetime: Only include messages after this time
            to_datetime: Only include messages before this time
            max_messages: Limit to last N messages (most recent first)
            max_conversations: Limit to N conversations (stops pagination after reaching limit)

        Returns:
            self for method chaining
        """
        if self.workspace_client is None:
            raise ValueError("workspace_client is required for from_space()")

        self._setup()
        target_space_id = space_id or self.space_id

        logger.info(
            "Starting from_space import",
            layer=self.name,
            space_id=target_space_id,
            include_all_messages=include_all_messages,
        )

        stats = {
            "conversations_processed": 0,
            "prompts_imported": 0,
            "prompts_skipped": 0,
            "cache_entries_imported": 0,
            "cache_entries_skipped": 0,
            "errors": 0,
        }

        from databricks.sdk.service.dashboards import GenieMessage

        all_messages: list[tuple[str, GenieMessage]] = []
        page_token: str | None = None

        while True:
            try:
                response = self.workspace_client.genie.list_conversations(
                    space_id=target_space_id,
                    include_all=include_all_messages,
                    page_token=page_token,
                )
            except Exception as e:
                logger.error(f"Failed to list conversations: {e}", layer=self.name)
                stats["errors"] += 1
                break

            if response.conversations is None:
                logger.debug(
                    "No conversations in response",
                    layer=self.name,
                    space_id=target_space_id,
                )
                break

            logger.debug(
                "Fetched conversations page",
                layer=self.name,
                conversations_in_page=len(response.conversations),
                total_conversations_so_far=stats["conversations_processed"],
                has_next_page=response.next_page_token is not None,
            )

            for conversation in response.conversations:
                if conversation.conversation_id is None:
                    continue

                stats["conversations_processed"] += 1

                try:
                    messages_response = (
                        self.workspace_client.genie.list_conversation_messages(
                            space_id=target_space_id,
                            conversation_id=conversation.conversation_id,
                        )
                    )
                except Exception as e:
                    logger.warning(f"Failed to fetch messages: {e}", layer=self.name)
                    stats["errors"] += 1
                    continue

                if messages_response.messages is None:
                    continue

                for message in messages_response.messages:
                    all_messages.append((conversation.conversation_id, message))

                if max_messages and len(all_messages) >= max_messages:
                    break

                if (
                    max_conversations
                    and stats["conversations_processed"] >= max_conversations
                ):
                    break

            if max_messages and len(all_messages) >= max_messages:
                break

            if (
                max_conversations
                and stats["conversations_processed"] >= max_conversations
            ):
                logger.debug(
                    "Reached max_conversations limit",
                    layer=self.name,
                    max_conversations=max_conversations,
                    total_conversations=stats["conversations_processed"],
                )
                break

            page_token = response.next_page_token
            if page_token is None:
                logger.debug(
                    "No more pages to fetch",
                    layer=self.name,
                    total_conversations=stats["conversations_processed"],
                    total_messages=len(all_messages),
                )
                break

        # Sort and limit
        all_messages.sort(
            key=lambda x: x[1].created_timestamp if x[1].created_timestamp else 0,
            reverse=True,
        )
        if max_messages:
            all_messages = all_messages[:max_messages]

        # Group messages by conversation_id for context building
        from collections import defaultdict

        messages_by_conversation: dict[str, list[tuple[str, GenieMessage]]] = (
            defaultdict(list)
        )
        for conv_id, msg in all_messages:
            messages_by_conversation[conv_id].append((conv_id, msg))

        # Sort each conversation's messages by timestamp (oldest first for context building)
        for conv_id in messages_by_conversation:
            messages_by_conversation[conv_id].sort(
                key=lambda x: x[1].created_timestamp if x[1].created_timestamp else 0
            )

        # Process messages
        for conversation_id, message in all_messages:
            if message.content is None:
                continue

            message_created_at = None
            if message.created_timestamp:
                # Inherit timezone from the filter parameters so comparisons
                # don't mix offset-naive and offset-aware datetimes.
                tz = (
                    from_datetime.tzinfo
                    if from_datetime
                    else (to_datetime.tzinfo if to_datetime else None)
                )
                message_created_at = datetime.fromtimestamp(
                    message.created_timestamp / 1000.0,
                    tz=tz,
                )

            if message_created_at:
                if from_datetime and message_created_at < from_datetime:
                    continue
                if to_datetime and message_created_at > to_datetime:
                    continue

            # Store prompt
            prompt_stored = self._store_prompt_if_not_exists(
                prompt=message.content,
                conversation_id=conversation_id,
                space_id=target_space_id,
                cache_hit=False,
                created_at=message_created_at,
            )

            if prompt_stored:
                stats["prompts_imported"] += 1
            else:
                stats["prompts_skipped"] += 1

            # Check for SQL attachments
            if message.attachments:
                for attachment in message.attachments:
                    if attachment.query and attachment.query.query:
                        try:
                            # Build conversation context from prior messages
                            # Uses same "Previous: {content}" format as normal operations
                            prior_messages: list[str] = []
                            conv_messages = messages_by_conversation.get(
                                conversation_id, []
                            )
                            for _, prior_msg in conv_messages:
                                if (
                                    prior_msg.created_timestamp
                                    and message.created_timestamp
                                ):
                                    if (
                                        prior_msg.created_timestamp
                                        < message.created_timestamp
                                    ):
                                        if prior_msg.content:
                                            content = prior_msg.content
                                            if len(content) > 500:
                                                content = content[:500] + "..."
                                            prior_messages.append(
                                                f"Previous: {content}"
                                            )

                            # Limit to context_window_size (most recent N messages)
                            context_window = self.context_window_size
                            if len(prior_messages) > context_window:
                                prior_messages = prior_messages[-context_window:]

                            conversation_context = "\n".join(prior_messages)

                            # Generate embeddings
                            question_embedding = self._embeddings.embed_query(
                                message.content
                            )
                            if conversation_context:
                                context_embedding = self._embeddings.embed_query(
                                    conversation_context
                                )
                            else:
                                # Zero vector when no prior context (first message)
                                context_embedding = [0.0] * len(question_embedding)

                            cache_stored = self._store_cache_entry_if_not_exists(
                                question=message.content,
                                conversation_context=conversation_context,
                                question_embedding=question_embedding,
                                context_embedding=context_embedding,
                                sql_query=attachment.query.query,
                                description=attachment.query.description or "",
                                conversation_id=conversation_id,
                                space_id=target_space_id,
                            )

                            if cache_stored:
                                stats["cache_entries_imported"] += 1
                            else:
                                stats["cache_entries_skipped"] += 1
                        except Exception as e:
                            logger.warning(
                                f"Failed to generate embeddings: {e}", layer=self.name
                            )
                            stats["errors"] += 1

        logger.info("Completed from_space import", layer=self.name, **stats)
        return self

    def _store_prompt_if_not_exists(
        self,
        prompt: str,
        conversation_id: str,
        space_id: str | None = None,
        cache_hit: bool = False,
        created_at: datetime | None = None,
    ) -> bool:
        """Store prompt with ON CONFLICT DO NOTHING."""
        target_space_id = space_id or self.space_id
        prompt_table_id = sql.Identifier(self.prompt_history_table)

        if created_at:
            insert_sql = sql.SQL("""
                INSERT INTO {} 
                (genie_space_id, conversation_id, prompt, cache_hit, created_at)
                VALUES (%s, %s, %s, %s, %s)
                ON CONFLICT (genie_space_id, conversation_id, prompt) DO NOTHING
            """).format(prompt_table_id)
            params = (target_space_id, conversation_id, prompt, cache_hit, created_at)
        else:
            insert_sql = sql.SQL("""
                INSERT INTO {} 
                (genie_space_id, conversation_id, prompt, cache_hit)
                VALUES (%s, %s, %s, %s)
                ON CONFLICT (genie_space_id, conversation_id, prompt) DO NOTHING
            """).format(prompt_table_id)
            params = (target_space_id, conversation_id, prompt, cache_hit)

        try:
            with self._pool.connection() as conn:
                with conn.cursor() as cur:
                    cur.execute(insert_sql, params)
                    return cur.rowcount > 0 if isinstance(cur.rowcount, int) else False
        except Exception:
            return False

    def _store_cache_entry_if_not_exists(
        self,
        question: str,
        conversation_context: str,
        question_embedding: list[float],
        context_embedding: list[float],
        sql_query: str,
        description: str | None = None,
        conversation_id: str | None = None,
        space_id: str | None = None,
    ) -> bool:
        """Store cache entry with ON CONFLICT DO NOTHING."""
        target_space_id = space_id or self.space_id

        insert_sql = sql.SQL("""
            INSERT INTO {} 
            (genie_space_id, question, conversation_context, context_string, 
             question_embedding, context_embedding, sql_query, description, conversation_id)
            VALUES (%s, %s, %s, %s, %s::vector, %s::vector, %s, %s, %s)
            ON CONFLICT (genie_space_id, question) DO NOTHING
        """).format(sql.Identifier(self.table_name))
        question_emb_str = _serialize_embedding(question_embedding)
        context_emb_str = _serialize_embedding(context_embedding)
        full_context = (
            f"{conversation_context}\nCurrent: {question}"
            if conversation_context
            else question
        )

        try:
            with self._pool.connection() as conn:
                with conn.cursor() as cur:
                    cur.execute(
                        insert_sql,
                        (
                            target_space_id,
                            question,
                            conversation_context,
                            full_context,
                            question_emb_str,
                            context_emb_str,
                            sql_query,
                            description or "",
                            conversation_id or "",
                        ),
                    )
                    return cur.rowcount > 0 if isinstance(cur.rowcount, int) else False
        except Exception:
            return False
