import json
import time
from decimal import Decimal
from urllib.parse import urljoin

import requests
import uuid

from dify_plugin.entities import I18nObject
from dify_plugin.entities.model import (
    AIModelEntity,
    EmbeddingInputType,
    FetchFrom,
    ModelPropertyKey,
    ModelType,
    PriceConfig,
    PriceType,
)
from dify_plugin.entities.model.text_embedding import (
    EmbeddingUsage,
    TextEmbeddingResult,
)
from dify_plugin.errors.model import (
    CredentialsValidateFailedError,
)
from dify_plugin.interfaces.model.openai_compatible.common import _CommonOaiApiCompat
from dify_plugin.interfaces.model.text_embedding_model import TextEmbeddingModel


class OAICompatEmbeddingModel(_CommonOaiApiCompat, TextEmbeddingModel):
    """
    Model class for an OpenAI API-compatible text embedding model.
    """

    def _invoke(
        self,
        model: str,
        credentials: dict,
        texts: list[str],
        user: str | None = None,
        input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
    ) -> TextEmbeddingResult:
        """
        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 = credentials.get("api_key")
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"

        endpoint_url = credentials.get("endpoint_url","")
        # if not endpoint_url.endswith("/"):
        #     endpoint_url += "/"

        # endpoint_url = urljoin(endpoint_url, "embeddings")
        endpoint_url = urljoin(endpoint_url, "")
        print("端点的url是{}".format(endpoint_url))
        extra_model_kwargs = {}
        if user:
            extra_model_kwargs["user"] = user

        extra_model_kwargs["encoding_format"] = "float"

        # get model properties
        context_size = self._get_context_size(model, credentials)
        max_chunks = self._get_max_chunks(model, credentials)

        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((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 = {
                "inputs": inputs[i : i + max_chunks],
                "model": credentials.get("endpoint_model_name", model),
                **extra_model_kwargs,
            }

            # Make the request to the OpenAI API
            print("词向量接口最后送出去的数据是：{} --{} --{}".format(endpoint_url,headers,payload))
            response = requests.post(
                endpoint_url,
                headers=headers,
                data=json.dumps(payload),
                timeout=(10, 300),
            )
            data = response.json()
            print("词向量返回的原始数据是:{}".format(data))
            arr = []
            if len(data) >0 :
                arr = data[0]
            response.raise_for_status()  # Raise an exception for HTTP errors
            response_data =  {
                'data':[{'embedding':arr ,'index':0,'object':'embedding'}],
                'object':'list',
                'model': credentials.get('model',''),
                'usage': {'prompt_tokens':2,'total_tokens':2},
                'id': uuid.uuid4()
            }
            print("词向量返回---加工的数据是:{}".format(response_data))
            # 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

        # calc usage
        usage = self._calc_response_usage(model=model, credentials=credentials, tokens=used_tokens)

        return TextEmbeddingResult(embeddings=batched_embeddings, usage=usage, model=model)

    def get_num_tokens(self, model: str, credentials: dict, texts: list[str]) -> list[int]:
        """
        Approximate number of tokens for given messages using GPT2 tokenizer

        :param model: model name
        :param credentials: model credentials
        :param texts: texts to embed
        :return:
        """
        return [self._get_num_tokens_by_gpt2(text) for text in texts]

    def validate_credentials(self, model: str, credentials: dict) -> None:
        """
        Validate model credentials

        :param model: model name
        :param credentials: model credentials
        :return:
        """
        try:
            headers = {"Content-Type": "application/json"}

            api_key = credentials.get("api_key")

            if api_key:
                headers["Authorization"] = f"Bearer {api_key}"

            endpoint_url = credentials.get("endpoint_url", "")
            # if not endpoint_url.endswith("/"):
            #     endpoint_url += "/"

            # endpoint_url = urljoin(endpoint_url, "embeddings")
            endpoint_url = urljoin(endpoint_url, "")

            payload = {"inputs": "ping", "model": credentials.get("endpoint_model_name", model)}

            response = requests.post(
                url=endpoint_url,
                headers=headers,
                data=json.dumps(payload),
                timeout=(10, 300),
            )

            if response.status_code != 200:
                raise CredentialsValidateFailedError(
                    f"Credentials validation failed with status code {response.status_code}"
                )

            try:
                data = response.json()
                arr = []
                if len(data) >0 :
                    arr = data[0]
                json_result = {
                    'data':[{'embedding':arr ,'index':0,'object':'embedding'}],
                    'object':'list',
                    'model': credentials.get('model',''),
                    'usage': {'prompt_tokens':2,'total_tokens':2},
                    'id': uuid.uuid4()
                }
            except json.JSONDecodeError as e:
                raise CredentialsValidateFailedError("Credentials validation failed: JSON decode error") from e

            if "model" not in json_result:
                raise CredentialsValidateFailedError("Credentials validation failed: invalid response")
        except CredentialsValidateFailedError:
            raise
        except Exception as ex:
            raise CredentialsValidateFailedError(str(ex)) from ex

    def get_customizable_model_schema(self, model: str, credentials: dict) -> AIModelEntity:
        """
        generate custom model entities from credentials
        """
        entity = AIModelEntity(
            model=model,
            label=I18nObject(en_US=model),
            model_type=ModelType.TEXT_EMBEDDING,
            fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,
            model_properties={
                ModelPropertyKey.CONTEXT_SIZE: int(credentials.get("context_size", 512)),
                ModelPropertyKey.MAX_CHUNKS: 1,
            },
            parameter_rules=[],
            pricing=PriceConfig(
                input=Decimal(credentials.get("input_price", 0)),
                unit=Decimal(credentials.get("unit", 0)),
                currency=credentials.get("currency", "USD"),
            ),
        )

        return entity

    def _calc_response_usage(self, model: str, credentials: dict, tokens: int) -> EmbeddingUsage:
        """
        Calculate response usage

        :param model: model name
        :param credentials: model credentials
        :param tokens: input tokens
        :return: usage
        """
        # get input price info
        input_price_info = self.get_price(
            model=model,
            credentials=credentials,
            price_type=PriceType.INPUT,
            tokens=tokens,
        )

        # transform usage
        usage = EmbeddingUsage(
            tokens=tokens,
            total_tokens=tokens,
            unit_price=input_price_info.unit_price,
            price_unit=input_price_info.unit,
            total_price=input_price_info.total_amount,
            currency=input_price_info.currency,
            latency=time.perf_counter() - self.started_at,
        )

        return usage
