from typing import List, Union

from dashscope.api_entities.dashscope_response import DashScopeAPIResponse
from dashscope.client.base_api import BaseApi
from dashscope.common.constants import TEXT_EMBEDDING_INPUT_KEY
from dashscope.common.utils import _get_task_group_and_task


class TextEmbedding(BaseApi):
    task = 'text-embedding'

    class Models:
        text_embedding_v1 = 'text-embedding-v1'
        text_embedding_v2 = 'text-embedding-v2'

    @classmethod
    def call(cls, model: str, input: Union[str, List[str]],
             **kwargs) -> DashScopeAPIResponse:
        """Get embedding of text input.

        Args:
            model (str): The embedding model name.
            input (Union[str, List[str], io.IOBase]): The text input,
                can be a text or list of text or opened file object,
                if opened file object, will read all lines,
                one embedding per line.
            **kwargs:
                text_type(str, `optional`): query or document.

        Returns:
            DashScopeAPIResponse: The embedding result.
        """
        embedding_input = {}
        if isinstance(input, str):
            embedding_input[TEXT_EMBEDDING_INPUT_KEY] = [input]
        else:
            embedding_input[TEXT_EMBEDDING_INPUT_KEY] = input
        kwargs.pop('stream', False)  # not support streaming output.
        task_group, function = _get_task_group_and_task(__name__)
        return super().call(model=model,
                            input=embedding_input,
                            task_group=task_group,
                            task=TextEmbedding.task,
                            function=function,
                            **kwargs)