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

import dashscope
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain_community.embeddings.dashscope import DashScopeEmbeddings
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 pydantic.main import BaseModel

from llm_sdk.model_providers.aes import AESCipher
from llm_sdk.model_providers.base import AIModel
from llm_sdk.model_providers.rerank import RerankDocument, RerankModel, RerankResult

import logging

logger = logging.getLogger(__name__)


class TongyiLargeLanguageModel(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("dashscope_api_key"):
                credential["dashscope_api_key"] = aes_client.decrypt(credential["dashscope_api_key"])
        except Exception as e:
            logger.warning(f"decrypt credential failed: {e}, use original key")

        self._model = ChatTongyi(
            model=model_name,
            api_key=credential.get("dashscope_api_key", ""),
            max_retries=credential.get("max_retries", 3),
        )
        self._model_name = model_name
        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("dashscope_api_key", ""),
            "openai_api_type": "tongyi",
            "openai_api_deployment": self._model_name,
            "timeout": self._credential.get("timeout", 60),
            "max_retries": self._credential.get("max_retries", 1),
            "temperature": self._credential.get("temperature", 0.7),
        }


class TongyiTextEmbeddingModel(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("dashscope_api_key"):
                credential["dashscope_api_key"] = aes_client.decrypt(credential["dashscope_api_key"])
        except Exception as e:
            logger.warning(f"decrypt credential failed: {e}, use original key")

        self._model = DashScopeEmbeddings(
            dashscope_api_key=credential.get("dashscope_api_key", ""),
            model=model_name,
            max_retries=credential.get("max_retries", 3),
        )
        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("dashscope_api_key", ""),
            "openai_api_type": "tongyi",
            "openai_api_deployment": self._model_name,
        }


class GTERerankModel(RerankModel):

    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("dashscope_api_key"):
                credential["dashscope_api_key"] = aes_client.decrypt(credential["dashscope_api_key"])
        except Exception as e:
            logger.warning(f"decrypt credential failed: {e}, use original key")

        self._credential = credential
        self._model_name = model_name

    def rerank(
        self,
        query: str,
        docs: list[str],
        score_threshold: Optional[float] = None,
        top_n: Optional[int] = None,
    ) -> RerankResult:
        if len(docs) == 0:
            return RerankResult(docs=[])

        # initialize client
        dashscope.api_key = self._credential.get("dashscope_api_key", "")

        response = dashscope.TextReRank.call(
            query=query,
            documents=docs,
            model=self._model_name,
            top_n=top_n,
            return_documents=True,
        )

        rerank_documents: list[RerankDocument] = []
        if not response.output:
            return RerankResult(docs=rerank_documents)
        for _, result in enumerate(response.output.results):
            # format document
            rerank_document = RerankDocument(
                index=result.index,
                score=result.relevance_score,
                text=result["document"]["text"],
            )

            # score threshold check
            if score_threshold is not None:
                if result.relevance_score >= score_threshold:
                    rerank_documents.append(rerank_document)
            else:
                rerank_documents.append(rerank_document)

        return RerankResult(docs=rerank_documents)

    def unified_credential(self) -> Dict:
        return {
            "openai_api_key": self._credential.get("dashscope_api_key", ""),
            "openai_api_type": "tongyi",
            "openai_api_deployment": self._model_name,
        }
