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

from langchain_core.embeddings.embeddings import Embeddings
from langchain_core.language_models.base import LanguageModelInput
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages.base import BaseMessage, BaseMessageChunk
from langchain_core.runnables.base import Runnable, RunnableLambda
from langchain_core.runnables.config import RunnableConfig
from langchain_core.tools.base import BaseTool
from langchain_ollama import ChatOllama
from langchain_ollama.embeddings import OllamaEmbeddings
from pydantic.main import BaseModel

from llm_sdk.model_providers.base import AIModel


class OllamaLargeLanguageModel(AIModel):
    """
    Model class for Ollama large language model.
    """

    _model: Optional[BaseChatModel] = None

    def __init__(self, model_name: str, credential: Dict, app_id: Optional[str] = None):
        super().__init__()

        self._credential = credential
        self._model_name = model_name
        self._model = ChatOllama(
            base_url=credential.get("base_url", ""),
            model=model_name,
        )

    def invoke(
        self,
        input_dict: Union[List[BaseMessage], List[str], str],
        config: Optional[RunnableConfig] = None,
        **kwargs
    ) -> Any:
        return self._model.invoke(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.ainvoke(input_dict, config=config, **kwargs)

    def stream(
        self, input_dict: List[BaseMessage], config: Optional[RunnableConfig] = None, **kwargs
    ) -> Iterator[BaseMessageChunk]:
        self._model.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.astream(input_dict, config=config, **kwargs):
            yield chunk

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

    def with_structured_output(self, schema: Union[Dict, Type[BaseModel]], **kwargs) -> Runnable:
        def wrapper(input_data: LanguageModelInput, config: Optional[RunnableConfig] = None):
            structured_llm = self._model.with_structured_output(schema=schema, **kwargs)

            return structured_llm.invoke(input_data, config=config)

        return RunnableLambda(wrapper)

    def unified_credential(self) -> Dict:
        return {
            "openai_api_base": self._credential.get("base_url", ""),
            "openai_api_type": "ollama",
            "openai_api_deployment": self._model_name,
        }


class OllamaEmbeddingModel(AIModel):

    _model: Optional[Embeddings] = None

    def __init__(self, mode_name: str, credential: Dict, app_id: Optional[str] = None):
        super().__init__()
        self._model_name = mode_name
        self._credential = credential

        self._model = OllamaEmbeddings(
            model=mode_name,
            base_url=credential.get("base_url", ""),
        )

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

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

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

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

    def unified_credential(self) -> Dict:
        return {
            "openai_api_base": self._credential.get("base_url", ""),
            "openai_api_type": "ollama",
            "openai_api_deployment": self._model_name,
        }
