from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, validator
from app.services.document_service import DocumentService
from app.services.vector_store_service import VectorStoreService
from app.chains.rag_chain import RAGChain
from app.core.logger import logger
import os
from pathlib import Path

router = APIRouter(prefix="/rag", tags=["RAG"])

# 初始化服务
document_service = DocumentService()
vector_store_service = VectorStoreService()
rag_chain = RAGChain(vector_store_service)

class DirectoryRequest(BaseModel):
    directory: str

    @validator('directory')
    def validate_directory(cls, v):
        # 统一路径分隔符
        path = Path(v)
        if not path.exists():
            raise ValueError(f"路径不存在: {v}")
        if not path.is_file() and not path.is_dir():
            raise ValueError(f"无效的路径: {v}")
        return str(path.absolute())

class QuestionRequest(BaseModel):
    question: str

@router.post("/upload")
async def upload_documents(request: DirectoryRequest):
    """上传并处理文档"""
    try:
        logger.info(f"开始处理文档: {request.directory}")
        # 处理文档
        documents = document_service.process_documents(request.directory)
        if not documents:
            return {"message": "没有新的文档需要处理"}
        # 添加到向量存储
        vector_store_service.add_documents(documents)
        logger.info(f"成功处理 {len(documents)} 个文档块")
        return {"message": f"成功处理 {len(documents)} 个文档块"}
    except Exception as e:
        logger.error(f"处理文档时发生错误: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/query")
async def query(request: QuestionRequest):
    """查询问题"""
    try:
        result = rag_chain.query(request.question)
        return {
            "answer": result["result"],
            "sources": [doc.page_content for doc in result["source_documents"]]
        }
    except Exception as e:
        logger.error(f"查询问题时发生错误: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/clear")
async def clear_vector_store():
    """清空向量存储"""
    try:
        vector_store_service.clear()
        document_service.clear_processed_files()  # 同时清空文件处理记录
        return {"message": "向量存储和文件处理记录已清空"}
    except Exception as e:
        logger.error(f"清空向量存储时发生错误: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))