import json
from typing import Literal
from urllib.parse import urlencode

from pydantic import BaseModel

from openkosmos_ai.common.model import SearchResult
from openkosmos_ai.nodes.api import ApiConfig, ApiNode


class PageParams(BaseModel):
    page: int = 1
    page_size: int = 9999
    orderby: Literal["create_time", "update_time"] = None
    desc: bool = None

    def to_http_params(self, params: dict = {}):
        return self.model_dump(exclude_none=True) | params


class RAGFlowClientConfig(ApiConfig):
    api_prefix: str = "/api/v1"


class RAGFlowClientNode(ApiNode[RAGFlowClientConfig]):
    def __init__(self, config: RAGFlowClientConfig):
        super().__init__(config)

    def create_dataset(self, name: str):
        return self.post("/datasets", json_data={
            "name": name
        })

    def list_datasets(self, dataset_id: str = None, dataset_name=None,
                      page_params: PageParams = PageParams()):
        return self.get("/datasets",
                        params=page_params.to_http_params({
                            "id": dataset_id,
                            "name": dataset_name
                        }))

    def delete_datasets(self, ids: list[str]):
        return self.delete(f"/datasets", json_data={
            "ids": ids
        })

    def upload_documents(self, dataset_id: str, files):
        return self.post(f"/datasets/{dataset_id}/documents", files=files)

    def parse_documents(self, dataset_id: str, document_ids: list[str]):
        return self.post(f"/datasets/{dataset_id}/chunks", json_data={
            "document_ids": document_ids
        })

    def list_chunks(self, dataset_id: str, document_id: str, page_params: PageParams = PageParams()):
        return self.get_api(f"/datasets/{dataset_id}/documents/{document_id}/chunks",
                            page_params=page_params,
                            params=page_params)

    def list_documents(self, dataset_id: str, document_id: str = None, document_name: str = None,
                       metadata_condition: dict = None, page_params: PageParams = PageParams()):
        return self.get(f"/datasets/{dataset_id}/documents",
                        params=page_params.to_http_params({
                            "id": document_id,
                            "name": document_name,
                            "metadata_condition": None if metadata_condition is None else urlencode(
                                json.dumps(metadata_condition))
                        }))

    def delete_document(self, dataset_id: str, ids: list[str] = None):
        return self.delete(f"/datasets/{dataset_id}/documents", json_data={
            "ids": ids
        })

    def list_chunks(self, dataset_id: str, document_id: str, page_params: PageParams = PageParams()):
        return self.get(f"/datasets/{dataset_id}/documents/{document_id}/chunks",
                        params=page_params)

    def download_document(self, dataset_id: str, document_id: str):
        return self.get(f"/datasets/{dataset_id}/documents/{document_id}", return_json=False).content

    def retrieve_chunks(self, question: str, dataset_ids: list[str] = None, document_ids: list[str] = None,
                        similarity_threshold=0.6, vector_similarity_weight=1.0,
                        metadata_condition: dict = None, top_k: int = 3):
        body = {
            "question": question,
            "top_k": top_k * 6,
            "page": 1,
            "page_size": top_k,
            "vector_similarity_weight": vector_similarity_weight,
            "similarity_threshold": similarity_threshold
        }
        if metadata_condition is not None:
            body["metadata_condition"] = metadata_condition
        if dataset_ids is not None:
            body["dataset_ids"] = dataset_ids
        if document_ids is not None:
            body["document_ids"] = document_ids
        return self.post("/retrieval", json_data=body)

    def search(self, query: str, dataset_ids: list[str] = None, document_ids: list[str] = None, threshold=0.6,
               vector_similarity_weight=1.0, filter_meta: dict = None, top_k: int = 3) -> list[SearchResult]:

        retrieve_result = self.retrieve_chunks(question=query,
                                               dataset_ids=dataset_ids,
                                               document_ids=document_ids,
                                               metadata_condition=filter_meta,
                                               similarity_threshold=threshold,
                                               vector_similarity_weight=vector_similarity_weight,
                                               top_k=top_k)

        return [SearchResult(id=chunk["id"], content=chunk["content"], similarity_score=chunk["similarity"],
                             metadata={
                                 "document_id": chunk["document_id"],
                                 "document_keyword": chunk["document_keyword"],
                                 # "kb_id": chunk["kb_id"],
                                 "similarity": chunk["similarity"],
                                 "term_similarity": chunk["term_similarity"],
                                 "vector_similarity": chunk["vector_similarity"],
                             }) for chunk in
                retrieve_result["data"]["chunks"]]
