# Copyright 2025 Iguazio
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import TYPE_CHECKING, Any, Optional, Union

import mlrun
from mlrun.datastore.model_provider.model_provider import (
    InvokeResponseFormat,
    ModelProvider,
    UsageResponseKeys,
)

if TYPE_CHECKING:
    from transformers.pipelines.base import Pipeline
    from transformers.pipelines.text_generation import ChatType


class HuggingFaceProvider(ModelProvider):
    """
    HuggingFaceProvider is a wrapper around the Hugging Face Transformers pipeline
    that provides an interface for interacting with a wide range of Hugging Face models.

    It supports synchronous operations, enabling flexible integration into various workflows.

    This class extends the ModelProvider base class and implements Hugging Face-specific
    functionality, including pipeline initialization, default text generation operations,
    and custom operations tailored to the Hugging Face Transformers pipeline API.
    """

    def __init__(
        self,
        parent,
        schema,
        name,
        endpoint="",
        secrets: Optional[dict] = None,
        default_invoke_kwargs: Optional[dict] = None,
    ):
        endpoint = endpoint or mlrun.mlconf.model_providers.huggingface_default_model
        if schema != "huggingface":
            raise mlrun.errors.MLRunInvalidArgumentError(
                "HuggingFaceProvider supports only 'huggingface' as the provider kind."
            )
        super().__init__(
            parent=parent,
            kind=schema,
            name=name,
            endpoint=endpoint,
            secrets=secrets,
            default_invoke_kwargs=default_invoke_kwargs,
        )
        self.options = self.get_client_options()
        self._expected_operation_type = None
        self.load_client()

    @staticmethod
    def _extract_string_output(response: list[dict]) -> str:
        """
        Extracts the first generated string from Hugging Face pipeline output,
        regardless of whether it's plain text-generation or chat-style output.
        """
        if not isinstance(response, list) or len(response) == 0:
            raise ValueError("Empty or invalid pipeline output")
        if len(response) != 1:
            raise mlrun.errors.MLRunInvalidArgumentError(
                "HuggingFaceProvider: extracting string from response is only supported for single-response outputs"
            )
        return response[0].get("generated_text")

    @classmethod
    def parse_endpoint_and_path(cls, endpoint, subpath) -> (str, str):
        if endpoint and subpath:
            endpoint = endpoint + subpath
            # In HuggingFace, "/" in a model name is part of the name — `subpath` is not used.
            subpath = ""
        return endpoint, subpath

    def _response_handler(
        self,
        response: Union[str, list],
        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
        messages: Union[str, list[str], "ChatType", list["ChatType"]] = None,
        **kwargs,
    ) -> Union[str, list, dict[str, Any]]:
        """
        Same as `ModelProvider._response_handler`.

        * Expected to receive the response with `return_full_text=False`.

        :param messages:                Same as in `ModelProvider._response_handler`.
        :param response:                Same as in `ModelProvider._response_handler`.
        :param invoke_response_format:  Same as in `ModelProvider._response_handler`, in full and string modes.

                                        For usage mode, generate 3 statistics:
                                        prompt_tokens, completion_tokens and total_tokens.

                                        NOTE: Token counts are estimated after answer generation and
                                        may differ from the actual tokens generated by the model due to
                                        internal decoding behavior and implementation details.

        :param kwargs:                  Same as in `ModelProvider._response_handler`.

        :return: The result formatted according to the `invoke_response_format`.

        :raises MLRunInvalidArgumentError: If extracting the string response fails.
        :raises MLRunRuntimeError: If applying the chat template to the model fails.
        """
        if InvokeResponseFormat.is_str_response(invoke_response_format.value):
            str_response = self._extract_string_output(response)
            if invoke_response_format == InvokeResponseFormat.STRING:
                return str_response
            if invoke_response_format == InvokeResponseFormat.USAGE:
                tokenizer = self.client.tokenizer
                if not isinstance(messages, str):
                    try:
                        messages = tokenizer.apply_chat_template(
                            messages, tokenize=False, add_generation_prompt=True
                        )
                    except Exception as e:
                        raise mlrun.errors.MLRunRuntimeError(
                            f"Failed to apply chat template using the tokenizer for model '{self.model}'. "
                            "This may indicate that the tokenizer does not support chat formatting, "
                            "or that the input format is invalid. "
                            f"Original error: {e}"
                        )
                prompt_tokens = len(tokenizer.encode(messages))
                completion_tokens = len(tokenizer.encode(str_response))
                total_tokens = prompt_tokens + completion_tokens
                usage = {
                    "prompt_tokens": prompt_tokens,
                    "completion_tokens": completion_tokens,
                    "total_tokens": total_tokens,
                }
                response = {
                    UsageResponseKeys.ANSWER: str_response,
                    UsageResponseKeys.USAGE: usage,
                }
        return response

    def load_client(self) -> None:
        """
        Initializes the Hugging Face pipeline using the provided options.

        This method imports the `pipeline` function from the `transformers` package,
        creates a pipeline instance with the specified task and model (from `self.options`),
        and assigns it to `self._client`.

        Note: Hugging Face pipelines are synchronous and do not support async invocation.

        :raises:
            ImportError: If the `transformers` package is not installed.
        """
        try:
            from transformers import pipeline, AutoModelForCausalLM  # noqa
            from transformers import AutoTokenizer  # noqa
            from transformers.pipelines.base import Pipeline  # noqa

            self._client = pipeline(model=self.model, **self.options)
            self._expected_operation_type = Pipeline
        except ImportError as exc:
            raise ImportError("transformers package is not installed") from exc

    def get_client_options(self):
        res = dict(
            task=self._get_secret_or_env("HF_TASK") or "text-generation",
            token=self._get_secret_or_env("HF_TOKEN"),
            device=self._get_secret_or_env("HF_DEVICE"),
            device_map=self._get_secret_or_env("HF_DEVICE_MAP"),
            trust_remote_code=self._get_secret_or_env("HF_TRUST_REMOTE_CODE"),
            model_kwargs=self._get_secret_or_env("HF_MODEL_KWARGS"),
        )
        return self._sanitize_options(res)

    def custom_invoke(
        self, operation: Optional["Pipeline"] = None, **invoke_kwargs
    ) -> Union[list, dict, Any]:
        """
        HuggingFace implementation of `ModelProvider.custom_invoke`.
        Use the default config in provider client/ user defined client:

        Example:
        ```python
            image = Image.open(image_path)
            pipeline_object =  pipeline("image-classification", model="microsoft/resnet-50")
            result = hf_provider.custom_invoke(
                pipeline_object,
                inputs=image,
            )
        ```


        :param operation:               A pipeline object
        :param invoke_kwargs:           Keyword arguments to pass to the operation.
        :return:                        The full response returned by the operation.

        """
        invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
        if operation:
            if not isinstance(operation, self._expected_operation_type):
                raise mlrun.errors.MLRunInvalidArgumentError(
                    "Huggingface operation must inherit" " from 'Pipeline' object"
                )
            return operation(**invoke_kwargs)
        else:
            return self.client(**invoke_kwargs)

    def invoke(
        self,
        messages: Union[str, list[str], "ChatType", list["ChatType"]],
        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
        **invoke_kwargs,
    ) -> Union[str, list, dict[str, Any]]:
        """
        HuggingFace-specific implementation of `ModelProvider.invoke`.
        Invokes a HuggingFace model operation using the synchronous client.
        For full details, see `ModelProvider.invoke`.

        :param messages:
            Same as `ModelProvider.invoke`.

        :param invoke_response_format: InvokeResponseFormat
            Specifies the format of the returned response. Options:

            - "string": Returns only the generated text content, extracted from a single response.
            - "usage":  Combines the generated text with metadata (e.g., token usage), returning a dictionary:

            .. code-block:: json
                {
                    "answer": "<generated_text>",
                    "usage": {
                        "prompt_tokens": <int>,
                        "completion_tokens": <int>,
                        "total_tokens": <int>
                    }
                }

            - "full":   Returns the raw response object from the HuggingFace model,
                        typically a list of generated sequences (dictionaries).
                        This format does not include token usage statistics.

        :param invoke_kwargs:
            Additional keyword arguments passed to the HuggingFace client. Same as in `ModelProvider.invoke`.

        :return:
            A string, dictionary, or list of model outputs, depending on `invoke_response_format`.
        """

        if self.client.task != "text-generation":
            raise mlrun.errors.MLRunInvalidArgumentError(
                "HuggingFaceProvider.invoke supports text-generation task only"
            )
        if InvokeResponseFormat.is_str_response(invoke_response_format.value):
            invoke_kwargs["return_full_text"] = False
        response = self.custom_invoke(text_inputs=messages, **invoke_kwargs)
        response = self._response_handler(
            messages=messages,
            response=response,
            invoke_response_format=invoke_response_format,
        )
        return response
