#!/usr/bin/env python3
"""Two-layer semantic cache: exact match first, then cosine similarity.

Self-contained (no external vector DB): uses a simple in-memory embedding
(simulated) + numpy cosine. Swap the embed() function for a real embedding
model in production.
"""

from __future__ import annotations

import hashlib
import time
from dataclasses import dataclass, field
from typing import Callable, Optional


# ── Simulated embedding (replace with sentence-transformers / API in prod) ─

def embed(text: str) -> list[float]:
    """Deterministic fake embedding: hashes token bigrams to vector dims.

    Real usage: model.encode(text) from sentence-transformers, or an embedding
    API call. This stands in so the demo runs with zero dependencies.
    """
    tokens = text.lower().split()
    vec = [0.0] * 64
    for i in range(len(tokens)):
        for gram in (tokens[i], " ".join(tokens[i:i + 2])):
            h = int(hashlib.md5(gram.encode()).hexdigest(), 16)
            vec[h % 64] += 1.0
    norm = sum(v * v for v in vec) ** 0.5 or 1.0
    return [v / norm for v in vec]


def cosine(a: list[float], b: list[float]) -> float:
    return sum(x * y for x, y in zip(a, b))


# ── Cache store ────────────────────────────────────────────────────────────

@dataclass
class Entry:
    question: str
    answer: str
    embedding: list[float]
    created_at: float


class SemanticCache:
    """Exact-match layer + semantic (cosine) layer, then fall back to LLM."""

    def __init__(
        self,
        threshold: float = 0.93,
        llm_fn: Optional[Callable[[str], str]] = None,
    ) -> None:
        self.threshold = threshold
        self.llm_fn = llm_fn or (lambda q: f"[LLM answer for: {q}]")
        self.exact: dict[str, str] = {}
        self.semantic: list[Entry] = []
        self.hits = 0
        self.misses = 0

    def _key(self, question: str, context: str = "") -> str:
        """Hash question + context so different tasks don't collide."""
        return hashlib.sha256(f"{context}||{question}".encode()).hexdigest()

    def get(self, question: str, context: str = "") -> tuple[str, str]:
        """Return (answer, provenance). Provenance: exact|semantic|miss."""
        key = self._key(question, context)

        # Layer 1: exact match (zero risk, instant)
        if key in self.exact:
            self.hits += 1
            return self.exact[key], "exact"

        # Layer 2: semantic similarity
        q_emb = embed(question)
        for entry in self.semantic:
            if cosine(q_emb, entry.embedding) >= self.threshold:
                self.hits += 1
                return entry.answer, "semantic"

        # Layer 3: LLM (last resort)
        self.misses += 1
        answer = self.llm_fn(question)
        self._store(key, question, answer, q_emb)
        return answer, "miss"

    def _store(self, key: str, question: str, answer: str, emb: list[float]) -> None:
        self.exact[key] = answer
        self.semantic.append(Entry(question, answer, emb, time.time()))

    def stats(self) -> dict[str, int]:
        return {
            "hits": self.hits,
            "misses": self.misses,
            "hit_rate": self.hits / (self.hits + self.misses) if (self.hits + self.misses) else 0,
        }


# ── Demo ───────────────────────────────────────────────────────────────────

def main() -> None:
    cache = SemanticCache(threshold=0.93)

    queries = [
        "How do I reset my password?",            # miss -> LLM
        "How do I reset my password?",            # exact hit
        "What's the way to change my password?",  # semantic hit (paraphrase)
        "Tell me how to recover my password",     # semantic hit
        "What's the capital of France?",          # miss (unrelated)
    ]

    for q in queries:
        answer, provenance = cache.get(q, context="helpdesk")
        print(f"  [{provenance:9s}] {q!r} -> {answer[:50]!r}")

    print("\n=== Stats ===")
    for k, v in cache.stats().items():
        print(f"  {k}: {v}")

    print("\nProduction notes:")
    print("  - Swap embed() for sentence-transformers or an embedding API")
    print("  - Use Redis + RediSearch or a vector DB for persistence")
    print("  - Scope cache keys by user/tenant to avoid cross-user leaks")
    print("  - Never cache mutable-state answers (inventory, prices)")


if __name__ == "__main__":
    main()