import aiohttp

from huibiao_framework.client.data_model.model_resp_vo import ModelBaseRespVo
from huibiao_framework.client.data_model.qa_search_new import (
    SearchQaDto,
    SearchQaVo,
    ReadQaDto,
    AddDocVo,
    AddDocDto,
    ReadQaVo,
    DeleteDbDto,
    DeleteDbVo,
)
from huibiao_framework.config.config import ClientConfig


class KnowledgeQaSearchClientNew:
    def __init__(self, session: aiohttp.ClientSession):
        self.base_url = ClientConfig.SEARCH_URL
        self.headers = {"Content-Type": "application/json"}
        self.session = session

    async def add_qa(
        self,
        request_body: AddDocDto,
    ) -> AddDocVo:
        """
        异步添加QA对到知识库
        """
        url = f"{self.base_url}/data/add_doc"

        async with self.session.post(
            url, json=request_body.model_dump(), headers=self.headers
        ) as response:
            response.raise_for_status()
            raw_response = await response.json()
            return ModelBaseRespVo[AddDocVo](**raw_response).result

    async def delete_db(
        self,
        request_body: DeleteDbDto,
    ) -> DeleteDbVo:
        """
        异步添加QA对到知识库
        """
        url = f"{self.base_url}/data/delete_db"

        async with self.session.post(
            url, json=request_body.model_dump(), headers=self.headers
        ) as response:
            response.raise_for_status()
            raw_response = await response.json()
            return ModelBaseRespVo[DeleteDbVo](**raw_response).result

    async def search_qa(self, request_body: SearchQaDto) -> SearchQaVo:
        """
        异步搜索QA对
        """
        url = f"{self.base_url}/data/search"

        async with self.session.post(
            url, json=request_body.model_dump(exclude_none=True), headers=self.headers
        ) as response:
            response.raise_for_status()
            raw_response = await response.json()
            return ModelBaseRespVo[SearchQaVo](**raw_response).result


    async def read_qa(self, request_body: ReadQaDto) -> ReadQaVo:
        """
        异步搜索QA对
        """
        url = f"{self.base_url}/data/chunks/query_by_doc"

        async with self.session.post(
            url, json=request_body.model_dump(exclude_none=True), headers=self.headers
        ) as response:
            response.raise_for_status()
            raw_response = await response.json()
            return ModelBaseRespVo[ReadQaVo](**raw_response).result
