from typing import List, Optional, Dict, Any, Union

import httpx

from scale_egp.sdk.enums import EmbeddingModelName, ChunkUploadStatus
from scale_egp.sdk.models import (
    KnowledgeBase,
    KnowledgeBaseRequest,
    EmbeddingConfig,
    Chunk,
    KnowledgeBaseQueryRequest,
    KnowledgeBaseUpload,
    LocalChunksSourceConfig,
    ChunkToUpload,
    KnowledgeBaseArtifact,
    S3DataSourceConfig,
    GoogleDriveDataSourceConfig,
    SharePointDataSourceConfig,
    SharePointDataSourceAuthConfig,
    CharacterChunkingStrategyConfig,
    KnowledgeBaseResponse,
    ListKnowledgeBasesResponse,
    ListKnowledgeBaseArtifactsResponse,
    KnowledgeBaseQueryResponse,
    KnowledgeBaseRemoteUploadRequest,
    KnowledgeBaseUploadResponse,
    KnowledgeBaseLocalChunkUploadRequest,
    ListKnowledgeBaseUploadsResponse,
    CancelKnowledgeBaseUploadResponse,
)
from scale_egp.utils.api_utils import APIEngine


class KnowledgeBaseCollection(APIEngine):

    _sub_path = "v2/knowledge-bases"

    def uploads(self) -> "KnowledgeBaseUploadsCollection":
        """
        Returns a KnowledgeBaseUploadsCollection object for uploads associated with a knowledge
        base.

        Returns:
            A KnowledgeBaseUploadsCollection object.
        """
        return KnowledgeBaseUploadsCollection(self._api_client)

    def artifacts(self) -> "KnowledgeBaseArtifactsCollection":
        """
        Returns a KnowledgeBaseArtifactsCollection object for artifacts associated with a
        knowledge base.

        Returns:
            A KnowledgeBaseArtifactsCollection object.
        """
        return KnowledgeBaseArtifactsCollection(self._api_client)

    def create(
        self,
        name: str,
        embedding_model_name: EmbeddingModelName,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> KnowledgeBase:
        """
        Create a new Knowledge Base.

        Args:
            name: The name of the Knowledge Base.
            embedding_model_name: The name of the embedding model to use for the Knowledge Base.
            metadata: The metadata of the Knowledge Base.

        Returns:
            The newly created Knowledge Base.
        """
        response = self._post(
            sub_path=self._sub_path,
            request=KnowledgeBaseRequest(
                account_id=self._api_client.account_id,
                knowledge_base_name=name,
                embedding_config=EmbeddingConfig(
                    embedding_model=embedding_model_name,
                ),
                metadata=metadata,
            ),
        )
        response_model = KnowledgeBaseResponse.from_dict(response.json())
        return self.get(id=response_model.knowledge_base_id)

    def get(
        self,
        id: str,
    ) -> KnowledgeBase:
        """
        Get an Knowledge Base by ID.

        Args:
            id: The ID of the Knowledge Base.

        Returns:
            The Knowledge Base.
        """
        response = self._get(
            sub_path=f"{self._sub_path}/{id}",
        )
        return KnowledgeBase.from_dict(response.json())

    def delete(
        self,
        id: str,
    ) -> bool:
        """
        Delete a Knowledge Base by ID.

        Args:
            id: The ID of the Knowledge Base.
        """
        response = self._delete(
            sub_path=f"{self._sub_path}/{id}",
        )
        return response.status_code == httpx.codes.ok

    def list(
        self,
    ) -> List[KnowledgeBase]:
        """
        List all Knowledge Bases.

        Returns:
            A list of Knowledge Bases.
        """
        response = self._get(
            sub_path=self._sub_path,
        )
        response_model = ListKnowledgeBasesResponse.from_dict(response.json())
        return response_model.items

    def query(
        self,
        knowledge_base: KnowledgeBase,
        query: str,
        top_k: int,
        include_embeddings: bool = False,
    ) -> List[Chunk]:
        """
        Query a Knowledge Base.

        Args:
            knowledge_base: The Knowledge Base to query.
            query: The query string.
            top_k: The number of results to return.
            include_embeddings: Whether to include embeddings in the response.

        Returns:
            The query response.
        """
        response = self._post(
            sub_path=f"{self._sub_path}/{knowledge_base.knowledge_base_id}/query",
            request=KnowledgeBaseQueryRequest(
                query=query,
                top_k=top_k,
                include_embeddings=include_embeddings,
            ),
        )
        response_model = KnowledgeBaseQueryResponse.from_dict(response.json())
        return response_model.chunks


class KnowledgeBaseArtifactsCollection(APIEngine):

    _sub_path = "v2/knowledge-bases/{knowledge_base_id}/artifacts"

    def get(
        self,
        id: str,
        knowledge_base: KnowledgeBase,
        status_filter: Optional[ChunkUploadStatus] = ChunkUploadStatus.COMPLETED.value,
    ) -> KnowledgeBaseArtifact:
        """
        Get a Knowledge Base Artifact by ID.

        Args:
            id: The ID of the Knowledge Base Artifact.
            knowledge_base: The Knowledge Base the artifact was created for.
            status_filter: Return only artifacts with the given status.

        Returns:
            The Knowledge Base Artifact.
        """
        response = self._get(
            sub_path=f"{self._sub_path.format(knowledge_base_id=knowledge_base.knowledge_base_id)}/"
            f"{id}",
            query_params=dict(
                status_filter=status_filter,
            ),
        )
        return KnowledgeBaseArtifact.from_dict(response.json())

    def list(
        self,
        knowledge_base: KnowledgeBase,
    ) -> List[KnowledgeBaseArtifact]:
        """
        List all Knowledge Base Artifacts.

        Returns:
            A list of Knowledge Base Artifacts.
        """
        response = self._get(
            sub_path=self._sub_path.format(knowledge_base_id=knowledge_base.knowledge_base_id),
        )
        response_model = ListKnowledgeBaseArtifactsResponse.from_dict(response.json())
        return response_model.artifacts


class KnowledgeBaseUploadsCollection(APIEngine):
    _sub_path = "v2/knowledge-bases/{knowledge_base_id}/uploads"

    def create_remote_upload(
        self,
        knowledge_base: KnowledgeBase,
        data_source_config: Union[
            S3DataSourceConfig, SharePointDataSourceConfig, GoogleDriveDataSourceConfig
        ],
        data_source_auth_config: Optional[Union[SharePointDataSourceAuthConfig]],
        chunking_strategy_config: Union[CharacterChunkingStrategyConfig],
    ) -> KnowledgeBaseUpload:
        """
        Create a new remote upload.

        Args:
            knowledge_base: The Knowledge Base to upload data to.
            data_source_config: The data source config.
            data_source_auth_config: The data source auth config.
            chunking_strategy_config: The chunking strategy config.

        Returns:
            The newly created remote upload.
        """
        response = self._post(
            sub_path=f"{self._sub_path.format(knowledge_base_id=knowledge_base.knowledge_base_id)}",
            request=KnowledgeBaseRemoteUploadRequest(
                data_source_config=data_source_config,
                data_source_auth_config=data_source_auth_config,
                chunking_strategy_config=chunking_strategy_config,
            ),
        )
        response_model = KnowledgeBaseUploadResponse.from_dict(response.json())
        return self.get(id=response_model.upload_id, knowledge_base=knowledge_base)

    def create_local_upload(
        self,
        knowledge_base: KnowledgeBase,
        data_source_config: LocalChunksSourceConfig,
        chunks: List[ChunkToUpload],
    ) -> KnowledgeBaseUpload:
        """
        Create a new local upload.

        Args:
            knowledge_base: The Knowledge Base to upload data to.
            data_source_config: The data source config.
            chunks: The chunks to upload.

        Returns:
            The newly created local upload.
        """
        response = self._post(
            sub_path=f"{self._sub_path.format(knowledge_base_id=knowledge_base.knowledge_base_id)}",
            request=KnowledgeBaseLocalChunkUploadRequest(
                data_source_config=data_source_config,
                chunks=chunks,
            ),
        )
        response_model = KnowledgeBaseUploadResponse.from_dict(response.json())
        return self.get(id=response_model.upload_id, knowledge_base=knowledge_base)

    def get(
        self,
        id: str,
        knowledge_base: KnowledgeBase,
    ) -> KnowledgeBaseUpload:
        """
        Get an Knowledge Base Upload by ID.

        Args:
            id: The ID of the Knowledge Base Upload.
            knowledge_base: The Knowledge Base the upload was created for.

        Returns:
            The Knowledge Base Upload.
        """
        response = self._get(
            sub_path=f"{self._sub_path.format(knowledge_base_id=knowledge_base.knowledge_base_id)}/"
            f"{id}",
        )
        return KnowledgeBaseUpload.from_dict(response.json())

    def list(
        self,
        knowledge_base: KnowledgeBase,
    ) -> List[KnowledgeBaseUpload]:
        """
        List all Knowledge Base Uploads.

        Returns:
            A list of Knowledge Base Uploads.
        """
        response = self._get(
            sub_path=self._sub_path.format(knowledge_base_id=knowledge_base.knowledge_base_id),
        )
        response_model = ListKnowledgeBaseUploadsResponse.from_dict(response.json())
        return response_model.uploads

    def cancel(
        self,
        knowledge_base: KnowledgeBase,
        id: str,
    ) -> bool:
        """
        Cancel an upload.

        Args:
            knowledge_base: The Knowledge Base the upload was created for.
            id: The ID of the upload to cancel.

        Returns:
            True if the upload was canceled, False otherwise.
        """
        response = self._post(
            sub_path=f"{self._sub_path.format(knowledge_base_id=knowledge_base.knowledge_base_id)}/"
            f"{id}/cancel",
            request=None,
        )
        response_model = CancelKnowledgeBaseUploadResponse.from_dict(response.json())
        return response_model.canceled
