"""
RAG Ingestion & Retrieval Service v4
====================================
Three-pillar hybrid search (dense + sparse + entity) with:
  - Knowledge Graph expansion (Pillar 3+)
  - MMR deduplication after RRF fusion
  - Parent-child retrieval for contextual chunks
  - SPLADE+BM25 sparse search (Pillar 2)
  - Query-adaptive fusion weights

Collections: wallet_profiles, token_analysis, scam_patterns,
             forensic_reports, market_intel, contract_audits, known_scams
"""

import asyncio
import hashlib
import json
import logging
import os
import re
from datetime import UTC, datetime
from typing import Any

from app.crypto_embeddings import (
    COLLECTIONS,
    KNOWN_SCAM_PATTERNS,
    CryptoEmbedder,
    EmbeddingResult,
    extract_contract_features,
    get_embedder,
)

logger = logging.getLogger(__name__)

REDIS_HOST = os.getenv("REDIS_HOST", "127.0.0.1")
REDIS_PORT = int(os.getenv("REDIS_PORT", "6379"))
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", "")

# Track seeding state
_seeded = False

# Scam pattern pre-embedding cache (Feature 3)
_pattern_cache: dict[str, EmbeddingResult] = {}
_pattern_cache_loaded = False

# Redis singleton - reuses connection across all calls
_redis_pool = None


# Singleton accessor for dependency injectiondef get_rag_service():    import sys    return sys.modules[__name__]
async def _get_redis():
    """Get or create a shared Redis connection (singleton pattern to prevent leaks)."""
    global _redis_pool
    import redis.asyncio as redis

    if _redis_pool is not None:
        try:
            await _redis_pool.ping()
        except Exception:
            _redis_pool = None
    if _redis_pool is None:
        _redis_pool = redis.Redis(
            host=REDIS_HOST,
            port=REDIS_PORT,
            password=REDIS_PASSWORD or None,
            db=0,
            decode_responses=True,
        )
    return _redis_pool


async def _preembed_scam_patterns():
    """
    Pre-embed all KNOWN_SCAM_PATTERNS and cache results.
    Loads from Redis first (persistent cache), then embeds any missing ones.
    """
    global _pattern_cache, _pattern_cache_loaded

    embedder = await get_embedder()
    r = await _get_redis()

    for pattern in KNOWN_SCAM_PATTERNS:
        pattern_name = pattern["name"]

        # 1. Try in-memory cache
        if pattern_name in _pattern_cache:
            continue

        # 2. Try Redis persistent cache (30-day TTL)
        redis_key = f"rag:pattern_cache:{pattern_name}"
        try:
            cached_data = await r.get(redis_key)
            if cached_data:
                cached = json.loads(cached_data)
                _pattern_cache[pattern_name] = EmbeddingResult(
                    vector=cached["vector"],
                    dims=cached["dims"],
                    model=cached["model"],
                    head=cached.get("head", "scam_pattern"),
                )
                logger.debug(f"Pattern cache hit (Redis): {pattern_name}")
                continue
        except Exception as e:
            logger.debug(f"Redis pattern cache read failed for {pattern_name}: {e}")

        # 3. Embed and cache
        try:
            result = await embedder.embed_scam_pattern(
                pattern_name=pattern["name"],
                description=pattern["description"],
                code_snippets=pattern.get("code_snippets", []),
                indicators=pattern.get("indicators", []),
                severity=pattern.get("severity", "high"),
            )
            _pattern_cache[pattern_name] = result

            # Store in Redis with 30-day TTL
            try:
                await r.setex(
                    redis_key,
                    86400 * 30,
                    json.dumps(
                        {
                            "vector": result.vector,
                            "dims": result.dims,
                            "model": result.model,
                            "head": result.head,
                        }
                    ),
                )
            except Exception as e:
                logger.debug(f"Redis pattern cache write failed for {pattern_name}: {e}")

            logger.info(f"Pre-embedded scam pattern: {pattern_name}")
        except Exception as e:
            logger.warning(f"Failed to pre-embed pattern {pattern_name}: {e}")

    _pattern_cache_loaded = True


# ═══════════════════════════════════════════════════════════════════
# CORE INGEST
# ═══════════════════════════════════════════════════════════════════


async def ingest_document(
    collection: str,
    content: str,
    metadata: dict | None = None,
    doc_id: str | None = None,
) -> dict:
    """
    Ingest a document with real semantic embedding.
    """
    embedder = await get_embedder()
    metadata = metadata or {}

    if not doc_id:
        doc_id = hashlib.sha256(f"{collection}:{content[:100]}:{datetime.now(UTC).isoformat()}".encode()).hexdigest()[
            :16
        ]

    # Choose embedding head based on collection
    result = await _embed_by_collection(embedder, collection, content, metadata)

    # Store
    r = await _get_redis()
    doc = {
        "id": doc_id,
        "collection": collection,
        "vector": result.vector,
        "dims": result.dims,
        "model": result.model,
        "metadata": {**metadata, "head": result.head},
        "content": content[:5000],
        "stored_at": datetime.now(UTC).isoformat(),
    }
    # TTL by collection type:
    #   Permanent (0 = no expiry): scam_patterns, contract_audits, transaction_patterns, forensic_reports
    #   Long-lived (365 days): wallet_profiles, known_scams, market_intel
    #   Short-lived (30 days): news_articles, general, token_analysis (volatile data)
    _TTL_MAP = {
        "scam_patterns": 0,
        "contract_audits": 0,
        "transaction_patterns": 0,
        "forensic_reports": 0,
        "wallet_profiles": 86400 * 365,
        "known_scams": 86400 * 365,
        "market_intel": 86400 * 365,
        "news_articles": 86400 * 30,
        "token_analysis": 86400 * 90,
        "general": 86400 * 30,
    }
    ttl = _TTL_MAP.get(collection, 86400 * 30)
    key = f"rag:{collection}:{doc_id}"
    if ttl == 0:
        await r.set(key, json.dumps(doc))  # permanent, no expiry
    else:
        await r.setex(key, ttl, json.dumps(doc))
    await r.sadd(f"rag:idx:{collection}", doc_id)

    # Bump ANN index version so it gets rebuilt on next search
    try:
        from app.ann_index import get_ann_index

        ann = get_ann_index()
        await ann.bump_version(collection)
    except Exception as e:
        logger.debug(f"ANN version bump skipped: {e}")

    # ── Update secondary indexes ──
    # BM25 index: invalidate cache so it gets rebuilt
    try:
        from app.splade_bm25 import _bm25_built_at, _bm25_index

        if _bm25_index is not None:
            # Invalidate cache so next search rebuilds
            _bm25_built_at = 0  # force rebuild on next call
            logger.debug("BM25 index cache invalidated for incremental rebuild")
    except Exception as e:
        logger.debug(f"BM25 cache invalidation skipped: {e}")

    # Entity extraction + indexing
    try:
        from app.entity_extraction import EntityLookup, extract_entities

        entities = extract_entities(content)
        if entities.has_entities:
            lookup = EntityLookup.get_instance()
            for entity in entities.all_entities():
                await lookup.index_entity(
                    entity["type"],
                    entity["value"],
                    doc_id,
                    metadata={**metadata, "collection": collection},
                )
            logger.debug(f"Entity indexing: {len(entities.all_entities())} entities")
    except Exception as e:
        logger.debug(f"Entity indexing skipped: {e}")

    # Knowledge Graph edge creation
    try:
        from app.knowledge_graph import get_knowledge_graph

        kg = await get_knowledge_graph()
        await kg.ingest_rag_document(collection, doc_id, content, metadata)
        logger.debug("KG edge creation complete")
    except Exception as e:
        logger.debug(f"KG edge creation skipped: {e}")

    # pgvector disabled - FAISS + Redis are the primary vector stores.
    # Supabase pgvector was a migration artifact consuming 500MB of free-tier quota.
    # All vector search goes through FAISS ANN → Redis hydration.
    logger.debug("pgvector upsert skipped (FAISS primary)")

    logger.info(f"Ingested {collection}/{doc_id}: {content[:60]}...")
    return {"id": doc_id, "dims": result.dims, "collection": collection}


async def _embed_by_collection(
    embedder: CryptoEmbedder,
    collection: str,
    content: str,
    metadata: dict,
) -> EmbeddingResult:
    """Route to the right embedding head based on collection type."""
    if collection == "scam_patterns" or collection == "known_scams":
        return await embedder.embed_scam_pattern(
            pattern_name=metadata.get("name", content[:50]),
            description=content,
            severity=metadata.get("severity", "medium"),
        )
    elif collection == "token_analysis":
        code = metadata.get("contract_code", "")
        return await embedder.embed_token_scam(
            name=metadata.get("name", ""),
            symbol=metadata.get("symbol", ""),
            description=content,
            contract_code=code,
            chain=metadata.get("chain", "solana"),
            metadata=metadata,
        )
    elif collection == "wallet_profiles":
        return await embedder.embed_wallet(
            address=metadata.get("address", ""),
            labels=metadata.get("labels", []),
            transactions=metadata.get("transactions", []),
            chain=metadata.get("chain", "solana"),
            balance_usd=float(metadata.get("balance_usd", 0) or 0),
        )
    elif collection == "contract_audits":
        vec = await embedder._semantic_embed_one(content, "semantic")
        return EmbeddingResult(
            vector=vec,
            dims=len(vec),
            model="local/bge-small-en-v1.5",
            head="contract_audit",
        )
    else:
        # Generic semantic
        vec = await embedder._semantic_embed_one(content, "semantic")
        return EmbeddingResult(
            vector=vec,
            dims=len(vec),
            model="openai/text-embedding-3-large",
            head="semantic",
        )


# ═══════════════════════════════════════════════════════════════════
# SEARCH - now backed by FAISS ANN index + semantic cache
# ═══════════════════════════════════════════════════════════════════


async def search_similar(
    query: str,
    collection: str = "wallet_profiles",
    limit: int = 10,
    min_similarity: float = 0.0,
) -> list[dict[str, Any]]:
    """
    Semantic search across any collection.
    Uses FAISS ANN index for fast retrieval, with semantic cache.
    """
    import time

    t0 = time.time()

    embedder = await get_embedder()
    query_vec = await embedder.embed_query(query)

    # 1. Check semantic cache
    cache_hit = False
    try:
        from app.semantic_cache import get_semantic_cache

        cache = get_semantic_cache()
        cached = await cache.check(query_vec)
        if cached is not None:
            cache_hit = True
            logger.info(f"Semantic cache hit for query: {query[:60]}")
            # Trace observability
            try:
                from app.rag_observability import trace_retrieval

                trace_retrieval((time.time() - t0) * 1000, collection, len(cached), cache_hit=True)
            except Exception:
                pass
            return cached
    except Exception as e:
        logger.debug(f"Semantic cache check skipped: {e}")

    # 2. Try FAISS ANN index search
    results = []
    try:
        from app.ann_index import get_ann_index

        ann = get_ann_index()
        if ann.is_built(collection):
            ann_results = await ann.search(
                query_embedding=query_vec,
                collection=collection,
                limit=limit,
                min_similarity=min_similarity,
            )
            if ann_results:
                # ANN search already hydrates content; skip re-enrichment for hits that have content
                needs_enrich = [r for r in ann_results if "content" not in r]
                enriched = await _enrich_ann_results(collection, needs_enrich) if needs_enrich else []
                results = [r for r in ann_results if "content" in r] + enriched
        else:
            # FAISS not built - fall back to embedder brute-force search
            logger.info(f"FAISS not built for {collection}, using brute-force fallback")
            results = await embedder.search(
                query=query,
                collection=collection,
                limit=limit,
                min_similarity=min_similarity,
            )
    except Exception as e:
        logger.warning(f"ANN/pgvector search failed, falling back to brute-force: {e}")
        # 3. Fallback to embedder brute-force search
        results = await embedder.search(
            query=query,
            collection=collection,
            limit=limit,
            min_similarity=min_similarity,
        )

    # 4. Store in semantic cache
    if results:
        try:
            from app.semantic_cache import get_semantic_cache

            cache = get_semantic_cache()
            await cache.store(query_vec, results)
        except Exception as e:
            logger.debug(f"Semantic cache store skipped: {e}")

    # Trace observability
    try:
        from app.rag_observability import trace_retrieval

        trace_retrieval((time.time() - t0) * 1000, collection, len(results), cache_hit=cache_hit)
    except Exception:
        pass

    return results


async def _enrich_ann_results(
    collection: str,
    ann_results: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    """
    Given ANN results (id + similarity), fetch full documents from Redis
    and merge the content/metadata.
    """
    if not ann_results:
        return []

    r = await _get_redis()
    pipe = r.pipeline()
    for ar in ann_results:
        pipe.get(f"rag:{collection}:{ar['id']}")

    raw_docs = await pipe.execute()
    enriched = []
    for ar, data in zip(ann_results, raw_docs, strict=False):
        if not data:
            # Return what we have even without enrichment
            enriched.append(ar)
            continue
        try:
            doc = json.loads(data)
            enriched.append(
                {
                    "id": doc.get("id", ar["id"]),
                    "similarity": ar["similarity"],
                    "content": doc.get("content", "")[:500],
                    "metadata": doc.get("metadata", {}),
                    "model": doc.get("model", ""),
                    "collection": doc.get("collection", collection),
                }
            )
        except json.JSONDecodeError:
            enriched.append(ar)

    return enriched


async def search_multi_collection(
    query: str,
    collections: list[str] | None = None,
    limit: int = 10,
    min_similarity: float = 0.5,
) -> list[dict[str, Any]]:
    """Search across multiple collections and merge results."""
    if collections is None:
        collections = COLLECTIONS

    async def _search_one(coll: str) -> list[dict[str, Any]]:
        try:
            results = await search_similar(query, collection=coll, limit=limit, min_similarity=min_similarity)
            for r in results:
                if "collection" not in r:
                    r["collection"] = coll
            return results
        except Exception as e:
            logger.warning(f"Search failed for {coll}: {e}")
            return []

    # Search all collections in parallel
    all_per_collection = await asyncio.gather(*[_search_one(c) for c in collections])
    all_results = [r for results in all_per_collection for r in results]

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


# ═══════════════════════════════════════════════════════════════════
# THREE-PILLAR HYBRID SEARCH
# ═══════════════════════════════════════════════════════════════════


async def three_pillar_search(
    query: str,
    collections: list[str] | None = None,
    limit: int = 10,
    min_similarity: float = 0.5,
    entity_boost: float = 1.5,
    use_mmr: bool = True,
    mmr_lambda: float = 0.6,
    use_kg: bool = True,
    use_parent_child: bool = True,
    use_reranker: bool = False,
) -> dict[str, Any]:
    """
    Three-pillar hybrid search with Knowledge Graph expansion, MMR dedup,
    and parent-child retrieval.

      Pillar 1 - Dense vector search via FAISS ANN index (semantic similarity)
      Pillar 2 - Sparse text search via SPLADE+BM25
      Pillar 3 - Entity exact-match + Knowledge Graph expansion
      Post-fusion - MMR deduplication + parent-child context expansion

    All pillar result sets are fused with Reciprocal Rank Fusion (k=60),
    then deduplicated with Maximal Marginal Relevance for diversity,
    and expanded with parent-child context for generation quality.

    Returns dict with merged results + pillar attribution + MMR metadata.
    """
    if collections is None:
        collections = COLLECTIONS

    # Cap limit to prevent unbounded MMR/parent-child expansion
    limit = min(limit, 100)

    embedder = await get_embedder()
    query_vec = await embedder.embed_query(query)

    # ── Pillar 1: Dense vector search (ANN) ──
    pillar1_results: list[dict[str, Any]] = []
    try:
        # Check semantic cache first
        from app.semantic_cache import get_semantic_cache

        cache = get_semantic_cache()
        cached = await cache.check(query_vec)
        if cached is not None:
            pillar1_results = cached
            for r in pillar1_results:
                r["pillar"] = "dense (cached)"
        else:
            # ANN search across collections
            for coll in collections:
                try:
                    from app.ann_index import get_ann_index

                    ann = get_ann_index()
                    ann_hits = await ann.search(
                        query_embedding=query_vec,
                        collection=coll,
                        limit=limit,
                        min_similarity=min_similarity,
                    )
                    enriched = await _enrich_ann_results(coll, ann_hits)
                    for r in enriched:
                        r["pillar"] = "dense"
                    pillar1_results.extend(enriched)
                except Exception as e:
                    logger.debug(f"ANN search for {coll} failed: {e}")
                    # Fallback to embedder search for this collection
                    try:
                        hits = await embedder.search(query, coll, limit=limit, min_similarity=min_similarity)
                        for r in hits:
                            r["pillar"] = "dense"
                        pillar1_results.extend(hits)
                    except Exception:
                        pass

            # Store in semantic cache
            if pillar1_results:
                await cache.store(query_vec, pillar1_results)
    except Exception as e:
        logger.warning(f"Pillar 1 (dense) failed: {e}")
        # Fallback: use search_multi_collection
        try:
            pillar1_results = await search_multi_collection(
                query, collections, limit=limit * 2, min_similarity=min_similarity
            )
            for r in pillar1_results:
                r["pillar"] = "dense"
        except Exception:
            pass

    # ── Pillar 2: Sparse text search (BM25 + SPLADE) ──
    pillar2_results: list[dict[str, Any]] = []
    try:
        from app.splade_bm25 import splade_search

        pillar2_results = await splade_search(query, collections, limit=limit * 2)
        for r in pillar2_results:
            r["pillar"] = "sparse"
    except ImportError:
        # Fallback to old sparse search if new module unavailable
        try:
            pillar2_results = await _sparse_text_search(query, collections, limit=limit * 2)
            for r in pillar2_results:
                r["pillar"] = "sparse"
        except Exception as e:
            logger.warning(f"Pillar 2 (sparse) failed: {e}")
    except Exception as e:
        logger.warning(f"Pillar 2 (SPLADE) failed: {e}")
        # Fallback to old sparse search
        try:
            pillar2_results = await _sparse_text_search(query, collections, limit=limit * 2)
            for r in pillar2_results:
                r["pillar"] = "sparse"
        except Exception as e2:
            logger.warning(f"Pillar 2 fallback also failed: {e2}")

    # ── Pillar 3: Entity exact-match + Knowledge Graph expansion ──
    pillar3_results: list[dict[str, Any]] = []
    entity_extraction = None
    try:
        from app.entity_extraction import EntityLookup, extract_entities

        entity_extraction = extract_entities(query)
        if entity_extraction.has_entities:
            lookup = EntityLookup.get_instance()
            for entity in entity_extraction.all_entities():
                try:
                    matches = await lookup.lookup(entity["value"])
                    for m in matches:
                        m["pillar"] = "entity"
                    pillar3_results.extend(matches)
                except Exception as e:
                    logger.debug(f"Entity lookup for {entity['value']} failed: {e}")

            # Knowledge Graph expansion: find related entities via graph traversal
            if use_kg:
                try:
                    from app.knowledge_graph import get_knowledge_graph

                    kg = await get_knowledge_graph()
                    kg_results = await kg.expand_query_entities(entity_extraction.all_entities(), max_depth=2)
                    if kg_results:
                        for r in kg_results:
                            r["pillar"] = "entity+kg"
                        pillar3_results.extend(kg_results)
                        logger.info(f"KG expansion: {len(kg_results)} related entities")
                except ImportError:
                    logger.debug("Knowledge Graph module not available, skipping KG expansion")
                except Exception as e:
                    logger.debug(f"KG expansion failed (non-critical): {e}")

    except Exception as e:
        logger.warning(f"Pillar 3 (entity) failed: {e}")

    # ── RRF fusion with query-adaptive weights ──
    # Classify query type to weight pillars appropriately
    query_type = _classify_query(query)
    # Dense (semantic) weight, Sparse (keyword) weight, Entity weight
    WEIGHTS = {
        "entity_heavy": (0.3, 0.2, 0.5),  # Address/symbol queries → entity-heavy
        "keyword_heavy": (0.3, 0.6, 0.1),  # Keyword/term queries → sparse-heavy
        "semantic_heavy": (0.7, 0.2, 0.1),  # Natural language → dense-heavy
        "balanced": (0.4, 0.4, 0.2),  # Mixed queries → balanced
    }
    w_dense, w_sparse, w_entity = WEIGHTS.get(query_type, WEIGHTS["balanced"])

    p1_ranked = sorted(pillar1_results, key=lambda x: x.get("similarity", 0), reverse=True)
    p2_ranked = sorted(
        pillar2_results,
        key=lambda x: x.get("text_score", 0) or x.get("sparse_score", 0) or x.get("bm25_score", 0),
        reverse=True,
    )
    p3_ranked = sorted(pillar3_results, key=lambda x: x.get("score", 0), reverse=True)

    entity_doc_ids = set()
    for r in pillar3_results:
        did = r.get("doc_id") or r.get("id")
        if did:
            entity_doc_ids.add(did)

    fused = _rrf_fuse(
        [p1_ranked, p2_ranked, p3_ranked],
        k=60,
        entity_boost=entity_boost,
        entity_doc_ids=entity_doc_ids,
        pillar_weights=[w_dense, w_sparse, w_entity],
    )

    # Add pillar attribution
    pillar_map: dict[str, list[str]] = {}
    for r in pillar1_results:
        did = r.get("id") or r.get("doc_id")
        if did:
            pillar_map.setdefault(did, [])
            if "dense" not in pillar_map[did]:
                pillar_map[did].append("dense")
    for r in pillar2_results:
        did = r.get("id") or r.get("doc_id")
        if did:
            pillar_map.setdefault(did, [])
            if "sparse" not in pillar_map[did]:
                pillar_map[did].append("sparse")
    for r in pillar3_results:
        did = r.get("id") or r.get("doc_id")
        if did:
            pillar_map.setdefault(did, [])
            if "entity" not in pillar_map[did]:
                pillar_map[did].append("entity")

    for r in fused:
        did = r.get("id") or r.get("doc_id")
        r["pillars"] = pillar_map.get(did, ["dense"])
        r["match_type"] = "+".join(r["pillars"])

    # ── MMR deduplication ──
    pre_mmr_count = len(fused)
    mmr_applied = False
    if use_mmr and len(fused) > limit:
        try:
            from app.mmr_dedup import mmr_dedup_results

            fused = await mmr_dedup_results(
                fused,
                query_score_field="rrf_score",
                content_field="content",
                lambda_param=mmr_lambda,
                top_k=limit * 3,  # Keep more for parent-child expansion
            )
            mmr_applied = True
            logger.info(f"MMR dedup: {pre_mmr_count} → {len(fused)} results (lambda={mmr_lambda})")
        except ImportError:
            logger.debug("MMR dedup module not available, skipping")
        except Exception as e:
            logger.debug(f"MMR dedup failed (non-critical): {e}")

    # ── Cross-encoder reranking (optional stage after MMR) ──
    reranker_applied = False
    if use_reranker and fused:
        try:
            from app.cross_encoder_reranker import get_reranker

            # Map rrf_score to similarity for the reranker
            for r in fused:
                r["similarity"] = r.get("rrf_score", 0.0)
            reranker = await get_reranker()
            pre_rerank = len(fused)
            fused = await reranker.rerank(query, fused, top_k=limit * 3)
            reranker_applied = True
            logger.info(f"Cross-encoder reranked: {pre_rerank} → {len(fused)} results")
        except ImportError:
            logger.debug("Cross-encoder module not available, skipping rerank")
        except Exception as e:
            logger.debug(f"Cross-encoder rerank failed (non-critical): {e}")

    # ── Parent-child context expansion ──
    parent_child_applied = False
    if use_parent_child and fused:
        try:
            from app.contextual_chunking import parent_child_chunk  # noqa: F401

            expanded_results = []
            for r in fused[: limit * 2]:
                content = r.get("content", "") or ""
                metadata = r.get("metadata", {}) or {}

                # If the result has a parent document reference, try to expand
                parent_id = metadata.get("parent_id") or r.get("parent_id")
                if parent_id:
                    # Fetch parent content from Redis
                    try:
                        pr = await _get_redis()
                        coll = r.get("collection", "wallet_profiles")
                        parent_data = await pr.get(f"rag:{coll}:{parent_id}")
                        if parent_data:
                            pdoc = json.loads(parent_data)
                            r["parent_content"] = pdoc.get("content", "")[:2000]
                            r["parent_id"] = parent_id
                            parent_child_applied = True
                    except Exception:
                        pass

                # If content is short, it might be a child chunk - include for generation
                if content and len(content) < 800:
                    r["chunk_type"] = "child"
                    r["retrieval_note"] = "Short chunk - consider parent context for generation"

                expanded_results.append(r)

            if expanded_results:
                fused = expanded_results
        except ImportError:
            logger.debug("Contextual chunking module not available")
        except Exception as e:
            logger.debug(f"Parent-child expansion failed (non-critical): {e}")

    # Final trim
    final_results = fused[:limit]

    # ── Confidence scoring ──
    confidence = None
    try:
        from app.confidence import score_confidence

        entity_list = list(entity_doc_ids)[:5] if entity_doc_ids else []
        confidence = score_confidence(final_results, query=query, entity_matches=entity_list)
    except ImportError:
        logger.debug("Confidence module not available")
    except Exception as e:
        logger.debug(f"Confidence scoring failed (non-critical): {e}")

    return {
        "results": final_results,
        "pillar_summary": {
            "dense_hits": len(pillar1_results),
            "sparse_hits": len(pillar2_results),
            "entity_hits": len(pillar3_results),
            "entity_doc_ids": list(entity_doc_ids),
            "pillars_used": [
                p
                for p, n in [
                    ("dense", len(pillar1_results)),
                    ("sparse", len(pillar2_results)),
                    ("entity", len(pillar3_results)),
                ]
                if n > 0
            ],
            "mmr_applied": mmr_applied,
            "mmr_lambda": mmr_lambda if mmr_applied else None,
            "pre_mmr_count": pre_mmr_count,
            "post_mmr_count": len(fused) if mmr_applied else pre_mmr_count,
            "kg_expansion": len([r for r in pillar3_results if r.get("pillar") == "entity+kg"]),
            "parent_child_applied": parent_child_applied,
            "reranker_applied": reranker_applied,
        },
        "entity_extraction": entity_extraction.to_dict() if entity_extraction else None,
        "query": query,
        "query_type": query_type,
        "fusion_weights": {"dense": w_dense, "sparse": w_sparse, "entity": w_entity},
        "collections": collections,
        "confidence": confidence,
    }


def _classify_query(query: str) -> str:
    """Classify a query into entity_heavy, keyword_heavy, semantic_heavy, or balanced."""
    q = query.lower().strip()

    # Entity-heavy: contains hex addresses, $tickers, or known crypto symbols
    if re.search(r"0x[a-f0-9]{6,}", q):
        return "entity_heavy"
    if re.search(r"\$[a-z]{2,10}", q):
        return "entity_heavy"

    # Keyword-heavy: short, term-based, no natural language flow
    words = q.split()
    if len(words) <= 3 and not any(
        w in q for w in ["what", "how", "why", "when", "explain", "describe", "compare", "analyze"]
    ):
        return "keyword_heavy"

    # Semantic-heavy: natural language questions
    if any(
        w in q
        for w in [
            "what is",
            "how does",
            "why did",
            "explain",
            "describe",
            "compare",
            "analyze",
            "tell me",
            "show me",
            "find",
        ]
    ):
        return "semantic_heavy"

    return "balanced"


def _rrf_fuse(
    ranked_lists: list[list[dict[str, Any]]],
    k: int = 60,
    entity_boost: float = 1.5,
    entity_doc_ids: set | None = None,
    pillar_weights: list[float] | None = None,
) -> list[dict[str, Any]]:
    """
    Reciprocal Rank Fusion of multiple result lists.
    Supports pillar-specific weights for query-adaptive fusion.
    """
    entity_doc_ids = entity_doc_ids or set()
    if pillar_weights is None:
        pillar_weights = [1.0] * len(ranked_lists)

    scores: dict[str, float] = {}
    doc_data: dict[str, dict[str, Any]] = {}

    for list_idx, rlist in enumerate(ranked_lists):
        weight = pillar_weights[list_idx] if list_idx < len(pillar_weights) else 1.0
        for rank, item in enumerate(rlist, start=1):
            doc_id = item.get("doc_id") or item.get("id")
            if not doc_id:
                continue
            rrf = weight / (k + rank)
            if doc_id in entity_doc_ids:
                rrf *= entity_boost
            scores[doc_id] = scores.get(doc_id, 0.0) + rrf
            if doc_id not in doc_data:
                doc_data[doc_id] = item

    sorted_ids = sorted(scores, key=lambda d: scores[d], reverse=True)
    results = []
    for doc_id in sorted_ids:
        entry = dict(doc_data[doc_id])
        entry["rrf_score"] = round(scores[doc_id], 6)
        entry["entity_match"] = doc_id in entity_doc_ids
        results.append(entry)

    return results


# ═══════════════════════════════════════════════════════════════════
# QUERY TRANSFORMATION SEARCH
# ═══════════════════════════════════════════════════════════════════


async def search_with_transform(
    query: str,
    collection: str = "wallet_profiles",
    limit: int = 10,
    strategy: str = "auto",
) -> dict[str, Any]:
    """
    Search with query transformation.
    Transforms the query first, then searches with all transformed queries,
    and merges results using RRF.
    """
    from app.query_transform import transform_query

    transformed = await transform_query(query, strategy=strategy)

    # Search with each transformed query
    all_result_lists = []
    for tq in transformed.transformed_queries:
        try:
            results = await search_similar(tq, collection=collection, limit=limit, min_similarity=0.3)
            if results:
                all_result_lists.append(results)
        except Exception as e:
            logger.debug(f"Transformed query search failed for '{tq[:60]}': {e}")

    # If only one list (or none), just return as-is
    if len(all_result_lists) <= 1:
        results = all_result_lists[0] if all_result_lists else []
        return {
            "query": query,
            "collection": collection,
            "results": results[:limit],
            "total": len(results),
            "strategy": transformed.strategy,
            "transformed_queries": transformed.transformed_queries,
            "metadata": transformed.metadata,
        }

    # RRF merge across result lists
    rrf_k = 60
    scores: dict[str, float] = {}
    doc_data: dict[str, dict[str, Any]] = {}

    for rlist in all_result_lists:
        sorted_list = sorted(rlist, key=lambda x: x.get("similarity", 0), reverse=True)
        for rank, item in enumerate(sorted_list, start=1):
            doc_id = item.get("doc_id") or item.get("id")
            if not doc_id:
                continue
            rrf = 1.0 / (rrf_k + rank)
            scores[doc_id] = scores.get(doc_id, 0.0) + rrf
            if doc_id not in doc_data:
                doc_data[doc_id] = item

    sorted_ids = sorted(scores, key=lambda d: scores[d], reverse=True)
    merged = []
    for doc_id in sorted_ids:
        entry = dict(doc_data[doc_id])
        entry["rrf_score"] = round(scores[doc_id], 6)
        merged.append(entry)

    return {
        "query": query,
        "collection": collection,
        "results": merged[:limit],
        "total": len(merged),
        "strategy": transformed.strategy,
        "transformed_queries": transformed.transformed_queries,
        "metadata": transformed.metadata,
    }


async def _sparse_text_search(
    query: str,
    collections: list[str],
    limit: int = 20,
) -> list[dict[str, Any]]:
    """
    Simple keyword-based text search over Redis-stored documents.
    Tokenizes the query and matches against content fields.
    This is a lightweight BM25-style fallback when pgvector FTS is not available.
    """
    r = await _get_redis()

    # Tokenize query
    query_tokens = set(query.lower().split())
    if not query_tokens:
        return []

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

            # Batch fetch (limit to avoid massive reads)
            sample_ids = doc_ids[:200]
            keys = [f"rag:{coll}:{did}" for did in sample_ids]
            pipe = r.pipeline()
            for k in keys:
                pipe.get(k)
            raw_docs = await pipe.execute()

            for data in raw_docs:
                if not data:
                    continue
                try:
                    doc = json.loads(data)
                except json.JSONDecodeError:
                    continue

                content = doc.get("content", "").lower()
                if not content:
                    continue

                # BM25-style scoring: count matching tokens
                content_tokens = set(content.split())
                matching = query_tokens & content_tokens
                if not matching:
                    continue

                # Simple BM25-ish score: match ratio * idf-like factor
                score = len(matching) / (len(query_tokens) + 0.5)
                # Boost exact phrase matches
                if query.lower() in content:
                    score *= 2.0

                results.append(
                    {
                        "id": doc.get("id", ""),
                        "doc_id": doc.get("id", ""),
                        "text_score": round(score, 4),
                        "content": doc.get("content", "")[:500],
                        "metadata": doc.get("metadata", {}),
                        "collection": coll,
                    }
                )
        except Exception as e:
            logger.debug(f"Sparse search for {coll} failed: {e}")

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


# ═══════════════════════════════════════════════════════════════════
# SCAM PATTERN DETECTION
# ═══════════════════════════════════════════════════════════════════


async def detect_scam_patterns(
    token_data: dict,
    threshold: float = 0.65,
) -> dict[str, Any]:
    """
    Compare a token against all known scam patterns.
    Returns matched patterns with similarity scores.
    """
    global _pattern_cache_loaded
    embedder = await get_embedder()

    # Pre-embed patterns on first call if not already done
    if not _pattern_cache_loaded:
        await _preembed_scam_patterns()

    # Build a token embedding
    code = token_data.get("contract_code", "")
    desc = token_data.get("description", "")
    name = token_data.get("name", "Unknown")
    symbol = token_data.get("symbol", "???")

    # Quick pre-filter: check code for exact keyword matches (fast, no API)
    code_lower = code.lower() if code else ""
    quick_matches = []
    for pattern in KNOWN_SCAM_PATTERNS:
        if not pattern["code_snippets"]:
            continue
        hits = sum(1 for s in pattern["code_snippets"] if s.lower() in code_lower)
        if hits > 0:
            quick_matches.append(
                {
                    "pattern": pattern["name"],
                    "severity": pattern["severity"],
                    "quick_match_score": hits / len(pattern["code_snippets"]),
                    "matched_snippets": hits,
                }
            )

    # Deep semantic comparison
    token_result = await embedder.embed_token_scam(
        name=name,
        symbol=symbol,
        description=desc,
        contract_code=code,
        chain=token_data.get("chain", "solana"),
        metadata=token_data,
    )

    # Vector layout:
    #   token  = [semantic(sem_dim) | code(128) | behavioral(64) | wallet(64)]
    #   pattern = [semantic(sem_dim) | code(128)]
    # Both have a common [semantic | code] prefix if they use the same model.
    # For the semantic comparison, use only the common prefix.
    # For the code comparison, compare the 128-dim code sections.

    token_total = len(token_result.vector)
    # If token vector has multi-head layout, extract heads by known offsets
    if token_total > 128 + 64 + 64:
        # Full multi-head: [semantic | code(128) | behavioral(64) | wallet(64)]
        token_sem_end = token_total - 128 - 64 - 64
        token_sem = token_result.vector[:token_sem_end]
        token_code = token_result.vector[token_sem_end : token_sem_end + 128]
    elif token_total > 128:
        # Partial: [semantic | code(128)] - no behavioral/wallet
        token_sem_end = token_total - 128
        token_sem = token_result.vector[:token_sem_end]
        token_code = token_result.vector[token_sem_end : token_sem_end + 128]
    else:
        # Only semantic - use as-is
        token_sem = token_result.vector
        token_code = [0.0] * 128

    deep_matches = []
    for pattern in KNOWN_SCAM_PATTERNS:
        # Use pre-embedded cache if available, otherwise embed on-demand and cache
        pattern_name = pattern["name"]
        if pattern_name in _pattern_cache:
            pattern_result = _pattern_cache[pattern_name]
        else:
            pattern_result = await embedder.embed_scam_pattern(
                pattern_name=pattern["name"],
                description=pattern["description"],
                code_snippets=pattern.get("code_snippets", []),
                indicators=pattern.get("indicators", []),
                severity=pattern.get("severity", "high"),
            )
            # Cache for future use
            _pattern_cache[pattern_name] = pattern_result
            # Also persist to Redis
            try:
                r = await _get_redis()
                redis_key = f"rag:pattern_cache:{pattern_name}"
                await r.setex(
                    redis_key,
                    86400 * 30,
                    json.dumps(
                        {
                            "vector": pattern_result.vector,
                            "dims": pattern_result.dims,
                            "model": pattern_result.model,
                            "head": pattern_result.head,
                        }
                    ),
                )
            except Exception:
                pass

        # Extract pattern semantic + code portions
        pat_total = len(pattern_result.vector)
        if pat_total > 128:
            pat_sem_end = pat_total - 128
            pattern_sem = pattern_result.vector[:pat_sem_end]
            pattern_result.vector[pat_sem_end : pat_sem_end + 128]
        else:
            pattern_sem = pattern_result.vector

        # Compare same-length semantic portions (may differ if different models)
        min_sem = min(len(token_sem), len(pattern_sem))
        if min_sem == 0:
            continue
        sem_sim = embedder.cosine_similarity(token_sem[:min_sem], pattern_sem[:min_sem])

        # Code similarity
        pat_code_features = (
            extract_contract_features("\n".join(pattern.get("code_snippets", []))).tolist()
            if pattern.get("code_snippets")
            else [0.0] * 128
        )
        code_sim = embedder.cosine_similarity(token_code, pat_code_features) if any(pat_code_features) else 0.0

        combined = 0.7 * sem_sim + 0.3 * code_sim

        if combined >= threshold:
            deep_matches.append(
                {
                    "pattern": pattern["name"],
                    "severity": pattern["severity"],
                    "description": pattern["description"],
                    "similarity": round(combined, 4),
                    "semantic_sim": round(sem_sim, 4),
                    "code_sim": round(code_sim, 4),
                }
            )

    deep_matches.sort(key=lambda x: x["similarity"], reverse=True)

    return {
        "quick_matches": quick_matches,
        "deep_matches": deep_matches,
        "highest_threat": deep_matches[0]["pattern"] if deep_matches else "none",
        "threat_severity": deep_matches[0]["severity"] if deep_matches else "low",
    }


# ═══════════════════════════════════════════════════════════════════
# SEEDING
# ═══════════════════════════════════════════════════════════════════


async def seed_known_scams() -> dict[str, Any]:
    """Seed the RAG database with known scam patterns."""
    global _seeded

    r = await _get_redis()
    existing = await r.scard("rag:idx:known_scams")
    if existing > 0 and _seeded:
        return {"status": "already_seeded", "count": existing}

    embedder = await get_embedder()
    count = 0

    for pattern in KNOWN_SCAM_PATTERNS:
        try:
            pid = hashlib.sha256(f"scam:{pattern['name']}".encode()).hexdigest()[:16]
            result = await embedder.embed_scam_pattern(
                pattern_name=pattern["name"],
                description=pattern["description"],
                code_snippets=pattern.get("code_snippets", []),
                indicators=pattern.get("indicators", []),
                severity=pattern.get("severity", "high"),
            )

            doc = {
                "id": pid,
                "collection": "known_scams",
                "vector": result.vector,
                "dims": result.dims,
                "model": result.model,
                "metadata": pattern,
                "content": pattern["description"],
                "stored_at": datetime.now(UTC).isoformat(),
            }
            key = f"rag:known_scams:{pid}"
            await r.set(key, json.dumps(doc))  # permanent - seed data should never expire
            await r.sadd("rag:idx:known_scams", pid)
            count += 1
            logger.info(f"Seeded scam pattern: {pattern['name']}")

        except Exception as e:
            logger.error(f"Failed to seed {pattern['name']}: {e}")

    _seeded = True
    return {
        "status": "seeded",
        "count": count,
        "patterns": [p["name"] for p in KNOWN_SCAM_PATTERNS],
    }


# ═══════════════════════════════════════════════════════════════════
# STATS
# ═══════════════════════════════════════════════════════════════════


async def ingest_forensic_report(
    report_text: str,
    report_name: str = "forensic_report",
    metadata: dict | None = None,
) -> dict:
    """
    Ingest a forensic report into the RAG knowledge base.
    Forensic reports are stored permanently (TTL=0) in the forensic_reports collection.
    """
    meta = {
        **(metadata or {}),
        "name": report_name,
        "type": "forensic_report",
    }
    return await ingest_document(
        collection="forensic_reports",
        content=report_text,
        metadata=meta,
    )


async def get_stats() -> dict[str, Any]:
    embedder = await get_embedder()
    r = await _get_redis()

    # Get ANN index stats FIRST - these have the authoritative vector counts
    ann_stats = {}
    try:
        from app.ann_index import get_ann_index

        ann = get_ann_index()
        ann_stats = ann.stats()
    except Exception:
        pass

    # Build collection sizes from ANN n_vectors (authoritative),
    # falling back to Redis SCARD only when ANN isn't built for that collection
    collection_sizes = {}
    for coll in COLLECTIONS:
        ann_entry = ann_stats.get(coll, {})
        if ann_entry.get("built") and ann_entry.get("n_vectors", 0) > 0:
            collection_sizes[coll] = ann_entry["n_vectors"]
        else:
            try:
                size = await r.scard(f"rag:idx:{coll}")
                collection_sizes[coll] = size
            except Exception:
                collection_sizes[coll] = 0

    # Get semantic cache stats
    cache_stats = {}
    try:
        from app.semantic_cache import get_semantic_cache

        cache = get_semantic_cache()
        cache_stats = await cache.stats()
    except Exception:
        pass

    return {
        "embedder": embedder.stats,
        "collections": collection_sizes,
        "total_docs": sum(collection_sizes.values()),
        "seeded": _seeded,
        "ann_index": ann_stats,
        "semantic_cache": cache_stats,
    }


# ═══════════════════════════════════════════════════════════════════
# GOPLUS SECURITY ANALYSIS
# ═══════════════════════════════════════════════════════════════════


async def get_goplus_analysis(token_address: str, chain: str = "1") -> dict[str, Any]:
    """
    Run GoPlus Security API token risk check and return structured results.
    Chain IDs: 1=ETH, 56=BSC, 137=Polygon, 42161=Arbitrum, 43114=Avalanche,
    10=Optimism, 8453=Base, 324=zkSync, 534352=Scroll, 59144=Linea.
    """
    from app.scam_sources import GoPlusConnector

    result = await GoPlusConnector.check_token(chain, token_address)
    if not result:
        return {"error": "GoPlus check returned no data", "address": token_address, "chain": chain}

    # Ingest the result into the RAG knowledge base for future retrieval
    try:
        content = (
            f"GoPlus Security Scan [{chain}]: {token_address} - "
            f"honeypot={result.get('is_honeypot')}, "
            f"buy_tax={result.get('buy_tax')}, sell_tax={result.get('sell_tax')}, "
            f"proxy={result.get('is_proxy')}, "
            f"blacklisted={result.get('is_blacklisted')}, "
            f"hidden_owner={result.get('hidden_owner')}, "
            f"risk_score={result.get('risk_score')}"
        )
        await ingest_document(
            collection="token_analysis",
            content=content,
            metadata={
                "source": "goplus",
                "address": token_address,
                "chain": chain,
                **result,
            },
        )
    except Exception as e:
        logger.warning(f"Failed to ingest GoPlus result for {token_address}: {e}")

    return result