#!/usr/bin/env python3
"""
RAG v2: Hybrid search (BM25 + Embedding + RRF + Rerank)
Uses local Ollama bge-m3 for embedding (free, no API cost).
Falls back to BM25-only if embedding unavailable.

Lambda: Tool("hybrid_search", λq. Par(BM25(q), Vector(q)) >> RRF >> Rerank)
"""
import json
import math
import os
import pickle
import re
import struct
import time
import urllib.request
from collections import Counter, defaultdict
from pathlib import Path

# Import config + BM25
try:
    from config import cfg as _app_cfg
except ImportError:
    _app_cfg = None

try:
    from search_engine import search as bm25_search, _load_index as _load_bm25_index
except ImportError:
    bm25_search = None

# ── Config from config.py ──
OLLAMA_EMBED_URL = _app_cfg.embed_url if _app_cfg else os.environ.get('OLLAMA_EMBED_URL', 'http://127.0.0.1:11435/api/embed')
EMBED_MODEL = _app_cfg.embed_model if _app_cfg else os.environ.get('EMBED_MODEL', 'bge-m3')
VECTOR_SEARCH_DISABLED = os.environ.get('DISABLE_VECTOR_SEARCH', '').lower() in ('1', 'true', 'yes')
RRF_K = 60
RERANK_CANDIDATES = 20
SIMILARITY_THRESHOLD = 0.0

_vector_cache = {}


def _get_base_dir():
    if _app_cfg:
        return _app_cfg.base_dir
    return Path(os.environ.get('KNOWLEDGE_DIR', str(Path(__file__).parent)))


# ── Embedding via Ollama ──
def _embed_texts(texts, model=None):
    """Embed one or more texts via Ollama API. Returns list of vectors."""
    if model is None:
        model = EMBED_MODEL
    body = json.dumps({
        "model": model,
        "input": texts if isinstance(texts, list) else [texts],
    }).encode('utf-8')
    req = urllib.request.Request(
        OLLAMA_EMBED_URL,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST"
    )
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = json.loads(resp.read())
            return data.get('embeddings', [])
    except Exception as e:
        print(f"[embed] Error: {e}")
        return []


def _embed_query(query):
    """Embed a single query. Returns vector or None."""
    vecs = _embed_texts(query)
    return vecs[0] if vecs else None


# ── Vector Index ──
def _load_vector_index():
    """Load pre-built vector index (numpy .npy + metadata pickle)"""
    if 'vectors' in _vector_cache:
        return _vector_cache

    base = _get_base_dir()
    vec_path = base / 'rag_vectors_v2.npy'
    meta_path = base / 'rag_vectors_meta_v2.pkl'

    if not vec_path.exists() or not meta_path.exists():
        return None

    try:
        import numpy as np
        vectors = np.load(str(vec_path), mmap_mode='r')  # memory-mapped for efficiency
        with open(meta_path, 'rb') as f:
            meta = pickle.load(f)
        _vector_cache['vectors'] = vectors
        _vector_cache['meta'] = meta
        _vector_cache['dim'] = vectors.shape[1]
        return _vector_cache
    except Exception as e:
        print(f"[vector] Load error: {e}")
        return None


def _cosine_similarity_batch(query_vec, matrix):
    """Compute cosine similarity between query and all vectors in matrix"""
    try:
        import numpy as np
        query = np.array(query_vec, dtype=np.float32)
        # Normalize
        query_norm = query / (np.linalg.norm(query) + 1e-10)
        # matrix is already loaded as numpy array
        norms = np.linalg.norm(matrix, axis=1, keepdims=True) + 1e-10
        normalized = matrix / norms
        scores = normalized @ query_norm
        return scores
    except ImportError:
        # Pure Python fallback
        return _cosine_similarity_batch_python(query_vec, matrix)


def _cosine_similarity_batch_python(query_vec, matrix):
    """Pure Python fallback for cosine similarity"""
    import math
    q_norm = math.sqrt(sum(x*x for x in query_vec)) + 1e-10
    scores = []
    for row in matrix:
        dot = sum(a*b for a, b in zip(query_vec, row))
        r_norm = math.sqrt(sum(x*x for x in row)) + 1e-10
        scores.append(dot / (q_norm * r_norm))
    return scores


def _vector_search(query, top_k=20):
    """Search using embedding vectors"""
    if VECTOR_SEARCH_DISABLED:
        return []
    vidx = _load_vector_index()
    if vidx is None:
        return []

    query_vec = _embed_query(query)
    if query_vec is None:
        return []

    scores = _cosine_similarity_batch(query_vec, vidx['vectors'])
    meta = vidx['meta']

    # Get top-k indices
    try:
        import numpy as np
        top_indices = np.argsort(scores)[::-1][:top_k]
        results = []
        for idx in top_indices:
            idx = int(idx)
            score = float(scores[idx])
            if score < SIMILARITY_THRESHOLD:
                continue
            m = meta[idx]
            results.append({
                'source': m['source'],
                'chunk_id': m['chunk_id'],
                'text': m['text'],
                'score': round(score, 4),
                '_index': idx,
            })
        return results
    except ImportError:
        # Python fallback
        indexed = list(enumerate(scores))
        indexed.sort(key=lambda x: -x[1])
        results = []
        for idx, score in indexed[:top_k]:
            if score < SIMILARITY_THRESHOLD:
                continue
            m = meta[idx]
            results.append({
                'source': m['source'],
                'chunk_id': m['chunk_id'],
                'text': m['text'],
                'score': round(score, 4),
                '_index': idx,
            })
        return results


# ── RRF Fusion ──
def _rrf_fuse(bm25_results, vector_results, k=RRF_K):
    """Reciprocal Rank Fusion: combine two ranked lists"""
    scores = defaultdict(float)
    chunk_map = {}

    for rank, r in enumerate(bm25_results):
        key = f"{r['source']}#{r['chunk_id']}"
        scores[key] += 1.0 / (k + rank + 1)
        chunk_map[key] = r

    for rank, r in enumerate(vector_results):
        key = f"{r['source']}#{r['chunk_id']}"
        scores[key] += 1.0 / (k + rank + 1)
        if key not in chunk_map:
            chunk_map[key] = r

    sorted_keys = sorted(scores.items(), key=lambda x: -x[1])
    results = []
    for key, score in sorted_keys:
        r = chunk_map[key].copy()
        r['score'] = round(score, 4)
        results.append(r)

    return results


# ── Rerank via LLM (optional, uses vLLM) ──
def _rerank_llm(query, candidates, top_n=5):
    """
    Rerank candidates using vLLM local model.
    Fallback: skip rerank, return candidates as-is.
    """
    vllm_url = os.environ.get('VLLM_URL', 'http://127.0.0.1:8000/v1/chat/completions')

    if not candidates:
        return candidates

    # Build rerank prompt
    docs_text = ""
    for i, c in enumerate(candidates[:RERANK_CANDIDATES]):
        docs_text += f"\n[{i}] {c['text'][:300]}"

    prompt = f"""Given a query and candidate passages, rank them by relevance.
Output ONLY a JSON array of passage indices in order of relevance (most relevant first).
Example: [3, 0, 7, 1, 5]

Query: {query}

Passages:{docs_text}

Ranking (JSON array):"""

    try:
        body = json.dumps({
            "model": "qwen2.5-32b",
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.0,
            "max_tokens": 100,
        }).encode('utf-8')
        req = urllib.request.Request(
            vllm_url, data=body,
            headers={"Content-Type": "application/json"},
            method="POST"
        )
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = json.loads(resp.read())
            response_text = data["choices"][0]["message"]["content"].strip()

        # Parse ranking
        match = re.search(r'\[[\d,\s]+\]', response_text)
        if match:
            ranking = json.loads(match.group())
            reranked = []
            for idx in ranking[:top_n]:
                if 0 <= idx < len(candidates):
                    r = candidates[idx].copy()
                    r['score'] = round(1.0 - len(reranked) * 0.1, 2)  # Assign decreasing scores
                    reranked.append(r)
            return reranked
    except Exception as e:
        print(f"[rerank] Error: {e}, using RRF order")

    # Fallback: return top_n from RRF ordering
    return candidates[:top_n]


# ── Main Search Function ──
def search(query, top_k=5):
    """
    Hybrid search: BM25 + Vector + RRF + Rerank
    Drop-in replacement for search_engine.search()
    Returns: [{source, chunk_id, score, text}, ...]
    """
    t0 = time.time()

    # 1. BM25 search
    bm25_results = []
    if bm25_search:
        try:
            bm25_results = bm25_search(query, top_k=RERANK_CANDIDATES)
        except Exception as e:
            print(f"[bm25] Error: {e}")

    # 2. Vector search
    vector_results = _vector_search(query, top_k=RERANK_CANDIDATES)

    # 3. Decide fusion strategy
    if bm25_results and vector_results:
        # Both available → RRF fusion
        fused = _rrf_fuse(bm25_results, vector_results)
    elif vector_results:
        fused = vector_results
    else:
        # Fallback to BM25 only
        return bm25_results[:top_k] if bm25_results else []

    # 4. Rerank top candidates
    reranked = _rerank_llm(query, fused[:RERANK_CANDIDATES], top_n=top_k)

    elapsed = time.time() - t0
    # Add timing metadata to first result
    if reranked:
        reranked[0]['_search_time'] = round(elapsed, 3)

    return reranked


if __name__ == '__main__':
    import sys
    query = sys.argv[1] if len(sys.argv) > 1 else '海上交通安全法与水污染防治法'

    print(f"Query: {query}")
    print("=" * 60)

    t0 = time.time()
    results = search(query, top_k=5)
    t1 = time.time()
    print(f"Search time: {(t1-t0)*1000:.0f}ms")
    print(f"Results: {len(results)}")
    for i, r in enumerate(results):
        print(f"\n[{i+1}] score={r['score']} {r['source'][:60]}")
        print(f"    {r['text'][:150]}...")