from typing import Any, AsyncGenerator, Dict, Iterator, List, Optional, Type, Union

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_provider_factory import ModelProviderFactory
from llm_sdk.model_providers.base import AIModel, ConfigType
import requests


class ModelInstance(Runnable):
    """
    A ModelInstance class that can be used to create a LLM instance based on the given tool name.
    """

    model_insntance: Optional[AIModel] = None

    def __init__(
        self,
        config_type: ConfigType = ConfigType.CHAT,
        model_params: Optional[Dict] = None,
        app_id: Optional[str] = "chatbot",
    ):
        self._tools: Optional[List[BaseTool]] = None
        self.model_insntance = ModelProviderFactory().build_provider_model(
            model_provider=model_params.get("provider"),
            model_name=model_params.get("model"),
            config_type=config_type,
            crendential=model_params.get("credentials", {}),
            app_id=app_id,
        )

    def invoke(
        self,
        input_dict: Union[List[BaseMessage], List[str], str],
        config: Optional[RunnableConfig] = None,
        **kwargs,
    ) -> Any:
        return self.model_insntance.invoke(input_dict=input_dict, config=config, **kwargs)

    async def ainvoke(
        self,
        input_dict: Union[List[BaseMessage], List[str], str],
        config: Optional[RunnableConfig] = None,
        **kwargs,
    ) -> Any:
        return await self.model_insntance.ainvoke(input_dict, config=config, **kwargs)

    def stream(
        self, input_dict: List[BaseMessage], config: Optional[RunnableConfig] = None, **kwargs
    ) -> Iterator[BaseMessageChunk]:
        return self.model_insntance.stream(input_dict, config=config, **kwargs)

    async def astream(
        self, input_dict: List[BaseMessage], config: Optional[RunnableConfig] = None, **kwargs
    ) -> AsyncGenerator[BaseMessageChunk, None]:
        async for chunk in self.model_insntance.astream(input_dict, config=config, **kwargs):
            yield chunk

    def embed_documents(self, texts: List[str], config: Optional[RunnableConfig] = None) -> List[List[float]]:
        return self.model_insntance.embed_documents(texts, config=config)

    async def aembed_documents(
        self, texts: List[str], config: Optional[RunnableConfig] = None
    ) -> List[List[float]]:
        return await self.model_insntance.aembed_documents(texts, config=config)

    def embed_query(self, text: str, config: Optional[RunnableConfig] = None) -> List[float]:
        return self.model_insntance.embed_query(text, config=config)

    async def aembed_query(self, text: str, config: Optional[RunnableConfig] = None) -> List[float]:
        return await self.model_insntance.aembed_query(text, config=config)

    def bind_tools(self, tools: List[BaseTool], **kwargs) -> "ModelInstance":
        self._tools = tools
        if self.model_insntance is not None:
            self.model_insntance = self.model_insntance.bind_tools(tools, **kwargs)
        return self

    def with_structured_output(self, schema: Union[Dict, Type[BaseModel]], **kwargs) -> Runnable:
        return self.model_insntance.with_structured_output(schema=schema, **kwargs)

    def rerank(
        self,
        input: Union[List[BaseMessage], List[str], str],
        documents: List[str],
        top_n: Optional[int] = None,
        score_threshold: Optional[float] = None,
    ):
        return self.model_insntance.rerank(
            query=input, docs=documents, top_n=top_n, score_threshold=score_threshold
        )

    def unified_credential(self) -> Dict:
        return self.model_insntance.unified_credential()


class ModelInstanceFactory:

    @staticmethod
    def get_model_instance(
        model_managment_url: str,
        config_type: ConfigType = ConfigType.CHAT,
        model_params: Optional[Dict] = None,
        app_id: Optional[str] = None,
        app_name: Optional[str] = None,
        token: Optional[str] = None,
    ) -> ModelInstance:
        """_summary_

        Args:
            model_managment_url (str): model management url
            config_type (ConfigType, optional): Defaults to ConfigType.CHAT.
            app_id (Optional[str], optional): Defaults to None.

        Returns:
            ModelInstance: _description_
        """
        if model_params:
            return ModelInstance(
                config_type=config_type,
                model_params=model_params,
                app_id=app_id,
            )

        request_url = f"{model_managment_url.rstrip('/')}/v1.0/system/provider-model"
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }
        if config_type == ConfigType.CHAT:
            model_type = "llm"
        elif config_type == ConfigType.EMBEDDING:
            model_type = "text-embedding"
        elif config_type == ConfigType.RERANK:
            model_type = "rerank"
        params = {
            "agent_name": app_name,
            "model_type": model_type,
        }

        response = requests.get(request_url, headers=headers, params=params)
        response.raise_for_status()

        model_params = response.json()

        model_instance = ModelInstance(
            config_type=config_type,
            model_params=model_params,
            app_id=app_id,
        )

        return model_instance
