# Copyright 2025 The MelowRAG Author @erfanzar (Erfan Zare Chavoshi).
#
# 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
#
#     https://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.

import inspect

from transformers import pipeline

from ..modeling import Models
from ..utilities import Resolver
from .tensors import Tensors


class HFPipeline(Tensors):
    """
    Light wrapper around Hugging Face Transformers pipeline component for selected tasks. Adds support for model
    quantization and minor interface changes.
    """

    def __init__(self, task, path=None, quantize=False, gpu=False, model=None, **kwargs):
        """
        Loads a new pipeline model.

        Args:
            task: pipeline task or category
            path: optional path to model, accepts Hugging Face model hub id, local path or (model, tokenizer) tuple.
                  uses default model for task if not provided
            quantize: if model should be quantized, defaults to False
            gpu: True/False if GPU should be enabled, also supports a GPU device id
            model: optional existing pipeline model to wrap
            kwargs: additional keyword arguments to pass to pipeline model
        """

        if model:
            self.pipeline = model.pipeline if isinstance(model, HFPipeline) else model
        else:
            deviceid = Models.deviceid(gpu) if "device_map" not in kwargs else None
            device = Models.device(deviceid) if deviceid is not None else None

            modelargs, kwargs = self.parseargs(**kwargs)

            if isinstance(path, list | tuple):
                config = path[1] if path[1] and isinstance(path[1], str) else None

                model = Models.load(path[0], config, task)

                self.pipeline = pipeline(
                    task, model=model, tokenizer=path[1], device=device, model_kwargs=modelargs, **kwargs
                )
            else:
                self.pipeline = pipeline(task, model=path, device=device, model_kwargs=modelargs, **kwargs)

            if deviceid == -1 and quantize:
                # pylint: disable=E1101
                self.pipeline.model = self.quantize(self.pipeline.model)

        Models.checklength(self.pipeline.model, self.pipeline.tokenizer)

    def parseargs(self, **kwargs):
        """
        Inspects the pipeline method and splits kwargs into model args and pipeline args.

        Args:
            kwargs: all keyword arguments

        Returns:
            (model args, pipeline args)
        """

        args = inspect.getfullargspec(pipeline).args

        dtype = kwargs.get("torch_dtype")
        if dtype and isinstance(dtype, str) and dtype != "auto":
            kwargs["torch_dtype"] = Resolver()(dtype)

        return (
            {arg: value for arg, value in kwargs.items() if arg not in args},
            {arg: value for arg, value in kwargs.items() if arg in args},
        )

    def maxlength(self):
        """
        Gets the max length to use for generate calls.

        Returns:
            max length
        """

        return Models.maxlength(self.pipeline.model, self.pipeline.tokenizer)
