# coding=utf-8
"""Tokenization classes for CoEncoder"""

from typing import List, Union, Optional
from transformers import AutoTokenizer
from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
from transformers.utils import logging
from transformers.feature_extraction_utils import BatchFeature

logger = logging.get_logger(__name__)


class CoEncoderDualTokenizerKwargs(ProcessingKwargs, total=False):
    _defaults = {
        "context_kwargs": {
            "padding": False,
        },
        "text_kwargs": {
            "padding": False,
        },
    }


class CoEncoderDualTokenizer(ProcessorMixin):
    r"""
    CoEncoderDualTokenizer is tokenizer for the CoEncoder model. It processes context and main text.

    Args:
        context_tokenizer ([`PreTrainedTokenizer`]):
            The tokenizer for context.
        text_tokenizer ([`PreTrainedTokenizer`]):
            The tokenizer for main text.
    """

    attributes = ["context_tokenizer", "text_tokenizer"]
    context_tokenizer_class = "AutoTokenizer"
    text_tokenizer_class = "AutoTokenizer"

    def __init__(self, context_tokenizer=None, text_tokenizer=None):
        super().__init__(context_tokenizer, text_tokenizer)
    
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
        """
        Load both context and text tokenizers from a given repository.

        Args:
            pretrained_model_name_or_path (str): The name or path of the Hugging Face repository.

        Returns:
            CoEncoderDualTokenizer: An instance of the tokenizer class.
        """
        # Load context_tokenizer from 'context_tokenizer' directory
        context_tokenizer = AutoTokenizer.from_pretrained(f"{pretrained_model_name_or_path}/context_tokenizer", **kwargs)

        # Load text_tokenizer from 'text_tokenizer' directory
        text_tokenizer = AutoTokenizer.from_pretrained(f"{pretrained_model_name_or_path}/text_tokenizer", **kwargs)

        # Return a new instance of CoEncoderDualTokenizer with both tokenizers loaded
        return cls(context_tokenizer=context_tokenizer, text_tokenizer=text_tokenizer)

    def __call__(
        self,
        context: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
        return_tensors: Optional[str] = None,
        **kwargs: Unpack[CoEncoderDualTokenizerKwargs]
    ) -> BatchFeature:
        """
        Main method to prepare inputs for the CoEncoder model.

        Args:
            context: Context text input.
            text: Main text input.
            return_tensors: Type of tensors to return.

        Returns:
            BatchFeature: A BatchFeature object containing model inputs.
        """
        if context is None and text is None:
            raise ValueError("You must provide either context or text.")

        features = {}

        if context is not None:
            context_features = self.context_tokenizer(
                context,
                return_tensors=return_tensors,
                **kwargs.get("context_kwargs", {})
            )
            features.update({f"context_{k}": v for k, v in context_features.items()})

        if text is not None:
            text_features = self.text_tokenizer(
                text,
                return_tensors=return_tensors,
                **kwargs.get("text_kwargs", {})
            )
            features.update({k: v for k, v in text_features.items()})

        return BatchFeature(data=features, tensor_type=return_tensors)

    def pad(
        self,
        encoded_inputs,
        padding=True,
        max_length=None,
        return_tensors=None,
        **kwargs
    ):
        """
        Pads the encoded inputs to the maximum length in the batch.

        Args:
            encoded_inputs: A list of dictionaries containing context and text features.
            padding: Whether to pad sequences.
            max_length: Maximum length for padding.
            return_tensors: Type of tensors to return.

        Returns:
            A dictionary with padded sequences.
        """
        # Separate context and text features
        context_features = []
        text_features = []

        for feature in encoded_inputs:
            # Extract context features
            context_feature = {
                k[len("context_"):]: v
                for k, v in feature.items()
                if k.startswith("context_")
            }
            if context_feature:
                context_features.append(context_feature)
            # Extract text features
            text_feature = {
                k[len("input_"):]: v
                for k, v in feature.items()
                if k.startswith("input_")
            }
            if text_feature:
                text_features.append(text_feature)

        # Pad context features
        if context_features:
            context_padded = self.context_tokenizer.pad(
                context_features,
                padding=padding,
                max_length=max_length,
                return_tensors=return_tensors,
                **kwargs.get("context_kwargs", {})
            )
            context_padded = {f"context_{k}": v for k, v in context_padded.items()}
        else:
            context_padded = {}

        # Pad text features
        if text_features:
            text_padded = self.text_tokenizer.pad(
                text_features,
                padding=padding,
                max_length=max_length,
                return_tensors=return_tensors,
                **kwargs.get("text_kwargs", {})
            )
            text_padded = {k: v for k, v in text_padded.items()}
        else:
            text_padded = {}

        # Combine padded features
        padded_features = {**context_padded, **text_padded}

        return BatchFeature(data=padded_features, tensor_type=return_tensors)

    def batch_decode(self, *args, **kwargs):
        """
        Calls the batch_decode method of the text_tokenizer.
        """
        return self.text_tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        """
        Calls the decode method of the text_tokenizer.
        """
        return self.text_tokenizer.decode(*args, **kwargs)

    @property
    def model_input_names(self):
        """
        Returns the model input names.
        """
        return list(dict.fromkeys(self.context_tokenizer.model_input_names + self.text_tokenizer.model_input_names))