"""Semantic LLM cache backed by SAP HANA Cloud.

This module implements the LangChain BaseCache interface using HANA Cloud's
REAL_VECTOR type and COSINE_SIMILARITY function for semantic matching,
enabling cache hits for semantically similar (not just identical) prompts.

Typical usage:

    from langchain_hana_cache import HANASemanticLLMCache
    from langchain_openai import OpenAIEmbeddings
    from langchain_core.globals import set_llm_cache

    cache = HANASemanticLLMCache(
        connection=hana_conn,
        embedding=OpenAIEmbeddings(model="text-embedding-3-small"),
    )
    set_llm_cache(cache)
"""

from __future__ import annotations

import logging
from typing import Any, Optional, Sequence

from langchain_core.caches import BaseCache
from langchain_core.embeddings import Embeddings
from langchain_core.outputs import Generation

from langchain_hana_cache.utils import (
    deserialize_generations,
    hash_prompt,
    serialize_generations,
)

logger = logging.getLogger(__name__)

RETURN_VAL_TYPE = Sequence[Generation]

# ──────────────────────────────────────────────────────────────────────
# DDL — table definition for HANA Cloud
# ──────────────────────────────────────────────────────────────────────

_DEFAULT_TABLE = "LLM_CACHE"

_CREATE_TABLE_TEMPLATE = """
CREATE TABLE "{table}" (
    "PROMPT_HASH"       NVARCHAR(64) PRIMARY KEY,
    "PROMPT_TEXT"        NCLOB,
    "PROMPT_EMBEDDING"   REAL_VECTOR({dim}),
    "LLM_STRING_HASH"   NVARCHAR(64),
    "LLM_STRING"         NCLOB,
    "RESPONSE"           NCLOB,
    "CREATED_AT"         TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    "LAST_ACCESSED"      TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    "HIT_COUNT"          INTEGER DEFAULT 0
)
"""

# ──────────────────────────────────────────────────────────────────────
# SQL templates
# ──────────────────────────────────────────────────────────────────────

_SELECT_SIMILAR = """
    SELECT "PROMPT_TEXT", "RESPONSE", "LLM_STRING",
           COSINE_SIMILARITY("PROMPT_EMBEDDING", TO_REAL_VECTOR(?)) AS "SIMILARITY"
    FROM "{table}"
    WHERE "LLM_STRING_HASH" = ?
      AND COSINE_SIMILARITY("PROMPT_EMBEDDING", TO_REAL_VECTOR(?)) > ?
"""

_SELECT_SIMILAR_TTL = """
      AND SECONDS_BETWEEN("CREATED_AT", CURRENT_TIMESTAMP) < ?
"""

_SELECT_SIMILAR_ORDER = """
    ORDER BY "SIMILARITY" DESC
    LIMIT 1
"""

_UPDATE_ACCESS = """
    UPDATE "{table}"
    SET "LAST_ACCESSED" = CURRENT_TIMESTAMP, "HIT_COUNT" = "HIT_COUNT" + 1
    WHERE "PROMPT_HASH" = ?
"""

_UPSERT_ENTRY = """
    UPSERT "{table}"
    ("PROMPT_HASH", "PROMPT_TEXT", "PROMPT_EMBEDDING", "LLM_STRING_HASH",
     "LLM_STRING", "RESPONSE", "CREATED_AT", "LAST_ACCESSED", "HIT_COUNT")
    VALUES (?, ?, TO_REAL_VECTOR(?), ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, 0)
    WITH PRIMARY KEY
"""

_DELETE_ALL = """
    DELETE FROM "{table}"
"""

_DELETE_EXPIRED = """
    DELETE FROM "{table}"
    WHERE SECONDS_BETWEEN("CREATED_AT", CURRENT_TIMESTAMP) > ?
"""

_DELETE_LRU = """
    DELETE FROM "{table}"
    WHERE "PROMPT_HASH" NOT IN (
        SELECT "PROMPT_HASH" FROM "{table}"
        ORDER BY "LAST_ACCESSED" DESC
        LIMIT ?
    )
"""

_CHECK_TABLE = """
    SELECT COUNT(*) FROM SYS.TABLES
    WHERE SCHEMA_NAME = CURRENT_SCHEMA AND TABLE_NAME = ?
"""


# ──────────────────────────────────────────────────────────────────────
# Implementation
# ──────────────────────────────────────────────────────────────────────


class HANASemanticLLMCache(BaseCache):
    """Semantic cache for LLM responses using SAP HANA Cloud.

    Stores prompt embeddings and responses in a HANA table. On lookup,
    uses COSINE_SIMILARITY to find semantically similar prompts and
    returns cached responses when similarity exceeds the threshold.

    Args:
        connection: An hdbcli.dbapi.Connection to SAP HANA Cloud.
        embedding: A LangChain Embeddings instance for encoding prompts.
        table_name: Name of the cache table (default: ``"LLM_CACHE"``).
        similarity_threshold: Minimum cosine similarity for a cache hit
            (default: ``0.95``).
        ttl_seconds: Optional time-to-live in seconds. Entries older than
            this are ignored on lookup.
    """

    def __init__(
        self,
        connection: Any,
        embedding: Embeddings,
        table_name: str = _DEFAULT_TABLE,
        similarity_threshold: float = 0.95,
        ttl_seconds: Optional[int] = None,
    ) -> None:
        self.conn = connection
        self.embedding = embedding
        self.table_name = table_name
        self.similarity_threshold = similarity_threshold
        self.ttl_seconds = ttl_seconds
        self._vector_dim: Optional[int] = None
        self.setup()

    # ── Context manager ──────────────────────────────────────────────

    def __enter__(self) -> HANASemanticLLMCache:
        return self

    def __exit__(self, *args: Any) -> None:
        try:
            self.conn.close()
        except Exception:
            pass

    # ── Setup ────────────────────────────────────────────────────────

    def _detect_vector_dim(self) -> int:
        """Detect embedding dimension from the configured embedding model."""
        sample = self.embedding.embed_query("dimension probe")
        return len(sample)

    def _table_exists(self, cur: Any) -> bool:
        """Check if the cache table already exists via SYS.TABLES."""
        cur.execute(_CHECK_TABLE, (self.table_name,))
        row = cur.fetchone()
        return bool(row and row[0] > 0)

    def setup(self) -> None:
        """Create the cache table if it doesn't exist."""
        cur = self.conn.cursor()
        try:
            if not self._table_exists(cur):
                if self._vector_dim is None:
                    self._vector_dim = self._detect_vector_dim()
                ddl = _CREATE_TABLE_TEMPLATE.format(
                    table=self.table_name, dim=self._vector_dim
                )
                cur.execute(ddl)
                self.conn.commit()
                logger.info("Created cache table %s", self.table_name)
            else:
                logger.debug("Cache table %s already exists", self.table_name)
        finally:
            cur.close()

    # ── BaseCache interface ──────────────────────────────────────────

    def lookup(self, prompt: str, llm_string: str) -> Optional[RETURN_VAL_TYPE]:
        """Find a cached response for a semantically similar prompt.

        Args:
            prompt: The user prompt to look up.
            llm_string: Identifier for the LLM (scopes cache per model).

        Returns:
            A list of Generation objects if a cache hit is found, else None.
        """
        embedding_vector = self.embedding.embed_query(prompt)
        vector_str = str(embedding_vector)

        sql = _SELECT_SIMILAR.format(table=self.table_name)
        llm_hash = hash_prompt(llm_string)
        params: list[Any] = [vector_str, llm_hash, vector_str, self.similarity_threshold]

        if self.ttl_seconds is not None:
            sql += _SELECT_SIMILAR_TTL
            params.append(self.ttl_seconds)

        sql += _SELECT_SIMILAR_ORDER

        cur = self.conn.cursor()
        try:
            cur.execute(sql, params)
            row = cur.fetchone()

            if row is None:
                return None

            prompt_text, response_data, _llm_str, similarity = row
            logger.debug(
                "Cache hit for prompt (similarity=%.4f): %s",
                similarity, prompt[:80]
            )

            # Update access metadata
            prompt_hash = hash_prompt(prompt_text)
            cur.execute(
                _UPDATE_ACCESS.format(table=self.table_name), (prompt_hash,)
            )
            self.conn.commit()

            return deserialize_generations(response_data)
        finally:
            cur.close()

    def update(self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE) -> None:
        """Store a prompt-response pair in the cache.

        Args:
            prompt: The original prompt text.
            llm_string: Identifier for the LLM.
            return_val: The list of Generation objects to cache.
        """
        prompt_hash = hash_prompt(prompt)
        llm_hash = hash_prompt(llm_string)
        embedding_vector = self.embedding.embed_query(prompt)
        vector_str = str(embedding_vector)
        response_data = serialize_generations(list(return_val))

        cur = self.conn.cursor()
        try:
            cur.execute(
                _UPSERT_ENTRY.format(table=self.table_name),
                (prompt_hash, prompt, vector_str, llm_hash, llm_string, response_data),
            )
            self.conn.commit()
        finally:
            cur.close()

    def clear(self, **kwargs: Any) -> None:
        """Delete all cached entries."""
        cur = self.conn.cursor()
        try:
            cur.execute(_DELETE_ALL.format(table=self.table_name))
            self.conn.commit()
        finally:
            cur.close()

    # ── Eviction methods ─────────────────────────────────────────────

    def evict_expired(self) -> int:
        """Delete entries older than TTL.

        Returns:
            Number of deleted entries.

        Raises:
            ValueError: If no TTL is configured.
        """
        if self.ttl_seconds is None:
            raise ValueError("Cannot evict expired entries: no TTL configured")

        cur = self.conn.cursor()
        try:
            cur.execute(
                _DELETE_EXPIRED.format(table=self.table_name), (self.ttl_seconds,)
            )
            deleted = cur.rowcount
            self.conn.commit()
            logger.info("Evicted %d expired entries from %s", deleted, self.table_name)
            return deleted
        finally:
            cur.close()

    def evict_lru(self, max_entries: int = 1000) -> int:
        """Keep only the N most recently accessed entries.

        Args:
            max_entries: Maximum number of entries to keep.

        Returns:
            Number of deleted entries.
        """
        cur = self.conn.cursor()
        try:
            cur.execute(
                _DELETE_LRU.format(table=self.table_name), (max_entries,)
            )
            deleted = cur.rowcount
            self.conn.commit()
            logger.info(
                "Evicted %d LRU entries from %s (kept %d)",
                deleted, self.table_name, max_entries,
            )
            return deleted
        finally:
            cur.close()
