"""
Cortex Embed API.

OpenAPI 3.0 specification for the Cortex Embed REST API  # noqa: E501

The version of the OpenAPI document: 0.1.0
Contact: support@snowflake.com
Generated by: https://openapi-generator.tech

Do not edit this file manually.
"""

import json
import logging
import os
import re
import tempfile

from concurrent.futures import Future
from datetime import date, datetime
from typing import Literal, Union, overload

from dateutil.parser import parse
from pydantic import Field, validate_call
from typing_extensions import Annotated

import snowflake.core._http_requests

from snowflake.core._generated import rest
from snowflake.core._generated.api_client import ApiClient
from snowflake.core._internal.utils import is_running_inside_stored_procedure
from snowflake.core.cortex.embed_service._generated import models
from snowflake.core.cortex.embed_service._generated.models.embed_request import EmbedRequest
from snowflake.core.cortex.embed_service._generated.models.embed_response import EmbedResponse
from snowflake.core.exceptions import (  # noqa: F401
    _APITypeError,
    _APIValueError,
)


logger = logging.getLogger(__name__)


class CortexEmbedApi:
    """NOTE: This class is auto generated by OpenAPI Generator.

    Ref: https://openapi-generator.tech

    Do not edit the class manually.
    """

    def __init__(self, root, resource_class, sproc_client):
        self._root = root
        self._resource_name = "embed_service"
        self._resource_class = resource_class

        if is_running_inside_stored_procedure():
            self._api_client = sproc_client
        else:
            self._api_client = ApiClient.get_default()

    @property
    def api_client(self):
        return self._api_client

    def deserialize(self, response, response_type):
        """Deserializes response into an object.

        :param response: RESTResponse object to be deserialized.
        :param response_type: class literal for
            deserialized object, or string of class name.

        :return: deserialized object.
        """
        # handle file downloading
        # save response body into a tmp file and return the instance
        if response_type == "file":
            return self.__deserialize_file(response)

        # fetch data from response object
        try:
            data = json.loads(response.data)
        except ValueError:
            data = response.data

        return self.__deserialize(data, response_type)

    def __deserialize(self, data, klass):
        """Deserializes dict, list, str into an object.

        :param data: dict, list or str.
        :param klass: class literal, or string of class name.

        :return: object.
        """
        if data is None:
            return None

        if isinstance(klass, str):
            if klass.startswith("Iterable["):
                sub_kls = re.match(r"Iterable\[(.*)]", klass).group(1)
                return [self.__deserialize(sub_data, sub_kls) for sub_data in data]

            if klass.startswith("Dict["):
                sub_kls = re.match(r"Dict\[([^,]*), (.*)]", klass).group(2)
                return {k: self.__deserialize(v, sub_kls) for k, v in data.items()}

            # convert str to class
            if klass in snowflake.core._http_requests.NATIVE_TYPES_MAPPING:
                klass = snowflake.core._http_requests.NATIVE_TYPES_MAPPING[klass]
            else:
                klass = getattr(models, klass)

        if klass in snowflake.core._http_requests.PRIMITIVE_TYPES:
            return self.__deserialize_primitive(data, klass)
        elif klass == object:  # noqa: E721
            return self.__deserialize_object(data)
        elif klass == date:
            return self.__deserialize_date(data)
        elif klass == datetime:
            return self.__deserialize_datetime(data)
        else:
            return self.__deserialize_model(data, klass)

    def __deserialize_file(self, response):
        """Deserializes body to file.

        Saves response body into a file in a temporary folder,
        using the filename from the `Content-Disposition` header if provided.

        :param response:  RESTResponse.
        :return: file path.
        """
        fd, path = tempfile.mkstemp(dir=self.api_client.configuration.temp_folder_path)
        os.close(fd)
        os.remove(path)

        content_disposition = response.getheader("Content-Disposition")
        if content_disposition:
            filename = re.search(r'filename=[\'"]?([^\'"\s]+)[\'"]?', content_disposition).group(1)
            path = os.path.join(os.path.dirname(path), filename)

        with open(path, "wb") as f:
            f.write(response.data)

        return path

    def __deserialize_primitive(self, data, klass):
        """Deserializes string to primitive type.

        :param data: str.
        :param klass: class literal.

        :return: int, long, float, str, bool.
        """
        try:
            return klass(data)
        except UnicodeEncodeError:
            return str(data)
        except TypeError:
            return data

    def __deserialize_object(self, value):
        """Return an original value.

        :return: object.
        """
        return value

    def __deserialize_date(self, string):
        """Deserializes string to date.

        :param string: str.
        :return: date.
        """
        try:
            return parse(string).date()
        except ImportError:
            return string
        except ValueError:
            raise rest.APIError(status=0, reason=f"Failed to parse `{string}` as date object")

    def __deserialize_datetime(self, string):
        """Deserializes string to datetime.

        The string should be in iso8601 datetime format.

        :param string: str.
        :return: datetime.
        """
        try:
            return parse(string)
        except ImportError:
            return string
        except ValueError:
            raise rest.APIError(status=0, reason=(f"Failed to parse `{string}` as datetime object"))

    def __deserialize_model(self, data, klass):
        """Deserializes list or dict to model.

        :param data: dict, list.
        :param klass: class literal.
        :return: model object.
        """
        return klass.from_dict(data)

    @overload
    def embed(
        self,
        embed_request: Annotated[
            EmbedRequest, Field(description="Request payload containing model identifier and input text.")
        ],
        async_req: Literal[False] = False,
        **kwargs,
    ) -> EmbedResponse: ...

    @overload
    def embed(
        self,
        embed_request: Annotated[
            EmbedRequest, Field(description="Request payload containing model identifier and input text.")
        ],
        async_req: Literal[True] = True,
        **kwargs,
    ) -> Future[EmbedResponse]: ...

    @overload
    def embed(
        self,
        embed_request: Annotated[
            EmbedRequest, Field(description="Request payload containing model identifier and input text.")
        ],
        async_req: bool = False,
        **kwargs,
    ) -> Union[EmbedResponse, Future[EmbedResponse]]: ...

    @validate_call
    def embed(
        self,
        embed_request: Annotated[
            EmbedRequest, Field(description="Request payload containing model identifier and input text.")
        ],
        **kwargs,
    ) -> Union[EmbedResponse, Future[EmbedResponse]]:
        r"""Perform LLM embed  # noqa: E501

        Perform LLM embedding for input text, similar to the Snowflake Cortex `EMBED_TEXT` functions. For more details on available functions, see: - [EMBED_TEXT_768](https://docs.snowflake.com/en/sql-reference/functions/embed_text-snowflake-cortex) - [EMBED_TEXT_1024](https://docs.snowflake.com/en/sql-reference/functions/embed_text_1024-snowflake-cortex)   # noqa: E501

        This method makes a synchronous HTTP request by default. To make an
        asynchronous HTTP request, please pass async_req=True

        >>> future = api.embed(embed_request, async_req=True)
        >>> result = future.result()
        :param embed_request: Request payload containing model identifier and input text. (required)
        :type embed_request: EmbedRequest
        :param async_req: Whether to execute the request asynchronously.
        :type async_req: bool, optional
        :param _preload_content: if False, the urllib3.HTTPResponse object will
                                 be returned without reading/decoding response
                                 data. Default is True.
        :type _preload_content: bool, optional
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :return: Returns the result object.
                 If the method is called asynchronously,
                 returns a Future object representing the execution of the method.
        :rtype: Union[EmbedResponse, Future[EmbedResponse]]
        """
        kwargs["_return_http_data_only"] = True
        return self.embed_with_http_info(embed_request, **kwargs)

    @validate_call
    def embed_with_http_info(
        self,
        embed_request: Annotated[
            EmbedRequest, Field(description="Request payload containing model identifier and input text.")
        ],
        **kwargs,
    ):
        r"""Perform LLM embed  # noqa: E501

        Perform LLM embedding for input text, similar to the Snowflake Cortex `EMBED_TEXT` functions. For more details on available functions, see: - [EMBED_TEXT_768](https://docs.snowflake.com/en/sql-reference/functions/embed_text-snowflake-cortex) - [EMBED_TEXT_1024](https://docs.snowflake.com/en/sql-reference/functions/embed_text_1024-snowflake-cortex)   # noqa: E501

        This method makes a synchronous HTTP request by default. To make an
        asynchronous HTTP request, please pass async_req=True

        >>> future = api.embed_with_http_info(embed_request, async_req=True)
        >>> result = future.result()
        :param embed_request: Request payload containing model identifier and input text. (required)
        :type embed_request: EmbedRequest
        :param async_req: Whether to execute the request asynchronously.
        :type async_req: bool, optional
        :param _return_http_data_only: response data without head status code
                                       and headers
        :type _return_http_data_only: bool, optional
        :param _preload_content: if False, the urllib3.HTTPResponse object will
                                 be returned without reading/decoding response
                                 data. Default is True.
        :type _preload_content: bool, optional
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the authentication
                              in the spec for a single request.
        :type _request_auth: dict, optional
        :type _content_type: string, optional: force content-type for the request
        :return: Returns the result object.
                 If the method is called asynchronously,
                 returns a Future object representing the execution of the method.
        :rtype: tuple(Union[EmbedResponse, Future[EmbedResponse]], status_code(int), headers(HTTPHeaderDict))
        """
        _params = locals()

        _all_params = ["embed_request"]
        _all_params.extend(
            [
                "async_req",
                "_return_http_data_only",
                "_preload_content",
                "_request_timeout",
                "_request_auth",
                "_content_type",
                "_headers",
            ]
        )

        # validate the arguments
        for _key, _val in _params["kwargs"].items():
            if _key not in _all_params:
                raise _APITypeError(f"Got an unexpected keyword argument '{_key}' to method embed")
            _params[_key] = _val
        del _params["kwargs"]

        _collection_formats = {}

        # process the path parameters
        _path_params = {}

        # process the query parameters
        _query_params = []

        # process the header parameters
        _header_params = dict(_params.get("_headers", {}))

        # process the form parameters
        _form_params = []
        _files = {}

        # process the body parameter
        _body_params = None

        if _params["embed_request"] is not None:
            _body_params = _params["embed_request"]

        # set the HTTP header `Accept`
        _header_params["Accept"] = self.api_client.select_header_accept(["application/json"])

        # set the HTTP header `Content-Type`
        _content_types_list = _params.get(
            "_content_type", self.api_client.select_header_content_type(["application/json"])
        )
        if _content_types_list:
            _header_params["Content-Type"] = _content_types_list

        # authentication setting
        _auth_settings = ["ExternalOAuth", "KeyPair", "ProgrammaticAccessToken", "SnowflakeOAuth"]

        _response_types_map = {
            "200": "EmbedResponse",
            "400": "ErrorResponse",
            "401": "ErrorResponse",
            "403": "ErrorResponse",
            "404": "ErrorResponse",
            "405": "ErrorResponse",
            "500": "ErrorResponse",
            "503": "ErrorResponse",
            "504": "ErrorResponse",
        }

        return self.api_client.call_api(
            self._root,
            "/api/v2/cortex/inference:embed",
            "POST",
            _path_params,
            _query_params,
            _header_params,
            body=_body_params,
            post_params=_form_params,
            files=_files,
            response_types_map=_response_types_map,
            auth_settings=_auth_settings,
            async_req=_params.get("async_req"),
            _return_http_data_only=_params.get("_return_http_data_only"),
            _preload_content=_params.get("_preload_content", True),
            _request_timeout=_params.get("_request_timeout"),
            collection_formats=_collection_formats,
            _request_auth=_params.get("_request_auth"),
            _deserialize_response_fn=self.deserialize,
        )
