#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@Time    :   2024/03/07 16:17:27
@Author  :   ChenHao
@Description  :   
@Contact :   jerrychen1990@gmail.com
'''
import os
from typing import List
import requests
from requests.auth import HTTPBasicAuth
from loguru import logger
from xagents.kb.common import DistanceStrategy


class XClient:
    def __init__(self, host: str, username: str, password: str, port=8001):
        self.host = host
        self.port = port
        self.auth = HTTPBasicAuth(username, password)

    def call_service(self, path: str, method='POST', **kwargs) -> dict:
        try:
            url = f" http://{self.host}:{self.port}/{path}"
            if method == "POST":
                resp = requests.post(url=url,  auth=self.auth, **kwargs)
            else:
                resp = requests.get(url=url, auth=self.auth, **kwargs, )
            logger.debug(f"request to {url} with {method=}, {kwargs}")
            resp.raise_for_status()
            resp = resp.json()
            code = resp["code"]
            msg = resp["msg"]
            resp_data = resp["data"]
            logger.debug(f"get resp:{code=}, {msg=}, {resp_data=}")
            return resp_data

        except Exception as e:
            logger.error(f"call service {url}failed\n error:{e}\nresp:{resp.text}")
            raise e

    def health(self):
        return self.call_service(path="health", method="GET")

    def get_job(self, job_id: str):
        return self.call_service(path=f"job", params=dict(job_id=job_id), method="GET")

    def list_kb(self):
        return self.call_service(path="kb/list", method="GET")

    def get_kb(self, kb_name: str):
        return self.call_service(path=f"kb", params=dict(kb_name=kb_name), method="GET")

    def create_kb(self, kb_name: str, description: str = None, vecstore_config: dict = dict(vs_cls='XFAISS'),
                  embedding_config: dict = dict(model_cls="ZhipuEmbedding", batch_size=16, norm=True),
                  distance_strategy=DistanceStrategy.MAX_INNER_PRODUCT):
        json_data = dict(kb_name=kb_name, description=description, distance_strategy=distance_strategy,
                         vecstore_config=vecstore_config, embedding_config=embedding_config)
        return self.call_service(path="kb/create", json=json_data)

    def update_kb(self, kb_name: str, kb_config: dict):
        json_data = dict(kb_name=kb_name, kb_config=kb_config)
        return self.call_service(path="kb/update", json=json_data)

    def delete_kb(self, kb_name: str):
        data = dict(kb_name=kb_name)
        return self.call_service(path="kb/delete", data=data)

    def build_index(self, kb_name: str):
        data = dict(kb_name=kb_name)
        return self.call_service(path="kb/build_index", data=data)

    def list_kb_files(self, kb_name: str):
        data = dict(kb_name=kb_name)
        return self.call_service(path="kb_file/list", data=data, method="GET")

    def create_kb_file(self, kb_name: str, file_path: str):
        file_name = os.path.basename(file_path)
        with open(file_path, "r") as f:
            files = {'file': (file_path, f)}
            data = dict(kb_name=kb_name, file_name=file_name)
            return self.call_service(path="kb_file/create", data=data, files=files)

    def delete_kb_file(self, kb_name: str, file_name: str):
        data = dict(kb_name=kb_name, file_name=file_name)
        return self.call_service(path="kb_file/delete", data=data)

    def get_kb_file_chunks(self, kb_name: str, file_name: str):
        data = dict(kb_name=kb_name, file_name=file_name)
        return self.call_service(path="kb_file/chunks", data=data, method="GET")

    def cut_kb_file(self,
                    kb_name: str,
                    file_name: str,
                    splitter_cls: str = "BaseSplitter",
                    splitter: str = "\n",
                    max_len: int = 200,
                    min_len: int = 10,
                    parse_table: bool = False):
        data = dict(kb_name=kb_name, file_name=file_name, splitter_cls=splitter_cls,
                    splitter=splitter, max_len=max_len, min_len=min_len, parse_table=parse_table)
        return self.call_service(path="kb_file/cut", data=data)

    def search_kb(self,
                  kb_name: str,
                  query: str,
                  file_names: List[str] = None,
                  top_k: int = 3,
                  score_threshold: float = 0.0,
                  pre_content_len: int = 0,
                  next_content_len: int = 0
                  ):
        data = dict(kb_name=kb_name, query=query, file_names=file_names, top_k=top_k, score_threshold=score_threshold,
                    pre_content_len=pre_content_len, next_content_len=next_content_len)
        return self.call_service(path="kb/search", data=data)
