from enum import Enum
from typing import (
    Any,
    AsyncGenerator,
    Dict,
    Iterator,
    List,
    Optional,
    Type,
    Union,
)

from langchain_core.embeddings.embeddings import Embeddings
from langchain_core.messages.base import BaseMessage, BaseMessageChunk
from langchain_core.runnables import Runnable
from langchain_core.runnables.config import RunnableConfig
from langchain_core.tools.base import BaseTool
from pydantic.main import BaseModel

from llm_sdk.model_providers.rerank import RerankModel


class ConfigType(Enum):
    CHAT = "CHAT"
    EMBEDDING = "EMBEDDING"
    RERANK = "RERANK"


class AIModel(Runnable, Embeddings, RerankModel):

    def __init__(self):
        super().__init__()

    def invoke(
        self,
        input_dict: Union[List[BaseMessage], List[str], str],
        config: Optional[RunnableConfig] = None,
        **kwargs
    ) -> Any:
        raise NotImplementedError("Chat is not supported for this model")

    async def ainvoke(
        self,
        input_dict: Union[List[BaseMessage], List[str], str],
        config: Optional[RunnableConfig] = None,
        **kwargs
    ) -> Any:
        raise NotImplementedError("Chat is not supported for this model")

    def stream(
        self, input_dict: List[BaseMessage], config: Optional[RunnableConfig] = None, **kwargs
    ) -> Iterator[BaseMessageChunk]:
        raise NotImplementedError("Chat is not supported for this model")

    async def astream(
        self, input_dict: List[BaseMessage], config: Optional[RunnableConfig] = None, **kwargs
    ) -> AsyncGenerator[BaseMessageChunk, None]:
        raise NotImplementedError("Chat is not supported for this model")

    def embed_documents(self, texts: List[str], config: Optional[RunnableConfig] = None) -> List[List[float]]:
        raise NotImplementedError("Embedding is not supported for this model")

    async def aembed_documents(
        self, texts: List[str], config: Optional[RunnableConfig] = None
    ) -> List[List[float]]:
        raise NotImplementedError("Embedding is not supported for this model")

    def embed_query(self, text: str, config: Optional[RunnableConfig] = None) -> List[float]:
        raise NotImplementedError("Embedding is not supported for this model")

    async def aembed_query(self, text: str, config: Optional[RunnableConfig] = None) -> List[float]:
        raise NotImplementedError("Embedding is not supported for this model")

    def bind_tools(self, tools: List[BaseTool], **kwargs) -> "AIModel":
        raise NotImplementedError("Chat is not supported for this model")

    def with_structured_output(self, schema: Union[Dict, Type[BaseModel]], **kwargs) -> Runnable:
        raise NotImplementedError("Chat is not supported for this model")

    def rerank(
        self,
        query: str,
        docs: list[str],
        score_threshold: Optional[float] = None,
        top_n: Optional[int] = None,
    ):
        raise NotImplementedError("Rerank is not supported for this model")

    def unified_credential(credential: Dict) -> Dict:
        raise NotImplementedError("Provider is not supported for this model")
