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

from langchain_core.embeddings.embeddings import Embeddings
from langchain_core.language_models import LanguageModelInput
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages.base import BaseMessage, BaseMessageChunk
from langchain_core.runnables import Runnable, RunnableLambda
from langchain_core.runnables.config import RunnableConfig
from langchain_core.tools import BaseTool
from langchain_openai.chat_models.azure import AzureChatOpenAI
from langchain_openai.embeddings.azure import AzureOpenAIEmbeddings
from pydantic.main import BaseModel

from llm_sdk.model_providers.aes import AESCipher
from llm_sdk.model_providers.base import AIModel

import logging

logger = logging.getLogger(__name__)


class AzureOpenAILargeLanguageModel(AIModel):

    _model: Optional[BaseChatModel] = None

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

        if app_id is None:
            raise ValueError("app_id is required")
        try:
            aes_client = AESCipher(app_id)
            # decrypt credential
            if credential.get("openai_api_key"):
                credential["openai_api_key"] = aes_client.decrypt(credential["openai_api_key"])
        except Exception as e:
            logger.warning(f"decrypt credential failed: {e}, use original key")

        self._model = AzureChatOpenAI(
            model=credential.get("base_model_name", ""),
            azure_endpoint=credential.get("openai_api_base", ""),
            api_key=credential.get("openai_api_key", ""),
            api_version=credential.get("openai_api_version", ""),
            temperature=credential.get("temperature", 0.7),
            timeout=credential.get("timeout", 60),
            max_retries=3,
        )
        self._credential = credential

    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_key": self._credential.get("openai_api_key", ""),
            "openai_api_base": self._credential.get("openai_api_base", ""),
            "openai_api_version": self._credential.get("openai_api_version", ""),
            "openai_api_type": "azure_openai",
            "openai_api_deployment": self._credential.get("base_model_name", ""),
            "timeout": self._credential.get("timeout", 60),
            "max_retries": self._credential.get("max_retries", 1),
            "temperature": self._credential.get("temperature", 0.7),
        }


class AzureOpenAIEmbeddingModel(AIModel):

    _model: Optional[Embeddings] = None

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

        if app_id is None:
            raise ValueError("app_id is required")
        try:
            aes_client = AESCipher(app_id)
            # decrypt credential
            if credential.get("openai_api_key"):
                credential["openai_api_key"] = aes_client.decrypt(credential["openai_api_key"])
        except Exception as e:
            logger.warning(f"decrypt credential failed: {e}, use original key")

        self._model = AzureOpenAIEmbeddings(
            model=credential.get("base_model_name", ""),
            azure_endpoint=credential.get("openai_api_base", ""),
            api_version=credential.get("openai_api_version", ""),
            api_key=credential.get("openai_api_key", ""),
            timeout=credential.get("timeout", 60),
        )
        self._credential = credential
        self._model_name = model_name

    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_key": self._credential.get("openai_api_key", ""),
            "openai_api_base": self._credential.get("openai_api_base", ""),
            "openai_api_version": self._credential.get("openai_api_version", ""),
            "openai_api_type": "azure_openai",
            "openai_api_deployment": self._credential.get("base_model_name", ""),
        }
