#!/usr/bin/env python3
"""
RAG集成模块 - 将gene-family-mas-v2的RAG功能集成到bbav5
"""

import os
import sys
from pathlib import Path
from typing import Optional, List, Dict, Any
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

# 添加gene-family-mas-v2路径
GFMAS_ROOT = Path("/home/xia-lab/gene-family-mas-v2")
sys.path.insert(0, str(GFMAS_ROOT))

# 导入RAG组件
try:
    from src.rag.base import SearchResult
    from src.rag.vector_store import InMemoryStore
    from src.rag.embeddings import BGEEmbedding
    from src.rag.agentic.orchestrator import RAGAgentOrchestrator
    RAG_AVAILABLE = True
    print("✅ RAG组件导入成功")
except ImportError as e:
    RAG_AVAILABLE = False
    print(f"⚠️ RAG组件导入失败: {e}")
    SearchResult = None
    InMemoryStore = None
    BGEEmbedding = None
    RAGAgentOrchestrator = None

# 创建路由
router = APIRouter(prefix="/rag", tags=["RAG"])

# 全局变量
_vector_store = None
_embedding_model = None
_orchestrator = None


class RAGQueryRequest(BaseModel):
    """RAG查询请求"""
    query: str
    mode: str = "agentic"  # agentic, hybrid, exact, substring, bm25
    top_k: int = 10
    max_results: int = 3  # 显示的最大结果数
    use_llm: bool = False
    enable_self_rag: bool = True
    enable_multi_hop: bool = True


class RAGSource(BaseModel):
    """RAG来源"""
    doc_id: str
    content: str
    score: float
    metadata: dict


class RAGQueryResponse(BaseModel):
    """RAG查询响应"""
    response: str
    sources: List[RAGSource]
    mode: str
    total_chunks: int
    self_rag_used: bool = False
    multi_hop_used: bool = False


def get_vector_store():
    """获取或初始化向量存储"""
    global _vector_store, _embedding_model
    
    if _vector_store is not None:
        return _vector_store
    
    if not RAG_AVAILABLE:
        raise HTTPException(status_code=500, detail="RAG组件不可用")
    
    # 加载索引
    index_path = GFMAS_ROOT / "data" / "rag_index_final" / "vector_index.pkl"
    
    if not index_path.exists():
        raise HTTPException(status_code=500, detail=f"RAG索引不存在: {index_path}")
    
    # 创建embedding model
    _embedding_model = BGEEmbedding()
    
    # 创建store
    _vector_store = InMemoryStore(embedding_model=_embedding_model)
    
    # 加载索引
    import pickle
    with open(index_path, 'rb') as f:
        data = pickle.load(f)
    
    _vector_store._chunks = data['chunks']
    _vector_store._embeddings = data['embeddings']
    
    print(f"✅ RAG索引加载完成: {len(_vector_store._chunks)} chunks")
    
    return _vector_store


def get_orchestrator():
    """获取或初始化Agent编排器"""
    global _orchestrator
    
    if _orchestrator is not None:
        return _orchestrator
    
    if not RAG_AVAILABLE:
        raise HTTPException(status_code=500, detail="RAG组件不可用")
    
    store = get_vector_store()
    
    # 初始化编排器
    _orchestrator = RAGAgentOrchestrator(
        vector_store=store,
        retriever=store,
        enable_self_rag=True,
        enable_multi_hop=True,
        multi_hop_threshold=0.3,
        max_iterations=3
    )
    
    print("✅ RAG Agent编排器初始化完成")
    
    return _orchestrator


@router.post("/chat", response_model=RAGQueryResponse)
async def rag_chat(request: RAGQueryRequest):
    """
    RAG智能查询端点
    
    支持Self-RAG和MultiHopRAG的Agentic检索
    """
    try:
        if not RAG_AVAILABLE:
            raise HTTPException(status_code=500, detail="RAG组件未安装")
        
        store = get_vector_store()
        orchestrator = get_orchestrator()
        
        # 执行Agentic检索
        if request.mode == "agentic":
            # 使用同步fallback方法避免LangGraph异步问题
            result = await orchestrator._run_fallback(request.query, max_results=request.max_results)
            
            # 格式化结果 - sources简化处理
            sources = []
            
            return RAGQueryResponse(
                response=result.get("answer", result.get("response", "未找到答案")),
                sources=sources,
                mode=request.mode,
                total_chunks=len(store._chunks),
                self_rag_used=result.get("reflection") is not None,
                multi_hop_used=result.get("multi_hop_used", False)
            )
        else:
            # 其他模式使用简单检索
            # TODO: 实现其他检索模式
            return RAGQueryResponse(
                response=f"模式 '{request.mode}' 暂未实现，请使用 'agentic' 模式",
                sources=[],
                mode=request.mode,
                total_chunks=len(store._chunks),
                self_rag_used=False,
                multi_hop_used=False
            )
            
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"RAG查询失败: {str(e)}")


@router.get("/status")
async def rag_status():
    """RAG系统状态"""
    return {
        "available": RAG_AVAILABLE,
        "vector_store_loaded": _vector_store is not None,
        "orchestrator_initialized": _orchestrator is not None,
        "chunks_count": len(_vector_store._chunks) if _vector_store else 0
    }


# 便捷函数
def init_rag(app):
    """初始化RAG集成"""
    if RAG_AVAILABLE:
        app.include_router(router)
        print("✅ RAG路由已挂载到 /rag")
    else:
        print("⚠️ RAG不可用，跳过挂载")