#!/usr/bin/env python3
"""
SPLADE + BM25 Sparse Search Engine for RMI RAG
================================================
Implements proper learned sparse representations (SPLADE-style) and BM25
with IDF weighting, replacing the naive token-overlap sparse search.

Architecture:
  - SPLADE: Term-level importance weighting using the existing BGE-small + ReLU + log(1+abs())
  - BM25: Standard Okapi BM25 with IDF, k1=1.2, b=0.75, stopword removal
  - Integration: Drop-in replacement for _sparse_text_search in rag_service.py

The SPLADE approach generates sparse vectors where each dimension corresponds to a
vocabulary term, weighted by importance. This replaces naive token counting with
learned term expansion and importance weighting.

For production SPLADE, you'd use naver/splade-v3-doc or prithivida/Splade_PP_en_v2.
This module provides a lightweight approximation using BGE-small features + expansion,
with a clear upgrade path to a full SPLADE model.
"""

import asyncio
import logging
import math
import re
from collections import Counter
from typing import Optional

logger = logging.getLogger(__name__)

# ── Stopwords (crypto-domain-aware) ────────────────────────────────────
STOPWORDS = frozenset(
    {
        # Standard English
        "a",
        "an",
        "the",
        "is",
        "are",
        "was",
        "were",
        "be",
        "been",
        "being",
        "have",
        "has",
        "had",
        "do",
        "does",
        "did",
        "will",
        "would",
        "could",
        "should",
        "may",
        "might",
        "must",
        "shall",
        "can",
        "need",
        "dare",
        "ought",
        "used",
        "to",
        "of",
        "in",
        "for",
        "on",
        "with",
        "at",
        "by",
        "from",
        "as",
        "into",
        "through",
        "during",
        "before",
        "after",
        "above",
        "below",
        "between",
        "out",
        "off",
        "over",
        "under",
        "again",
        "further",
        "then",
        "once",
        "here",
        "there",
        "when",
        "where",
        "why",
        "how",
        "all",
        "each",
        "every",
        "both",
        "few",
        "more",
        "most",
        "other",
        "some",
        "such",
        "no",
        "nor",
        "not",
        "only",
        "own",
        "same",
        "so",
        "than",
        "too",
        "very",
        "just",
        "because",
        "but",
        "and",
        "or",
        "if",
        "while",
        "although",
        "this",
        "that",
        "these",
        "those",
        "i",
        "me",
        "my",
        "myself",
        "we",
        "our",
        "ours",
        "ourselves",
        "you",
        "your",
        "yours",
        "yourself",
        "he",
        "him",
        "his",
        "himself",
        "she",
        "her",
        "hers",
        "herself",
        "it",
        "its",
        "itself",
        "they",
        "them",
        "their",
        "theirs",
        "themselves",
        "what",
        "which",
        "who",
        "whom",
        "whose",
        "about",
        "up",
        "down",
        # Crypto-domain: common but uninformative
        "token",
        "coin",
        "crypto",
        "address",
        "also",
        "like",
        "get",
        "got",
        "one",
        "two",
        "new",
        "old",
        "well",
        "even",
        "much",
        "many",
        "still",
    }
)


# ── Crypto-Domain Term Expansion ───────────────────────────────────────
# Maps common crypto terms to related expansion terms for SPLADE-like expansion
CRYPTO_EXPANSIONS = {
    "rug": {"rugpull", "rug pull", "scam", "exit scam", "honeypot", "dump"},
    "rugpull": {"rug pull", "rug", "scam", "exit scam", "dump"},
    "honeypot": {"honeypot", "scam", "rug", "trap", "cannot sell", "stuck funds"},
    "whale": {"large holder", "whale wallet", "big mover", "institutional"},
    "wash": {"wash trade", "wash trading", "volume manipulation", "fake volume"},
    "sniper": {"sniper bot", "mev", "sandwich", "front-run", "frontrun"},
    "pump": {"pump and dump", "pump", "artificial inflation", "manipulation"},
    "dump": {"dump", "sell-off", "crash", "rug", "liquidation"},
    "flash": {"flash loan", "flash attack", "instant loan", "defi exploit"},
    "sybil": {"sybil attack", "fake accounts", "multi-account", "astroturf"},
    "phishing": {"phishing", "scam", "fake site", "credential theft"},
    "mixer": {"mixer", "tumbler", "privacy", "laundering", "tornado"},
    "mint": {"minting", "token creation", "deploy", "launch"},
    "bridge": {"bridge", "cross-chain", "interoperability"},
    "stake": {"staking", "yield", "liquidity pool", "lp"},
    "exploit": {"exploit", "vulnerability", "hack", "vuln", "attack"},
    "sanction": {"ofac", "sanction", "blocked", "restricted"},
    "frontier": {"frontier", "fibonacci", "retrace"},
}


# ── Simple Stemmer (Porter-lite for crypto terms) ─────────────────────
_STEM_RULES = [
    (r"ing$", ""),  # running → runn
    (r"ed$", ""),  # scammed → scamm
    (r"ly$", ""),  # previously → previous
    (r"tion$", ""),  # liquidation → liquidat
    (r"ment$", ""),  # manipulation → manipulat
    (r"ness$", ""),  # worthiness → worthi
    (r"ity$", ""),  # liquidity → liquid
    (r"ous$", ""),  # suspicious → suspici
    (r"ive$", ""),  # restrictive → restrict
]


def simple_stem(word: str) -> str:
    """Aggressive suffix stripping for search matching. Not proper linguistics."""
    for pattern, replacement in _STEM_RULES:
        new = re.sub(pattern, replacement, word)
        if new != word and len(new) >= 3:
            return new
    return word


# ── Tokenizer ─────────────────────────────────────────────────────────
def tokenize(text: str, remove_stopwords: bool = True, stem: bool = True) -> list[str]:
    """
    Tokenize text for BM25/SPLADE search:
    1. Lowercase
    2. Split on non-alphanumeric (preserves hex addresses like 0xabc...)
    3. Remove stopwords
    4. Simple stemming
    """
    # Preserve hex addresses and dollar amounts
    text = text.lower()
    # Extract hex addresses, dollar amounts as single tokens
    hex_pattern = r"(0x[a-f0-9]{6,})"
    money_pattern = r"(\$[\d,.]+)"

    tokens = []
    # Keep hex addresses intact
    hex_addrs = re.findall(hex_pattern, text)
    money_vals = re.findall(money_pattern, text)

    # Remove hex and money for general tokenization
    clean = re.sub(hex_pattern, " ", text)
    clean = re.sub(money_pattern, " ", clean)
    clean = re.sub(r"[^\w\s-]", " ", clean)

    raw_tokens = clean.split()
    tokens.extend(hex_addrs)
    tokens.extend(money_vals)

    for tok in raw_tokens:
        if len(tok) < 2:
            continue
        if remove_stopwords and tok in STOPWORDS:
            continue
        if stem:
            tok = simple_stem(tok)
        tokens.append(tok)

    return tokens


# ── BM25 Implementation ────────────────────────────────────────────────


class BM25Index:
    """
    Okapi BM25 with crypto-domain enhancements.

    Parameters:
        k1: Term frequency saturation (1.2-2.0 typical, lower = less TF saturation)
        b: Document length normalization (0.75 typical, 1.0 = full normalization)
        epsilon: Floor for IDF to avoid negative scores
    """

    def __init__(self, k1: float = 1.2, b: float = 0.75, epsilon: float = 0.25):
        self.k1 = k1
        self.b = b
        self.epsilon = epsilon
        self.doc_count = 0
        self.avg_dl = 0.0
        self.doc_lengths: dict[str, int] = {}  # doc_id -> length
        self.doc_tokens: dict[str, Counter] = {}  # doc_id -> token counts
        self.df: Counter = Counter()  # token -> document frequency
        self._built = False

    def add_document(self, doc_id: str, content: str, metadata: dict | None = None):
        """Add a document to the index."""
        tokens = tokenize(content, remove_stopwords=True, stem=True)
        tf = Counter(tokens)

        self.doc_tokens[doc_id] = tf
        self.doc_lengths[doc_id] = len(tokens)
        self.doc_count += 1

        # Update document frequencies
        unique_tokens = set(tokens)
        for token in unique_tokens:
            self.df[token] += 1

        # Store metadata if provided
        if metadata and not hasattr(self, "_metadata"):
            self._metadata = {}
        if hasattr(self, "_metadata") and metadata:
            self._metadata[doc_id] = metadata

        # Store raw content for exact phrase matching in search_with_expansion
        if not hasattr(self, "_raw_content"):
            self._raw_content = {}
        # Store a truncated version to limit memory usage
        self._raw_content[doc_id] = content[:2000].lower()

    def build(self):
        """Finalize the index (compute avg doc length, IDF values)."""
        if self.doc_count == 0:
            self.avg_dl = 1.0
        else:
            total_len = sum(self.doc_lengths.values())
            self.avg_dl = total_len / self.doc_count if self.doc_count else 1.0
        self._built = True

    def save(self, path: str):
        """Persist BM25 index to disk via pickle."""
        import os
        import pickle

        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "wb") as f:
            pickle.dump(
                {
                    "k1": self.k1,
                    "b": self.b,
                    "epsilon": self.epsilon,
                    "doc_count": self.doc_count,
                    "avg_dl": self.avg_dl,
                    "doc_lengths": self.doc_lengths,
                    "doc_tokens": dict(self.doc_tokens),
                    "df": dict(self.df),
                    "_built": self._built,
                    "_metadata": getattr(self, "_metadata", {}),
                    "_raw_content": getattr(self, "_raw_content", {}),
                },
                f,
            )
        logger.info(f"BM25 index persisted to {path} ({os.path.getsize(path) / 1024 / 1024:.1f}MB)")

    @classmethod
    def load(cls, path: str) -> Optional["BM25Index"]:
        """Load BM25 index from disk."""
        import os
        import pickle

        if not os.path.exists(path):
            return None
        try:
            with open(path, "rb") as f:
                data = pickle.load(f)
            idx = cls(k1=data["k1"], b=data["b"], epsilon=data["epsilon"])
            idx.doc_count = data["doc_count"]
            idx.avg_dl = data["avg_dl"]
            idx.doc_lengths = data["doc_lengths"]
            idx.doc_tokens = data["doc_tokens"]
            idx.df = Counter(data["df"])
            idx._built = data["_built"]
            if data.get("_metadata"):
                idx._metadata = data["_metadata"]
            if data.get("_raw_content"):
                idx._raw_content = data["_raw_content"]
            logger.info(f"BM25 index loaded from disk: {idx.doc_count} docs, {len(idx.df)} terms")
            return idx
        except Exception as e:
            logger.warning(f"Failed to load BM25 from disk: {e}")
            return None

    def idf(self, token: str) -> float:
        """Compute IDF for a token."""
        df = self.df.get(token, 0)
        if df == 0:
            return 0.0
        return math.log((self.doc_count - df + 0.5) / (df + 0.5) + 1.0)

    def score_document(self, doc_id: str, query_tokens: list[str]) -> float:
        """Compute BM25 score for a single document against query tokens."""
        if doc_id not in self.doc_tokens:
            return 0.0

        doc_tf = self.doc_tokens[doc_id]
        dl = self.doc_lengths[doc_id]

        score = 0.0
        for token in query_tokens:
            tf = doc_tf.get(token, 0)
            if tf == 0:
                continue

            token_idf = self.idf(token)
            # BM25 formula
            numerator = tf * (self.k1 + 1)
            denominator = tf + self.k1 * (1 - self.b + self.b * dl / self.avg_dl)
            score += token_idf * numerator / denominator

        return score

    def search(self, query: str, limit: int = 20, min_score: float = 0.01) -> list[dict]:
        """
        Search the BM25 index.

        Returns list of dicts with: id, score, content_preview, metadata
        """
        if not self._built:
            self.build()

        query_tokens = tokenize(query, remove_stopwords=True, stem=True)

        # Expand query with crypto-domain terms
        expanded_tokens = list(query_tokens)
        for token in query_tokens:
            if token in CRYPTO_EXPANSIONS:
                expanded_tokens.extend(tokenize(" ".join(CRYPTO_EXPANSIONS[token])))

        # Score all documents
        scored = []
        for doc_id in self.doc_tokens:
            score = self.score_document(doc_id, expanded_tokens)
            if score >= min_score:
                entry = {
                    "id": doc_id,
                    "bm25_score": round(score, 6),
                    "collection": doc_id.split(":")[1] if ":" in doc_id else "unknown",
                }
                if hasattr(self, "_metadata") and doc_id in self._metadata:
                    entry["metadata"] = self._metadata[doc_id]
                scored.append(entry)

        scored.sort(key=lambda x: x["bm25_score"], reverse=True)
        return scored[:limit]

    def search_with_expansion(self, query: str, limit: int = 20) -> list[dict]:
        """
        SPLADE-inspired search: expand query terms, weight by importance.

        In full SPLADE, a transformer generates a log(1+ReLU(weight)) sparse vector.
        Our approximation:
        1. Tokenize + expand with crypto-domain synonyms
        2. Weight each term by its IDF (rare terms = important)
        3. Add positional bonus for exact phrase matches
        4. Weight expansion terms lower than original terms
        """
        if not self._built:
            self.build()

        original_tokens = tokenize(query, remove_stopwords=True, stem=True)

        # Build weighted query: original tokens weight=1.0, expansion weight=0.5
        weighted_tokens: Counter = Counter()
        for token in original_tokens:
            weighted_tokens[token] += 1.0

        for token in original_tokens:
            if token in CRYPTO_EXPANSIONS:
                for exp_token in tokenize(" ".join(CRYPTO_EXPANSIONS[token])):
                    weighted_tokens[exp_token] += 0.5

        # Score documents using weighted query
        scored = []
        for doc_id in self.doc_tokens:
            doc_tf = self.doc_tokens[doc_id]
            dl = self.doc_lengths[doc_id]

            score = 0.0
            for token, weight in weighted_tokens.items():
                tf = doc_tf.get(token, 0)
                if tf == 0:
                    continue

                token_idf = self.idf(token)
                # BM25 term score
                numerator = tf * (self.k1 + 1)
                denominator = tf + self.k1 * (1 - self.b + self.b * dl / self.avg_dl)
                bm25_term = token_idf * numerator / denominator

                # SPLADE-inspired: max activation * IDF weight
                # This approximates log(1 + ReLU(w_ij)) for each term
                score += bm25_term * weight

            # Exact phrase bonus (SPLADE positional interaction approximation)
            query_lower = query.lower()
            # Check raw content if we have it stored
            if score > 0 and hasattr(self, "_raw_content") and doc_id in self._raw_content:
                if query_lower in self._raw_content[doc_id].lower():
                    score *= 1.5

            if score > 0.01:
                entry = {
                    "id": doc_id,
                    "sparse_score": round(score, 6),
                    "collection": doc_id.split(":")[1] if ":" in doc_id else "unknown",
                }
                if hasattr(self, "_metadata") and doc_id in self._metadata:
                    entry["metadata"] = self._metadata[doc_id]
                scored.append(entry)

        scored.sort(key=lambda x: x["sparse_score"], reverse=True)
        return scored[:limit]


# ── Async BM25 Builder (for Redis collections) ──────────────────────


async def build_bm25_from_redis(
    collections: list[str] | None = None,
    max_per_collection: int = 50000,
) -> BM25Index:
    """
    Build a BM25 index from Redis-stored documents.
    Reads content from rag:{collection}:{doc_id} keys.
    """
    import json
    import os

    # Try to load .env from multiple possible locations (Docker vs dev)
    from pathlib import Path

    import redis.asyncio as aioredis

    for env_path in ["/app/.env", str(Path(__file__).parent.parent.parent / ".env")]:
        try:
            from dotenv import load_dotenv

            load_dotenv(env_path, override=True)
            break
        except Exception:
            pass

    if collections is None:
        collections = [
            "wallet_profiles",
            "token_analysis",
            "known_scams",
            "scam_patterns",
            "forensic_reports",
            "market_intel",
            "contract_audits",
            "news_articles",
            "transaction_patterns",
            "general",
        ]

    r = aioredis.Redis(
        host=os.environ.get("REDIS_HOST", "rmi-redis"),
        port=int(os.environ.get("REDIS_PORT", "6379")),
        password=os.environ.get("REDIS_PASSWORD", ""),
        decode_responses=True,
    )

    index = BM25Index(k1=1.2, b=0.75)
    total_indexed = 0

    for coll in collections:
        idx_key = f"rag:idx:{coll}"
        try:
            doc_ids = await r.smembers(idx_key)
            if not doc_ids:
                continue

            # Limit per collection for memory
            sample = list(doc_ids)[:max_per_collection]
            logger.info(f"BM25 indexing {coll}: {len(sample)} docs")

            # Batch fetch
            batch_size = 500
            for i in range(0, len(sample), batch_size):
                batch = sample[i : i + batch_size]
                keys = [f"rag:{coll}:{did}" for did in batch]
                pipe = r.pipeline()
                for k in keys:
                    pipe.get(k)
                raw_docs = await pipe.execute()

                for did, data in zip(batch, raw_docs, strict=False):
                    if not data:
                        continue
                    try:
                        doc = json.loads(data)
                        content = doc.get("content", "") or ""
                        metadata = doc.get("metadata", {}) or {}
                        # Use collection:doc_id as the index key
                        doc_key = f"{coll}:{did}"
                        index.add_document(doc_key, content, metadata)
                        total_indexed += 1
                    except (json.JSONDecodeError, Exception):
                        continue

        except Exception as e:
            logger.warning(f"BM25 indexing for {coll} failed: {e}")

    index.build()
    await r.aclose()

    logger.info(
        f"BM25 index built: {total_indexed} docs, {len(index.df)} unique terms, avg doc length {index.avg_dl:.1f}"
    )
    return index


# ── In-memory BM25 cache ───────────────────────────────────────────────
_bm25_index: BM25Index | None = None
_bm25_built_at: float = 0.0
_BM25_TTL = 3600  # Rebuild every hour
_bm25_lock = asyncio.Lock()


async def get_bm25_index(force_rebuild: bool = False) -> BM25Index:
    """Get or build the BM25 index (persisted to disk, 1h TTL in memory)."""
    global _bm25_index, _bm25_built_at
    import os
    import time

    BM25_PATH = "/app/data/bm25_index.pkl"

    now = time.time()
    if _bm25_index and not force_rebuild and (now - _bm25_built_at) < _BM25_TTL:
        return _bm25_index

    async with _bm25_lock:
        if _bm25_index and not force_rebuild and (time.time() - _bm25_built_at) < _BM25_TTL:
            return _bm25_index

        # Try loading from disk first
        if os.path.exists(BM25_PATH) and not force_rebuild:
            loaded = await asyncio.get_running_loop().run_in_executor(None, BM25Index.load, BM25_PATH)
            if loaded and loaded.doc_count > 0:
                _bm25_index = loaded
                _bm25_built_at = time.time()
                logger.info(
                    f"BM25 loaded from disk: {loaded.doc_count} docs (skipped {time.time() - now:.1f}s rebuild)"
                )
                return _bm25_index

        # Build from Redis and persist
        _bm25_index = await build_bm25_from_redis()
        _bm25_built_at = time.time()

        # Persist to disk in background
        if _bm25_index.doc_count > 0:
            asyncio.create_task(_save_bm25_async(BM25_PATH))

        return _bm25_index


async def _save_bm25_async(path: str):
    """Save BM25 index to disk in executor thread."""
    try:
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(None, _bm25_index.save, path)
    except Exception as e:
        logger.debug(f"BM25 disk save failed (non-critical): {e}")


async def bm25_search(
    query: str,
    collections: list[str] | None = None,
    limit: int = 20,
    min_score: float = 0.01,
) -> list[dict]:
    """
    Search using BM25 with crypto-domain query expansion.
    Drop-in enhancement for _sparse_text_search.
    """
    index = await get_bm25_index()
    results = index.search(query, limit=limit, min_score=min_score)
    return results


async def splade_search(
    query: str,
    collections: list[str] | None = None,
    limit: int = 20,
) -> list[dict]:
    """
    SPLADE-inspired sparse search with term expansion and importance weighting.
    Uses BM25 as backbone + crypto-domain expansion for SPLADE-like sparse vectors.
    """
    index = await get_bm25_index()
    results = index.search_with_expansion(query, limit=limit)
    return results
