from sqlalchemy.ext.asyncio import AsyncSession
from utils.ragflow_asy_util import AsyncRAGFlowClient
from module_admin.entity.vo.ragflow_vo import RagflowListQueryModel, ListDocumentsQueryModel, UpdateFileModel, DeleteFileModel
from config.env import RAGFlowConfig
from typing import List

class RAGFlowService:
    """ 
        RAGFlow服务
    """
    
    # 获取数据集列表
    @classmethod
    async def get_ragflow_dataset_list_services(cls, query_db: AsyncSession, rage_flow_query: RagflowListQueryModel):
        """
        获取数据集列表
        """
        
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.list_datasets(**(rage_flow_query.model_dump()))
        
        # 获取分页数据
        return result
    
    # 获取数据集中文档列表
    @classmethod
    async def list_documents_services(
        cls,
        query_db: AsyncSession,
        dataset_id: str,
        list_documents_query: ListDocumentsQueryModel,
    ):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.list_documents(dataset_id=dataset_id, **(list_documents_query.model_dump()))
             
        return result
    
    # 上传文档到数据集
    @classmethod
    async def upload_file_dataset_services(
        cls,
        dataset_id: str,
        files,
    ):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.upload_documents_bytes(dataset_id=dataset_id, file_bytes=files)
             
        return result
    
    # 更新文档内容
    @classmethod
    async def update_file_dataset_services(
        cls,
        dataset_id: str,
        document_id: str,
        update_params: UpdateFileModel,
    ):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.update_document(dataset_id=dataset_id, document_id=document_id, **(update_params.model_dump()))
             
        return result
    
    # 删除文档
    @classmethod
    async def delete_file_services(
        cls,
        dataset_id: str,
        delete_params: DeleteFileModel,
    ):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.delete_documents(dataset_id=dataset_id, **(delete_params.model_dump()))
        return result 
    
    # 删除数据集
    @classmethod
    async def delete_datasets_services(
        cls,
        delete_params: DeleteFileModel,
    ):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.delete_datasets(**(delete_params.model_dump()))
        return result