import torch
from transformers.processing_utils import ProcessorMixin


class MSPProcessor(ProcessorMixin):
    attributes = ["feature_extractor", "video_processor", "tokenizer"]
    feature_extractor_class = "MSPAudioFeatureExtractor"
    video_processor_class = "MSPVisualVideoProcessor"
    tokenizer_class = "AutoTokenizer"

    def __call__(self, images=None, text=None, videos=None, audio=None, **kwargs):

        if getattr(self, "feature_extractor", None) and audio is not None:
            if "sampling_rate" in kwargs:
                sampling_rate = kwargs["sampling_rate"]
            else:
                sampling_rate = self.feature_extractor.sampling_rate

            if isinstance(audio, (str, bytes)):
                audio = self.feature_extractor._load_audio(
                    audio, sample_rate=sampling_rate
                )
            elif (
                isinstance(audio, list) and audio and isinstance(audio[0], (str, bytes))
            ):
                audio = [
                    self.feature_extractor._load_audio(a, sample_rate=sampling_rate)
                    for a in audio
                ]

        if getattr(self, "video_processor", None) and videos is not None:
            if isinstance(videos, (str, bytes)):
                videos = self.video_processor._load_video(videos)
            elif (
                isinstance(videos, list)
                and videos
                and isinstance(videos[0], (str, bytes))
            ):
                videos = [self.video_processor._load_video(v) for v in videos]

        if getattr(self, "tokenizer", None) and text is not None:
            if isinstance(text, bytes):
                text = text.decode("utf-8")
            elif isinstance(text, list) and text and isinstance(text[0], bytes):
                text = [t.decode("utf-8") for t in text]

        outputs = super().__call__(
            images=images, text=text, videos=videos, audio=audio, **kwargs
        )

        if "device" in kwargs:
            device = kwargs["device"]
            for key, value in outputs.items():
                if isinstance(value, torch.Tensor):
                    outputs[key] = value.to(device)
                elif (
                    isinstance(value, list)
                    and value
                    and isinstance(value[0], torch.Tensor)
                ):
                    outputs[key] = [v.to(device) for v in value]

        return outputs