"""
RAG v10专用路由 - 使用苦荞SPL附表增强版索引
"""

import os
import sys
import json
import pickle
import numpy as np
from pathlib import Path
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from typing import List, Optional
import requests
from src.rag.project_config import PROJECT_ROOT

sys.path.insert(0, str(PROJECT_ROOT))

router = APIRouter(prefix="/v10")

# ========== 配置 ==========
OLLAMA_URL = "http://localhost:11434"
EMBEDDING_MODEL = "bge-m3:567m"
LLM_MODEL = "qwen3:0.6b"

# v10索引路径
V10_INDEX_PATH = PROJECT_ROOT / "paper" / "v10" / "rag_v10_faiss_final.index"
V10_CHUNKS_PATH = PROJECT_ROOT / "paper" / "v10" / "rag_v10_chunks_with_summary.json"
V10_ENHANCED_PATH = PROJECT_ROOT / "paper" / "v10" / "rag_v10_chunks_enhanced.json"

# 全局变量
_v10_index = None
_v10_chunks = None

def get_embedding(text):
    """获取文本向量"""
    response = requests.post(
        f"{OLLAMA_URL}/api/embeddings",
        json={"model": EMBEDDING_MODEL, "prompt": text}
    )
    vec = np.array(response.json()['embedding'], dtype=np.float32)
    return vec / np.linalg.norm(vec)

def load_v10_index():
    """加载v10索引"""
    global _v10_index, _v10_chunks
    
    if _v10_index is not None:
        return _v10_index, _v10_chunks
    
    print(f"[V10 RAG] 加载v10索引...")
    
    # 加载FAISS索引
    import faiss
    _v10_index = faiss.read_index(str(V10_INDEX_PATH))
    print(f"[V10 RAG] 索引加载完成: {_v10_index.ntotal} 个向量")
    
    # 加载文本块
    with open(V10_CHUNKS_PATH, 'r') as f:
        chunks1 = json.load(f)
    with open(V10_ENHANCED_PATH, 'r') as f:
        chunks2 = json.load(f)
    
    _v10_chunks = chunks1 + chunks2
    print(f"[V10 RAG] 文本块加载完成: {len(_v10_chunks)} 个")
    
    return _v10_index, _v10_chunks

# ========== 请求/响应模型 ==========
class V10RAGQueryRequest(BaseModel):
    query: str
    top_k: int = 3
    use_llm: bool = True

class V10RAGSource(BaseModel):
    chunk_id: str
    chunk_type: str
    content: str
    score: float

class V10RAGResponse(BaseModel):
    query: str
    answer: str
    sources: List[V10RAGSource]
    total_chunks: int

# ========== API端点 ==========
@router.post("/chat", response_model=V10RAGResponse)
async def v10_rag_chat(request: V10RAGQueryRequest):
    """V10专用RAG查询"""
    try:
        # 加载索引
        index, chunks = load_v10_index()
        
        # 1. 查询向量化
        query_vec = get_embedding(request.query).reshape(1, -1)
        
        # 2. FAISS检索
        distances, indices = index.search(query_vec, request.top_k)
        
        # 3. 构建sources
        sources = []
        context_parts = []
        for i, (dist, idx) in enumerate(zip(distances[0], indices[0])):
            chunk = chunks[idx]
            sources.append(V10RAGSource(
                chunk_id=chunk['id'],
                chunk_type=chunk['metadata']['type'],
                content=chunk['text'],
                score=float(dist)
            ))
            context_parts.append(f"[{i+1}] {chunk['text']}")
        
        # 4. LLM生成（可选）
        if request.use_llm:
            context = "\n\n".join(context_parts)
            prompt = f"""你是基因家族分析助手，基于以下资料回答问题。
请准确回答，并标注引用来源（如[1]、[2]等）。

参考资料：
{context}

用户问题：{request.query}

请基于以上资料回答："""
            
            response = requests.post(
                f"{OLLAMA_URL}/api/generate",
                json={
                    "model": LLM_MODEL,
                    "prompt": prompt,
                    "stream": False,
                    "options": {"temperature": 0.3}
                }
            )
            answer = response.json()['response']
        else:
            # 直接返回检索结果
            answer = "\n\n".join([f"[{i+1}] {s.content}" for i, s in enumerate(sources)])
        
        return V10RAGResponse(
            query=request.query,
            answer=answer,
            sources=sources,
            total_chunks=len(chunks)
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"V10 RAG查询失败: {str(e)}")

@router.get("/info")
async def v10_info():
    """获取v10索引信息"""
    try:
        index, chunks = load_v10_index()
        
        # 统计
        type_counts = {}
        for chunk in chunks:
            t = chunk['metadata']['type']
            type_counts[t] = type_counts.get(t, 0) + 1
        
        return {
            "index_type": "FAISS v10专用",
            "total_chunks": len(chunks),
            "embedding_model": EMBEDDING_MODEL,
            "llm_model": LLM_MODEL,
            "chunk_types": type_counts,
            "data_source": "苦荞SPL附表增强版"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

print("[V10 RAG] 路由加载完成")