"""LLMHosts Tier 1 cache -- entity-aware namespace bucketing.

Extracts programming languages, frameworks, and domain keywords from
prompts and routes cache lookups to per-namespace buckets.  This prevents
cross-domain false positives: "Python sorting algorithm" and "JavaScript
sorting algorithm" land in different buckets.

Requires the ``[smart]`` pip tier for cosine-similarity matching within
buckets (numpy).  Falls back to exact-key-only matching otherwise.
"""

from __future__ import annotations

import hashlib
import json
import logging
import re
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, ClassVar

import aiosqlite
from pydantic import BaseModel, ConfigDict

if TYPE_CHECKING:
    from pathlib import Path

    from llmhosts.cache.models import CacheEntry

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Optional dependency guard
# ---------------------------------------------------------------------------

NAMESPACE_AVAILABLE = False
try:
    import numpy as np

    NAMESPACE_AVAILABLE = True
except ImportError:  # pragma: no cover
    np = None  # type: ignore[assignment]

# ---------------------------------------------------------------------------
# Pydantic models
# ---------------------------------------------------------------------------


class NamespaceKey(BaseModel):
    """Combined entity key used for namespace bucketing."""

    model_config = ConfigDict(frozen=True)

    languages: list[str]
    frameworks: list[str]
    domains: list[str]
    key: str  # deterministic string key for bucketing

    @staticmethod
    def from_entities(
        langs: list[str],
        frameworks: list[str],
        domains: list[str],
    ) -> NamespaceKey:
        """Generate a namespace key from extracted entities.

        Entities are sorted alphabetically within each category so that
        ``{"python", "sql"}`` and ``{"sql", "python"}`` produce the same
        key.  When no entities are detected, the key falls back to
        ``"__default__"``.
        """
        s_langs = sorted(set(langs))
        s_fw = sorted(set(frameworks))
        s_dom = sorted(set(domains))

        parts: list[str] = []
        if s_langs:
            parts.append("lang:" + "+".join(s_langs))
        if s_fw:
            parts.append("fw:" + "+".join(s_fw))
        if s_dom:
            parts.append("dom:" + "+".join(s_dom))

        key = "|".join(parts) if parts else "__default__"

        return NamespaceKey(languages=s_langs, frameworks=s_fw, domains=s_dom, key=key)


class NamespaceStats(BaseModel):
    """Per-namespace bucket statistics."""

    total_namespaces: int
    total_entries: int
    eviction_count: int = 0
    entries_by_namespace: dict[str, int]
    hit_rate_by_namespace: dict[str, float]


# ---------------------------------------------------------------------------
# SQL schema
# ---------------------------------------------------------------------------

_CREATE_TABLE = """
CREATE TABLE IF NOT EXISTS namespace_cache (
    cache_key      TEXT PRIMARY KEY,
    namespace_key  TEXT NOT NULL,
    model          TEXT NOT NULL,
    prompt_hash    TEXT NOT NULL,
    response_json  TEXT NOT NULL,
    embedding_blob BLOB,
    created_at     TEXT NOT NULL,
    expires_at     TEXT NOT NULL,
    hit_count      INTEGER NOT NULL DEFAULT 0,
    last_hit_at    TEXT,
    size_bytes     INTEGER NOT NULL
);
"""

_CREATE_INDEXES = [
    "CREATE INDEX IF NOT EXISTS idx_ns_namespace ON namespace_cache (namespace_key);",
    "CREATE INDEX IF NOT EXISTS idx_ns_model ON namespace_cache (model);",
    "CREATE INDEX IF NOT EXISTS idx_ns_expires ON namespace_cache (expires_at);",
]

_CREATE_STATS_TABLE = """
CREATE TABLE IF NOT EXISTS namespace_stats (
    namespace_key  TEXT PRIMARY KEY,
    hit_count      INTEGER NOT NULL DEFAULT 0,
    miss_count     INTEGER NOT NULL DEFAULT 0
);
"""

# ---------------------------------------------------------------------------
# Namespace filter
# ---------------------------------------------------------------------------


class NamespaceFilter:
    """Tier 1 cache: entity-aware namespace bucketing.

    Extracts key entities from prompts and routes to namespace buckets.
    'Python sorting algorithm' and 'JavaScript sorting algorithm' go to
    different buckets, preventing cross-domain false positives.
    """

    # Entity categories for namespace generation
    PROGRAMMING_LANGS: ClassVar[set[str]] = {
        "python",
        "javascript",
        "typescript",
        "java",
        "c++",
        "c#",
        "go",
        "rust",
        "ruby",
        "php",
        "swift",
        "kotlin",
        "scala",
        "r",
        "sql",
        "html",
        "css",
        "shell",
        "bash",
        "perl",
        "lua",
        "dart",
        "elixir",
        "haskell",
        "clojure",
        "matlab",
        "fortran",
    }

    FRAMEWORKS: ClassVar[set[str]] = {
        "react",
        "vue",
        "angular",
        "django",
        "flask",
        "fastapi",
        "express",
        "nextjs",
        "next.js",
        "spring",
        "rails",
        "laravel",
        "dotnet",
        ".net",
        "tensorflow",
        "pytorch",
        "pandas",
        "numpy",
        "kubernetes",
        "docker",
        "aws",
        "gcp",
        "azure",
        "svelte",
        "tailwind",
        "bootstrap",
    }

    DOMAINS: ClassVar[set[str]] = {
        "web",
        "mobile",
        "backend",
        "frontend",
        "devops",
        "ml",
        "ai",
        "data",
        "security",
        "networking",
        "database",
        "api",
        "microservices",
        "cloud",
        "embedded",
        "game",
        "blockchain",
    }

    # Pre-compiled regex patterns for word-boundary matching.
    # Special cases: "c++", "c#", ".net", "next.js" need careful handling.
    _SPECIAL_TOKENS: ClassVar[dict[str, str]] = {
        "c++": r"c\+\+",
        "c#": r"c#",
        ".net": r"\.net",
        "next.js": r"next\.js",
    }

    def __init__(
        self,
        db_path: Path,
        similarity_threshold: float = 0.88,
        ttl_seconds: int = 3600,
        max_namespaces: int = 500,
    ) -> None:
        self._db_path = db_path
        self._similarity_threshold = similarity_threshold
        self._ttl_seconds = ttl_seconds
        self._max_namespaces = max_namespaces
        self._db: aiosqlite.Connection | None = None
        self._evictions = 0

        # Build compiled patterns once
        self._lang_patterns: list[tuple[str, re.Pattern[str]]] = self._build_patterns(self.PROGRAMMING_LANGS)
        self._fw_patterns: list[tuple[str, re.Pattern[str]]] = self._build_patterns(self.FRAMEWORKS)
        self._domain_patterns: list[tuple[str, re.Pattern[str]]] = self._build_patterns(self.DOMAINS)

    # ------------------------------------------------------------------
    # Pattern building
    # ------------------------------------------------------------------

    @classmethod
    def _build_patterns(cls, tokens: set[str]) -> list[tuple[str, re.Pattern[str]]]:
        """Build word-boundary regex patterns for a set of entity tokens."""
        patterns: list[tuple[str, re.Pattern[str]]] = []
        for token in sorted(tokens):
            raw = cls._SPECIAL_TOKENS[token] if token in cls._SPECIAL_TOKENS else re.escape(token)
            # Word boundary on both sides, case-insensitive
            patterns.append((token, re.compile(rf"(?<![a-zA-Z0-9_]){raw}(?![a-zA-Z0-9_])", re.IGNORECASE)))
        return patterns

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def initialize(self) -> None:
        """Create namespace cache tables."""
        self._db_path.parent.mkdir(parents=True, exist_ok=True)
        self._db = await aiosqlite.connect(str(self._db_path))
        await self._db.execute("PRAGMA journal_mode=WAL;")
        await self._db.execute("PRAGMA synchronous=NORMAL;")
        await self._db.execute(_CREATE_TABLE)
        for idx_sql in _CREATE_INDEXES:
            await self._db.execute(idx_sql)
        await self._db.execute(_CREATE_STATS_TABLE)
        await self._db.commit()
        logger.info("NamespaceFilter initialised (db=%s)", self._db_path)

    async def close(self) -> None:
        """Close database connection."""
        if self._db is not None:
            await self._db.close()
            self._db = None
            logger.debug("NamespaceFilter closed")

    def _ensure_db(self) -> aiosqlite.Connection:
        if self._db is None:
            raise RuntimeError("NamespaceFilter not initialised -- call initialize() first")
        return self._db

    # ------------------------------------------------------------------
    # Entity extraction
    # ------------------------------------------------------------------

    def _extract_text(self, messages: list[dict[str, Any]]) -> str:
        """Concatenate all message content into a single search string."""
        parts: list[str] = []
        for msg in messages:
            content = msg.get("content", "")
            if isinstance(content, list):
                for block in content:
                    if isinstance(block, dict) and block.get("type") == "text":
                        parts.append(str(block.get("text", "")))
                    elif isinstance(block, str):
                        parts.append(block)
            else:
                parts.append(str(content))
        return " ".join(parts)

    def _match_patterns(self, text: str, patterns: list[tuple[str, re.Pattern[str]]]) -> list[str]:
        """Return all tokens whose patterns match in *text*."""
        found: list[str] = []
        for token, pat in patterns:
            if pat.search(text):
                found.append(token)
        return found

    def extract_entities(self, messages: list[dict[str, Any]]) -> NamespaceKey:
        """Extract entities from prompt messages.

        Scans all message content for:
        - Programming languages mentioned
        - Frameworks/tools mentioned
        - Domain keywords

        Returns a :class:`NamespaceKey` combining the top entities.
        """
        text = self._extract_text(messages)

        langs = self._match_patterns(text, self._lang_patterns)
        fws = self._match_patterns(text, self._fw_patterns)
        doms = self._match_patterns(text, self._domain_patterns)

        return NamespaceKey.from_entities(langs, fws, doms)

    # ------------------------------------------------------------------
    # Hashing helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _prompt_hash(model: str, messages: list[dict[str, Any]]) -> str:
        """SHA-256 of normalised prompt for de-duplication within a namespace."""
        norm_model = model.strip().lower()
        normalised_msgs: list[dict[str, str]] = []
        for msg in messages:
            role = str(msg.get("role", "")).strip().lower()
            if role == "system":
                continue
            content = msg.get("content", "")
            if isinstance(content, list):
                text_parts = [
                    str(b.get("text", "")).strip().lower()
                    if isinstance(b, dict) and b.get("type") == "text"
                    else str(b).strip().lower()
                    for b in content
                    if isinstance(b, (dict, str))
                ]
                content_str = " ".join(text_parts)
            else:
                content_str = str(content).strip().lower()
            normalised_msgs.append({"role": role, "content": content_str})
        normalised_msgs.sort(key=lambda m: (m["role"], m["content"]))
        canonical = json.dumps(
            {"model": norm_model, "messages": normalised_msgs}, sort_keys=True, separators=(",", ":")
        )
        return hashlib.sha256(canonical.encode("utf-8")).hexdigest()

    @staticmethod
    def _now() -> datetime:
        return datetime.now(timezone.utc)

    @staticmethod
    def _iso(dt: datetime) -> str:
        return dt.isoformat()

    # ------------------------------------------------------------------
    # Cosine similarity
    # ------------------------------------------------------------------

    @staticmethod
    def _cosine_similarity(a: Any, b: Any) -> float:
        """Cosine similarity between two numpy arrays."""
        if np is None:
            return 0.0
        a_arr = np.asarray(a, dtype=np.float32).ravel()
        b_arr = np.asarray(b, dtype=np.float32).ravel()
        dot = float(np.dot(a_arr, b_arr))
        norm = float(np.linalg.norm(a_arr) * np.linalg.norm(b_arr))
        if norm < 1e-12:
            return 0.0
        return dot / norm

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def get(
        self,
        model: str,
        messages: list[dict[str, Any]],
        embedding: Any | None = None,
    ) -> CacheEntry | None:
        """Look up in namespace-filtered cache.

        1. Extract entities to determine namespace
        2. Search only within that namespace
        3. If embedding provided, use cosine similarity within namespace
        4. Return best match if similarity > threshold
        """
        from llmhosts.cache.models import CacheEntry

        db = self._ensure_db()
        ns_key = self.extract_entities(messages)
        now = self._now()

        # First, try exact prompt hash within namespace
        p_hash = self._prompt_hash(model, messages)
        async with db.execute(
            "SELECT cache_key, model, prompt_hash, response_json, created_at, "
            "expires_at, hit_count, last_hit_at, size_bytes "
            "FROM namespace_cache "
            "WHERE namespace_key = ? AND prompt_hash = ? AND model = ?",
            (ns_key.key, p_hash, model.strip().lower()),
        ) as cursor:
            row = await cursor.fetchone()

        if row is not None:
            expires_at = datetime.fromisoformat(row[5]) if row[5] else None
            if expires_at is None or expires_at > now:
                # Hit on exact match within namespace
                new_hits = row[6] + 1
                await db.execute(
                    "UPDATE namespace_cache SET hit_count = ?, last_hit_at = ? WHERE cache_key = ?",
                    (new_hits, self._iso(now), row[0]),
                )
                await self._record_hit(db, ns_key.key)
                await db.commit()
                logger.debug("Namespace HIT (exact) ns=%s key=%s", ns_key.key, row[0][:12])
                return CacheEntry(
                    cache_key=row[0],
                    model=row[1],
                    prompt_hash=row[2],
                    response_json=row[3],
                    created_at=datetime.fromisoformat(row[4]),
                    expires_at=datetime.fromisoformat(row[5]),
                    hit_count=new_hits,
                    last_hit_at=now,
                    size_bytes=row[8],
                )
            else:
                # Expired -- remove
                await db.execute("DELETE FROM namespace_cache WHERE cache_key = ?", (row[0],))
                await db.commit()

        # If no embedding provided or numpy unavailable, nothing else we can do
        if embedding is None or np is None:
            await self._record_miss(db, ns_key.key)
            await db.commit()
            return None

        # Semantic search within namespace: load all embeddings for this namespace + model
        async with db.execute(
            "SELECT cache_key, model, prompt_hash, response_json, created_at, "
            "expires_at, hit_count, last_hit_at, size_bytes, embedding_blob "
            "FROM namespace_cache "
            "WHERE namespace_key = ? AND model = ? AND expires_at > ?",
            (ns_key.key, model.strip().lower(), self._iso(now)),
        ) as cursor:
            rows = await cursor.fetchall()

        if not rows:
            await self._record_miss(db, ns_key.key)
            await db.commit()
            return None

        best_row = None
        best_sim = -1.0
        for r in rows:
            blob = r[9]
            if blob is None:
                continue
            cached_emb = np.frombuffer(blob, dtype=np.float32)
            sim = self._cosine_similarity(embedding, cached_emb)
            if sim > best_sim:
                best_sim = sim
                best_row = r

        if best_row is not None and best_sim >= self._similarity_threshold:
            new_hits = best_row[6] + 1
            await db.execute(
                "UPDATE namespace_cache SET hit_count = ?, last_hit_at = ? WHERE cache_key = ?",
                (new_hits, self._iso(now), best_row[0]),
            )
            await self._record_hit(db, ns_key.key)
            await db.commit()
            logger.debug(
                "Namespace HIT (semantic, sim=%.4f) ns=%s key=%s",
                best_sim,
                ns_key.key,
                best_row[0][:12],
            )
            return CacheEntry(
                cache_key=best_row[0],
                model=best_row[1],
                prompt_hash=best_row[2],
                response_json=best_row[3],
                created_at=datetime.fromisoformat(best_row[4]),
                expires_at=datetime.fromisoformat(best_row[5]),
                hit_count=new_hits,
                last_hit_at=now,
                size_bytes=best_row[8],
            )

        await self._record_miss(db, ns_key.key)
        await db.commit()
        return None

    async def put(
        self,
        model: str,
        messages: list[dict[str, Any]],
        response_json: str,
        embedding: Any | None = None,
    ) -> str:
        """Store response in appropriate namespace bucket."""
        db = self._ensure_db()
        ns_key = self.extract_entities(messages)
        p_hash = self._prompt_hash(model, messages)
        now = self._now()
        expires_at = now + timedelta(seconds=self._ttl_seconds)

        cache_key = hashlib.sha256(
            f"{ns_key.key}:{p_hash}".encode(),
        ).hexdigest()

        # Serialise embedding as bytes if available
        emb_blob: bytes | None = None
        if embedding is not None and np is not None:
            emb_blob = np.asarray(embedding, dtype=np.float32).tobytes()

        size_bytes = len(response_json.encode("utf-8"))

        await db.execute(
            "INSERT OR REPLACE INTO namespace_cache "
            "(cache_key, namespace_key, model, prompt_hash, response_json, "
            "embedding_blob, created_at, expires_at, hit_count, last_hit_at, size_bytes) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, NULL, ?)",
            (
                cache_key,
                ns_key.key,
                model.strip().lower(),
                p_hash,
                response_json,
                emb_blob,
                self._iso(now),
                self._iso(expires_at),
                size_bytes,
            ),
        )
        await db.commit()
        await self._evict_if_needed()
        logger.debug("Namespace PUT ns=%s key=%s model=%s", ns_key.key, cache_key[:12], model)
        return cache_key

    async def stats(self) -> NamespaceStats:
        """Stats per namespace bucket."""
        db = self._ensure_db()

        # Entry counts by namespace
        entries_by_ns: dict[str, int] = {}
        async with db.execute(
            "SELECT namespace_key, COUNT(*) FROM namespace_cache GROUP BY namespace_key",
        ) as cur:
            async for row in cur:
                entries_by_ns[row[0]] = row[1]

        # Hit rates by namespace
        hit_rate_by_ns: dict[str, float] = {}
        async with db.execute("SELECT namespace_key, hit_count, miss_count FROM namespace_stats") as cur:
            async for row in cur:
                total = row[1] + row[2]
                hit_rate_by_ns[row[0]] = round(row[1] / total, 4) if total > 0 else 0.0

        total_entries = sum(entries_by_ns.values())

        return NamespaceStats(
            total_namespaces=len(entries_by_ns),
            total_entries=total_entries,
            eviction_count=self._evictions,
            entries_by_namespace=entries_by_ns,
            hit_rate_by_namespace=hit_rate_by_ns,
        )

    # ------------------------------------------------------------------
    # LRU eviction
    # ------------------------------------------------------------------

    async def _evict_if_needed(self) -> None:
        """Evict least-recently-used namespaces when namespace count exceeds limit."""
        db = self._ensure_db()

        async with db.execute("SELECT COUNT(DISTINCT namespace_key) FROM namespace_cache") as cur:
            row = await cur.fetchone()
            ns_count = row[0] if row else 0

        if ns_count <= self._max_namespaces:
            return

        # Find LRU namespaces: ordered by their most recent hit (oldest first)
        evicted = 0
        while ns_count > self._max_namespaces:
            async with db.execute(
                "SELECT namespace_key, MAX(COALESCE(last_hit_at, '1970-01-01')) as last_access "
                "FROM namespace_cache "
                "GROUP BY namespace_key "
                "ORDER BY last_access ASC "
                "LIMIT 10",
            ) as cur:
                rows = await cur.fetchall()

            if not rows:
                break

            for row in rows:
                ns_key_to_evict = row[0]
                await db.execute("DELETE FROM namespace_cache WHERE namespace_key = ?", (ns_key_to_evict,))
                await db.execute("DELETE FROM namespace_stats WHERE namespace_key = ?", (ns_key_to_evict,))
                evicted += 1
                ns_count -= 1
                if ns_count <= self._max_namespaces:
                    break

        if evicted > 0:
            self._evictions += evicted
            await db.commit()
            try:
                from llmhosts.metrics.prometheus import record_cache_eviction

                record_cache_eviction("namespace", evicted)
            except ImportError:
                pass
            logger.info("Evicted %d namespace buckets (max_namespaces=%d)", evicted, self._max_namespaces)

    # ------------------------------------------------------------------
    # Internal stats tracking
    # ------------------------------------------------------------------

    async def _record_hit(self, db: aiosqlite.Connection, ns_key: str) -> None:
        await db.execute(
            "INSERT INTO namespace_stats (namespace_key, hit_count, miss_count) VALUES (?, 1, 0) "
            "ON CONFLICT(namespace_key) DO UPDATE SET hit_count = hit_count + 1",
            (ns_key,),
        )

    async def _record_miss(self, db: aiosqlite.Connection, ns_key: str) -> None:
        await db.execute(
            "INSERT INTO namespace_stats (namespace_key, hit_count, miss_count) VALUES (?, 0, 1) "
            "ON CONFLICT(namespace_key) DO UPDATE SET miss_count = miss_count + 1",
            (ns_key,),
        )
