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

from langchain_community.embeddings.xinference import XinferenceEmbeddings
from langchain_community.llms.xinference import Xinference
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_openai.chat_models.base import ChatOpenAI
from pydantic import SecretStr
from pydantic.main import BaseModel
from xinference_client.client.restful.restful_client import (
    RESTfulRerankModelHandle,
)

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

import logging

logger = logging.getLogger(__name__)


class XinferenceAILargeLanguageModel(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("server_url"):
                credential["server_url"] = aes_client.decrypt(credential["server_url"])
            else:
                raise ValueError("server_url is required")
            if credential.get("api_key"):
                credential["api_key"] = aes_client.decrypt(credential["api_key"])
        except Exception as e:
            logger.warning(f"decrypt credential failed: {e}, use original key")

        api_base = credential.get("server_url")
        if api_base:
            api_base = api_base.removesuffix("/")
            if not api_base.endswith("/v1"):
                api_base = api_base + "/v1"
        self._model = ChatOpenAI(
            model=model_name,
            api_key=credential.get("api_key") or SecretStr("sk-xxxxxx"),
            base_url=api_base,
            temperature=credential.get("temperature", 0.7),
            timeout=credential.get("timeout", 60),
            max_retries=3,
        )
        self._credential = credential
        self._model_name = 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_key": self._credential.get("api_key") or SecretStr("sk-xxxxxx"),
            "openai_api_base": self._credential.get("server_url"),
            "openai_api_version": self._credential.get("model_uid", ""),
            "openai_api_type": "xinference",
            "openai_api_deployment": self._model_name,
        }


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

        self._model = XinferenceEmbeddings(
            server_url=credential.get("server_url"),
            model_uid=credential.get("model_uid", ""),
        )
        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("api_key") or SecretStr("sk-xxxxxx"),
            "openai_api_base": self._credential.get("server_url"),
            "openai_api_version": self._credential.get("model_uid", ""),
            "openai_api_type": "xinference",
            "openai_api_deployment": self._model_name,
        }


class XinferenceRerankModel(AIModel):

    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("server_url"):
                credential["server_url"] = aes_client.decrypt(credential["server_url"])
            if credential.get("api_key"):
                credential["api_key"] = aes_client.decrypt(credential["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,
    ) -> Any:
        """
        Invoke rerank model

        :param model: model name
        :param credentials: model credentials
        :param query: search query
        :param docs: docs for reranking
        :param score_threshold: score threshold
        :param top_n: top n
        :param user: unique user id
        :return: rerank result
        """
        if len(docs) == 0:
            return RerankResult(docs=[])

        server_url = self._credential.get("server_url", "")
        if not server_url:
            raise InvokeServerUnavailableError("server_url is required")
        model_uid = self._credential.get("model_uid", "")
        api_key = self._credential.get("api_key", "")
        server_url = server_url.removesuffix("/")
        auth_headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}

        params = {"documents": docs, "query": query, "top_n": top_n, "return_documents": True}
        try:
            handle = RESTfulRerankModelHandle(model_uid, server_url, auth_headers)
            response = handle.rerank(**params)
        except RuntimeError as e:
            if "rerank hasn't support extra parameter" not in str(e):
                raise InvokeServerUnavailableError(str(e))

            # compatible xinference server between v0.10.1 - v0.12.1, not support 'return_len'
            handle = RESTfulRerankModelHandleWithoutExtraParameter(model_uid, server_url, auth_headers)
            response = handle.rerank(**params)

        rerank_documents = []
        for idx, result in enumerate(response["results"]):
            # format document
            index = result["index"]
            page_content = (
                result["document"] if isinstance(result["document"], str) else result["document"]["text"]
            )
            rerank_document = RerankDocument(
                index=index,
                text=page_content,
                score=result["relevance_score"],
            )

            # score threshold check
            if score_threshold is None or result["relevance_score"] >= score_threshold:
                rerank_documents.append(rerank_document)

        return RerankResult(docs=rerank_documents)

    def unified_credential(self) -> Dict:
        return {
            "openai_api_key": self._credential.get("api_key") or SecretStr("sk-xxxxxx"),
            "openai_api_base": self._credential.get("server_url"),
            "openai_api_version": self._credential.get("model_uid", ""),
            "openai_api_type": "xinference",
            "openai_api_deployment": self._model_name,
        }


class RESTfulRerankModelHandleWithoutExtraParameter(RESTfulRerankModelHandle):
    def rerank(
        self,
        documents: list[str],
        query: str,
        top_n: Optional[int] = None,
        max_chunks_per_doc: Optional[int] = None,
        return_documents: Optional[bool] = None,
        **kwargs,
    ):
        url = f"{self._base_url}/v1/rerank"
        request_body = {
            "model": self._model_uid,
            "documents": documents,
            "query": query,
            "top_n": top_n,
            "max_chunks_per_doc": max_chunks_per_doc,
            "return_documents": return_documents,
        }

        import requests

        response = requests.post(url, json=request_body, headers=self.auth_headers)
        if response.status_code != 200:
            raise InvokeServerUnavailableError(
                f"Failed to rerank documents, detail: {response.json()['detail']}"
            )
        response_data = response.json()
        return response_data
