import json
from json import dumps
from threading import Lock
from typing import Any, AsyncGenerator, Dict, Iterator, List, Optional, Type, Union
from urllib.parse import urljoin

import httpx
import numpy as np
import requests
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.main import BaseModel
from requests import post
from yarl import URL

from llm_sdk.model_providers.aes import AESCipher
from llm_sdk.model_providers.base import AIModel
from llm_sdk.model_providers.errors import (
    CredentialsValidateFailedError,
    InvokeServerUnavailableError,
)
from llm_sdk.model_providers.rerank import RerankDocument, RerankResult
import logging

logger = logging.getLogger(__name__)


_tokenizer: Any = None
_lock = Lock()


class OAIAPICompatLargeLanguageModel(AIModel):
    """
    Model class for Openai compatible large language model.
    """

    _model: Optional[BaseChatModel] = None

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

        try:
            aes_client = AESCipher(app_id)
            # decrypt credential
            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
        self._model = ChatOpenAI(
            model=model_name,
            api_key=credential.get("api_key") or "empty",
            base_url=credential.get("endpoint_url", ""),
        )

    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("endpoint_url", ""),
            "openai_api_type": "openai",
            "openai_api_key": self._credentialget("api_key", ""),
            "openai_api_deployment": self._model_name,
        }


class OAIAPICompatEmbeddingModel(AIModel):

    _model: Optional[Embeddings] = None

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

        try:
            aes_client = AESCipher(app_id)
            # decrypt credential
            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 = OAICompatEmbeddingModel(
            model_name=model_name,
            credential=credential,
        )
        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_base": self._credential.get("endpoint_url", ""),
            "openai_api_type": "openai",
            "openai_api_key": self._credential.get("api_key", ""),
            "openai_api_deployment": self._model_name,
        }


class OAICompatRerankModel(AIModel):

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

        try:
            aes_client = AESCipher(app_id)
            # decrypt credential
            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,
    ) -> RerankResult:
        if len(docs) == 0:
            return RerankResult(docs=[])

        server_url = self._credential["endpoint_url"]
        model_name = self._model_name

        if not server_url:
            raise CredentialsValidateFailedError("server_url is required")
        if not model_name:
            raise CredentialsValidateFailedError("model_name is required")

        url = server_url
        headers = {
            "Content-Type": "application/json",
        }
        if self._credential.get("api_key"):
            headers["Authorization"] = f"Bearer {self._credential['api_key']}"

        # TODO: Do we need truncate docs to avoid llama.cpp return error?

        data = {
            "model": model_name,
            "query": query,
            "documents": docs,
            "top_n": top_n,
            "return_documents": True,
        }

        try:
            response = post(str(URL(url) / "rerank"), headers=headers, data=dumps(data), timeout=60)
            response.raise_for_status()
            results = response.json()

            rerank_documents = []
            scores = [result["relevance_score"] for result in results["results"]]

            # Min-Max Normalization: Normalize scores to 0 ~ 1.0 range
            min_score = min(scores)
            max_score = max(scores)
            score_range = max_score - min_score if max_score != min_score else 1.0  # Avoid division by zero

            for result in results["results"]:
                index = result["index"]

                # Retrieve document text (fallback if llama.cpp rerank doesn't return it)
                text = docs[index]
                document = result.get("document", {})
                if document:
                    if isinstance(document, dict):
                        text = document.get("text", docs[index])
                    elif isinstance(document, str):
                        text = document

                # Normalize the score
                normalized_score = (result["relevance_score"] - min_score) / score_range

                # Create RerankDocument object with normalized score
                rerank_document = RerankDocument(
                    index=index,
                    text=text,
                    score=normalized_score,
                )

                # Apply threshold (if defined)
                if score_threshold is None or normalized_score >= score_threshold:
                    rerank_documents.append(rerank_document)

            # Sort rerank_documents by normalized score in descending order
            rerank_documents.sort(key=lambda doc: doc.score, reverse=True)

            return RerankResult(docs=rerank_documents)

        except httpx.HTTPStatusError as e:
            raise InvokeServerUnavailableError(str(e))

    def unified_credential(self) -> Dict:
        return {
            "openai_api_base": self._credential.get("endpoint_url", ""),
            "openai_api_type": "openai",
            "openai_api_key": self._credential.get("api_key", ""),
            "openai_api_deployment": self._model_name,
        }


class OAICompatEmbeddingModel(Embeddings):
    """
    Model class for Openai compatible embedding model.
    """

    def __init__(self, model_name: str, credential: Dict):
        self._model_name = model_name
        self._credentials = credential

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        """Embed search docs.

        Args:
            texts: List of text to embed.

        Returns:
            List of embeddings.
        """
        return self._invoke_documents(texts)

    def embed_query(self, text: str) -> list[float]:
        """Embed query text.

        Args:
            text: Text to embed.

        Returns:
            Embedding.
        """
        return self._invoke_query(text)

    async def aembed_documents(self, texts: list[str]) -> list[list[float]]:
        """Asynchronous Embed search docs.

        Args:
            texts: List of text to embed.

        Returns:
            List of embeddings.
        """
        return self._invoke_documents(texts)

    async def aembed_query(self, text: str) -> list[float]:
        """Asynchronous Embed query text.

        Args:
            text: Text to embed.

        Returns:
            Embedding.
        """
        return self._invoke_query(text)

    def _invoke_query(self, query: str) -> list[float]:
        headers = {"Content-Type": "application/json"}

        api_key = self._credentials.get("api_key")
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"

        endpoint_url = self._credentials.get("endpoint_url")
        assert endpoint_url is not None, "endpoint_url is required in credentials"
        if not endpoint_url.endswith("/"):
            endpoint_url += "/"

        endpoint_url = urljoin(endpoint_url, "embeddings")

        extra_model_kwargs = {}

        extra_model_kwargs["encoding_format"] = "float"

        payload = {"input": query, "model": self._model_name, **extra_model_kwargs}

        # Make the request to the OpenAI API
        response = requests.post(endpoint_url, headers=headers, data=json.dumps(payload), timeout=(10, 300))

        response.raise_for_status()  # Raise an exception for HTTP errors
        response_data = response.json()

        # Extract embeddings and used tokens from the response
        embeddings_batch = [data["embedding"] for data in response_data["data"]]
        if len(embeddings_batch) != 1:
            raise ValueError("embeddings_batch length less than 1")

        return embeddings_batch[0]

    def _invoke_documents(self, texts: list[str]) -> list[list[float]]:
        """
        Invoke text embedding model

        :param model: model name
        :param credentials: model credentials
        :param texts: texts to embed
        :param user: unique user id
        :param input_type: input type
        :return: embeddings result
        """

        # Prepare headers and payload for the request
        headers = {"Content-Type": "application/json"}

        api_key = self._credentials.get("api_key")
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"

        endpoint_url = self._credentials.get("endpoint_url")
        assert endpoint_url is not None, "endpoint_url is required in credentials"
        if not endpoint_url.endswith("/"):
            endpoint_url += "/"

        endpoint_url = urljoin(endpoint_url, "embeddings")

        extra_model_kwargs = {}
        extra_model_kwargs["encoding_format"] = "float"

        # get model properties
        context_size: int = int(self._credentials.get("context_size", 4096))
        max_chunks: int = self._credentials.get("max_chunks", 1000)

        inputs = []
        indices = []
        used_tokens = 0

        for i, text in enumerate(texts):
            # Here token count is only an approximation based on the GPT2 tokenizer
            # TODO: Optimize for better token estimation and chunking
            num_tokens = self._get_num_tokens_by_gpt2(text)

            if num_tokens >= context_size:
                cutoff = int(np.floor(len(text) * (context_size / num_tokens)))
                # if num tokens is larger than context length, only use the start
                inputs.append(text[0:cutoff])
            else:
                inputs.append(text)
            indices += [i]

        batched_embeddings = []
        _iter = range(0, len(inputs), max_chunks)

        for i in _iter:
            # Prepare the payload for the request
            payload = {"input": inputs[i : i + max_chunks], "model": self._model_name, **extra_model_kwargs}

            # Make the request to the OpenAI API
            response = requests.post(
                endpoint_url, headers=headers, data=json.dumps(payload), timeout=(10, 300)
            )

            response.raise_for_status()  # Raise an exception for HTTP errors
            response_data = response.json()

            # Extract embeddings and used tokens from the response
            embeddings_batch = [data["embedding"] for data in response_data["data"]]
            embedding_used_tokens = response_data["usage"]["total_tokens"]

            used_tokens += embedding_used_tokens
            batched_embeddings += embeddings_batch

        return batched_embeddings

    def _get_num_tokens_by_gpt2(self, text: str) -> int:
        """
        use gpt2 tokenizer to get num tokens
        """
        _tokenizer = OAICompatEmbeddingModel.get_encoder()
        tokens = _tokenizer.encode(text)
        return len(tokens)

    @staticmethod
    def get_encoder() -> Any:
        global _tokenizer, _lock
        with _lock:
            if _tokenizer is None:
                # Try to use tiktoken to get the tokenizer because it is faster
                #
                try:
                    import tiktoken

                    _tokenizer = tiktoken.get_encoding("gpt2")
                except Exception:
                    from os.path import abspath, dirname, join

                    from transformers import (
                        GPT2Tokenizer as TransformerGPT2Tokenizer,  # type: ignore
                    )

                    base_path = abspath(__file__)
                    gpt2_tokenizer_path = join(dirname(base_path), "gpt2")
                    _tokenizer = TransformerGPT2Tokenizer.from_pretrained(gpt2_tokenizer_path)

            return _tokenizer
