#!/usr/bin/env python3
"""
LambdaRAG v2: True Fusion Architecture

All 4 retrieval paths run in parallel for EVERY query.
Classifier adjusts weights, NOT routes.
Results from all paths are fused via Weighted RRF, then reranked.

Architecture:
  λ = Classify(q) → weights
      >> Par(BM25(q), Vector(q), GraphExpand(q), WikiRead(q))
      >> WeightedRRF(results, weights)
      >> BuildContext(fused_results, relation_context, wiki_context)

The key insight: Fusion means ALL paths contribute to EVERY answer.
The classifier only controls HOW MUCH each path contributes.
"""
import json
import os
import re
import sys
import time
import urllib.request
from collections import defaultdict
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

# ── Skill imports ──
try:
    from search_engine import search as _bm25_search_fn
except ImportError:
    _bm25_search_fn = None

try:
    from search_engine_v2 import _vector_search as _vector_search_fn, _embed_query
except ImportError:
    _vector_search_fn = None
    _embed_query = None

try:
    from graph_engine import get_graph
except ImportError:
    get_graph = None

# ── Config ──
VLLM_URL = os.environ.get('VLLM_URL', 'http://127.0.0.1:8000/v1/chat/completions')
WIKI_DIR = os.environ.get('WIKI_DIR', '/app/agentexample/qaagent67wiki/wiki')

# ── Weight profiles per query type ──
# Each weight controls how much that source contributes to final ranking
WEIGHT_PROFILES = {
    'fact':      {'bm25': 1.0, 'vector': 0.5, 'graph': 0.1, 'wiki': 0.1},
    'compare':   {'bm25': 0.5, 'vector': 0.8, 'graph': 0.4, 'wiki': 0.7},
    'relation':  {'bm25': 0.5, 'vector': 0.3, 'graph': 0.7, 'wiki': 0.8},
    'synthesis': {'bm25': 0.6, 'vector': 0.5, 'graph': 0.5, 'wiki': 0.6},
    'temporal':  {'bm25': 1.0, 'vector': 0.2, 'graph': 0.05, 'wiki': 0.2},
}
DEFAULT_WEIGHTS = {'bm25': 0.7, 'vector': 0.5, 'graph': 0.3, 'wiki': 0.3}
RRF_K = 60


def _llm_call(prompt, max_tokens=200):
    body = json.dumps({
        "model": "qwen2.5-32b",
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.0, "max_tokens": max_tokens,
    }).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())
        return data["choices"][0]["message"]["content"].strip()


# ══════════════════════════════════════════
# Step 1: Classify → Weights (not routes!)
# ══════════════════════════════════════════

def classify_query(query):
    """Classify query type and return weight profile"""
    q = query.lower()

    # Rule-based fast classification
    if re.search(r'\d{4}年.*多少|数量|人数|统计|及格率|总数', q):
        qtype = 'fact'
    elif re.search(r'异同|比较|区别|对比|不同|相同', q):
        qtype = 'compare'
    elif re.search(r'关系|隶属|上级|下级|依据|谁.*管|属于|归.*管', q):
        qtype = 'relation'
    elif re.search(r'最新|修订|变化|趋势|演变|历史|何时.*施行', q):
        qtype = 'temporal'
    elif re.search(r'总结|综合|梳理|分析.*体系|框架|概述|全面', q):
        qtype = 'synthesis'
    elif len(query) < 15:
        qtype = 'fact'
    else:
        # LLM fallback
        try:
            result = _llm_call(
                f"分类(只输出一个词): fact/compare/relation/synthesis/temporal\n问题: {query}\n类型:", 10)
            qtype = 'fact'
            for t in ['fact', 'compare', 'relation', 'synthesis', 'temporal']:
                if t in result.lower():
                    qtype = t
                    break
        except:
            qtype = 'fact'

    weights = WEIGHT_PROFILES.get(qtype, DEFAULT_WEIGHTS)
    return qtype, weights


# ══════════════════════════════════════════
# Step 2: Four parallel retrieval paths
# ══════════════════════════════════════════

def path_bm25(query, top_k=10):
    """Path A: BM25 keyword search"""
    if _bm25_search_fn:
        try:
            results = _bm25_search_fn(query, top_k=top_k)
            for r in results:
                r['_path'] = 'bm25'
            return results
        except Exception as e:
            print(f"[bm25] Error: {e}")
    return []


def path_vector(query, top_k=10):
    """Path B: Embedding vector search"""
    if _vector_search_fn:
        try:
            results = _vector_search_fn(query, top_k=top_k)
            for r in results:
                r['_path'] = 'vector'
            return results
        except Exception as e:
            print(f"[vector] Error: {e}")
    return []


def path_graph(query, top_k=10):
    """Path C: Entity graph expansion → BM25+Vector search"""
    if not get_graph:
        return []
    try:
        graph = get_graph(wiki_dir=WIKI_DIR)
        entities = graph.find_entities(query, max_results=5)
        if not entities:
            return []

        expand_result = graph.expand_query(query, seed_entities=entities, max_hops=2)
        expanded_query = expand_result.get('expanded_query', query)

        # Search with expanded query
        results = []
        if _bm25_search_fn:
            bm25_r = _bm25_search_fn(expanded_query, top_k=top_k // 2)
            results.extend(bm25_r)
        if _vector_search_fn:
            vec_r = _vector_search_fn(expanded_query, top_k=top_k // 2)
            results.extend(vec_r)

        # Also add relation triples as pseudo-results
        for rel in expand_result.get('relations', [])[:5]:
            results.append({
                'source': f"relation:{rel['subject']}→{rel['object']}",
                'chunk_id': 0,
                'text': f"{rel['subject']} {rel['relation']} {rel['object']}。{rel.get('description', '')}",
                'score': 0.5,
            })

        for r in results:
            r['_path'] = 'graph'

        # Store relation context for LLM prompt enrichment
        if results:
            results[0]['_relation_context'] = expand_result.get('relation_context', '')
            results[0]['_entities'] = expand_result.get('entities', [])

        return results
    except Exception as e:
        print(f"[graph] Error: {e}")
        return []


def path_wiki(query, top_k=5):
    """Path D: Wiki compiled knowledge pages"""
    wiki_path = Path(WIKI_DIR)
    if not wiki_path.exists():
        return []

    # Find entities for better matching
    entities = []
    if get_graph:
        try:
            graph = get_graph(wiki_dir=WIKI_DIR)
            entities = graph.find_entities(query, max_results=5)
        except:
            pass

    pages = []
    for subdir in ['sources', 'entities', 'topics', 'analyses']:
        d = wiki_path / subdir
        if not d.exists():
            continue
        for f in d.glob('*.md'):
            try:
                content = f.read_text(encoding='utf-8')
                score = 0
                for e in entities:
                    if e in f.stem or e in content[:500]:
                        score += 10
                score += sum(1 for c in set(query) if c in content[:500])
                if score > len(query) * 0.2:
                    pages.append({
                        'source': f'wiki:{subdir}/{f.stem}',
                        'chunk_id': 0,
                        'text': content[:2000],
                        'score': score / 100.0,
                        '_path': 'wiki',
                    })
            except:
                pass

    pages.sort(key=lambda x: -x['score'])
    return pages[:top_k]


# ══════════════════════════════════════════
# Step 3: Weighted RRF Fusion
# ══════════════════════════════════════════

def weighted_rrf_fuse(path_results, weights, k=RRF_K):
    """
    Weighted Reciprocal Rank Fusion.
    Each path's contribution is scaled by its weight.

    score(doc) = Σ weight_i / (k + rank_i(doc))
    """
    scores = defaultdict(float)
    chunk_map = {}
    path_tags = defaultdict(set)  # Track which paths found each doc

    for path_name, results in path_results.items():
        w = weights.get(path_name, 0.5)
        if w == 0:
            continue
        for rank, r in enumerate(results):
            key = f"{r.get('source', '')}#{r.get('chunk_id', 0)}"
            scores[key] += w / (k + rank + 1)
            if key not in chunk_map:
                chunk_map[key] = r.copy()
            path_tags[key].add(path_name)

    # Sort by fused score
    sorted_keys = sorted(scores.items(), key=lambda x: -x[1])

    fused = []
    for key, score in sorted_keys:
        r = chunk_map[key]
        r['score'] = round(score, 4)
        r['_fusion_paths'] = list(path_tags[key])
        r['_fusion_score'] = round(score, 4)
        fused.append(r)

    return fused


# ══════════════════════════════════════════
# Step 4: Build enriched context for LLM
# ══════════════════════════════════════════

def build_enriched_context(fused_results, top_k=5):
    """
    Build LLM context from fused results.
    Includes: document chunks + relation context + wiki summaries.
    This is the key advantage of fusion - multiple knowledge sources in one context.
    """
    context_parts = []
    relation_context = ""
    wiki_summaries = []

    for i, r in enumerate(fused_results[:top_k]):
        source = r.get('source', '')
        text = r.get('text', '')

        if source.startswith('wiki:'):
            wiki_summaries.append(f"[Wiki: {source[5:]}]\n{text[:800]}")
        elif source.startswith('relation:'):
            # Relations go into a separate section
            if relation_context == "":
                relation_context = "## 实体关系 (来自知识图谱)\n"
            relation_context += f"- {text}\n"
        else:
            context_parts.append(f"[doc{i+1}: {source}]\n{text}")

        # Extract relation context from graph path results
        if '_relation_context' in r and r['_relation_context']:
            if relation_context == "":
                relation_context = "## 实体关系 (来自知识图谱)\n"
            relation_context += r['_relation_context'] + "\n"

    enriched = ""
    if relation_context:
        enriched += relation_context + "\n"
    if wiki_summaries:
        enriched += "## Wiki 编译知识\n" + "\n\n".join(wiki_summaries) + "\n\n"
    if context_parts:
        enriched += "## 检索文档\n" + "\n\n".join(context_parts)

    return enriched


# ══════════════════════════════════════════
# Main search function
# ══════════════════════════════════════════

def search(query, top_k=5):
    """
    True Fusion search:
    1. Classify → get weights (not route!)
    2. All 4 paths run in parallel
    3. Weighted RRF fusion
    4. Return enriched results

    Drop-in replacement for search_engine.search()
    """
    t0 = time.time()

    # Step 1: Classify and get weights
    qtype, weights = classify_query(query)

    # Step 2: Run all paths in parallel
    path_results = {}
    with ThreadPoolExecutor(max_workers=4) as pool:
        futures = {
            pool.submit(path_bm25, query, 10): 'bm25',
            pool.submit(path_vector, query, 10): 'vector',
            pool.submit(path_graph, query, 10): 'graph',
            pool.submit(path_wiki, query, 5): 'wiki',
        }
        for future in as_completed(futures):
            path_name = futures[future]
            try:
                result = future.result(timeout=15)
                path_results[path_name] = result
            except Exception as e:
                print(f"[{path_name}] Timeout/Error: {e}")
                path_results[path_name] = []

    # Step 3: Weighted RRF fusion
    fused = weighted_rrf_fuse(path_results, weights)

    # Step 4: Time filter for temporal queries
    if qtype == 'temporal':
        fused = _time_filter(fused, query)

    elapsed = time.time() - t0

    # Add metadata
    if fused:
        fused[0]['_query_type'] = qtype
        fused[0]['_weights'] = weights
        fused[0]['_strategy'] = f"fusion({'+'.join(f'{k}:{v}' for k,v in weights.items() if v > 0)})"
        fused[0]['_total_time'] = round(elapsed, 3)
        fused[0]['_path_counts'] = {k: len(v) for k, v in path_results.items()}

        # Build enriched context (stored for LLM prompt building)
        fused[0]['_enriched_context'] = build_enriched_context(fused, top_k)

    return fused[:top_k]


def _time_filter(results, query):
    """Boost temporally relevant results"""
    scored = []
    for r in results:
        boost = 0
        text = r.get('text', '') + r.get('source', '')
        if any(k in text for k in ['最新修订', '现行有效', '起施行']):
            boost += 2
        src_years = re.findall(r'(20\d{2})', r.get('source', ''))
        if src_years:
            boost += (max(int(y) for y in src_years) - 2000) * 0.1
        scored.append((r, r.get('score', 0) + boost))
    scored.sort(key=lambda x: -x[1])
    return [r for r, _ in scored]


# ══════════════════════════════════════════
# CLI test
# ══════════════════════════════════════════

if __name__ == '__main__':
    query = sys.argv[1] if len(sys.argv) > 1 else '海上交通安全法与水污染防治法在船舶管理方面有何异同？'

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

    qtype, weights = classify_query(query)
    print(f"Type: {qtype}")
    print(f"Weights: {weights}")

    print(f"\nSearching (all 4 paths parallel)...")
    results = search(query, top_k=5)
    print(f"Results: {len(results)}")
    if results:
        meta = results[0]
        print(f"Time: {meta.get('_total_time', '?')}s")
        print(f"Strategy: {meta.get('_strategy', '?')}")
        print(f"Path counts: {meta.get('_path_counts', {})}")
    for i, r in enumerate(results):
        paths = r.get('_fusion_paths', [])
        print(f"\n[{i+1}] score={r.get('score','?')} paths={paths} {r.get('source','?')[:60]}")
        print(f"    {r.get('text','')[:150]}...")