from typing import Optional, Union, List, Dict
import numpy as np
from transformers import AutoTokenizer
from kithara.dataset.utils import initialize_tokenizer
from kithara.model.hf_compatibility import get_model_name_from_preset_handle
from kithara.model.maxtext.conversion_utils import MaxTextConversionMixin
from kithara.model.maxtext.ckpt_compatibility import (
    save_maxtext_model_in_hf_format,
    load_hf_weights_into_maxtext_model,
)
from kithara.model import (
    Model,
    set_precision,
    set_global_model_implementation_type,
    ModelImplementationType,
)


class MaxTextModel(Model, MaxTextConversionMixin):
    """
    MaxTextModel is class that represents a MaxText model via the
    Kithara.Model interface. It is a thin wrapper around the underlying
    MaxText model instance.

    Methods
    -------
    from_random: Create a randomly initialized MaxText model with the
        given configuration.
    from_preset: Create a MaxText model initialized with weights from
        HuggingFace Hub.
    generate: Generate text based on the input tokens, with an option to
        stop at specific token IDs.
    save_in_hf_format: Save the MaxText model in HuggingFace format.
    """

    @classmethod
    def from_random(
        cls,
        model_name: str,
        seq_len: int = 2048,
        per_device_batch_size: int = 1,
        precision: str = "mixed_float16",
        scan_layers: bool = False,
        maxtext_config_args: Optional[dict] = None,
    ) -> "MaxTextModel":
        """Create a randomly initialized MaxText model with the given configuration.

        Args:
            model_name (str): Name of the MaxText model configuration to use.
            seq_len (int, optional): Maximum sequence length. Defaults to 2048.
            per_device_batch_size (int, optional): Batch size per device.
                Defaults to 1.
            precision (str, optional): Precision mode for computations.
                Defaults to "mixed_float16".
            scan_layers (bool, optional): Whether to use scan layers for memory efficiency.
                Defaults to False.
            maxtext_config_args (Optional[dict], optional): Additional configuration arguments.
                Defaults to None.

        Returns:
            MaxTextModel: A new instance of MaxTextModel with random initialization.
        """
        set_global_model_implementation_type(ModelImplementationType.MAXTEXT)

        set_precision(precision)
        weight_dtype = cls._weight_dtype(precision)
        activation_dtype = cls._activation_dtype(precision)

        sharding_strategy, model = cls.initialize_random_maxtext_model(
            model_name,
            seq_len,
            per_device_batch_size,
            weight_dtype,
            activation_dtype,
            scan_layers,
            maxtext_config_args,
        )

        return cls(
            model,
            sharding_strategy,
            model_name=model_name,
            precision=precision,
            scan_layers=scan_layers,
        )

    @classmethod
    def from_preset(
        cls,
        preset_handle: str,
        seq_len: int = 2048,
        per_device_batch_size: int = 1,
        precision: str = "mixed_float16",
        scan_layers: bool = False,
        maxtext_config_args: Optional[dict] = None,
    ) -> "MaxTextModel":
        """Create a MaxText model initialized with weights from HuggingFace Hub.

        Args:
            preset_handle (str): HuggingFace model identifier. This could be a
                HuggingFace Hub path (e.g "gs://google/gemma-2-2b), or
                a local HuggingFace checkpoint path (e.g. tmp/my_model/checkpoint), or
                a GCS HuggingFace checkpoint path (e.g. gs://bucket_name/my_model/checkpoint)
            seq_len (int): Maximum sequence length.
            per_device_batch_size (int): Batch size per device.
            precision (str, optional): Precision mode for computations.
                Defaults to "mixed_float16".
            scan_layers (bool, optional): Whether to use scan layers. Defaults to False.
            maxtext_config_args (Optional[dict], optional): Additional configuration arguments.
                Defaults to None.
            **kwargs: Additional keyword arguments.

        Returns:
            MaxTextModel: A new instance of MaxTextModel initialized with pretrained weights.
        """
        set_global_model_implementation_type(ModelImplementationType.MAXTEXT)

        set_precision(precision)
        weight_dtype = cls._weight_dtype(precision)
        activation_dtype = cls._activation_dtype(precision)

        model_name = get_model_name_from_preset_handle(preset_handle)
        sharding_strategy, model = cls.initialize_random_maxtext_model(
            model_name,
            seq_len,
            per_device_batch_size,
            weight_dtype,
            activation_dtype,
            scan_layers,
            maxtext_config_args,
        )
        model = load_hf_weights_into_maxtext_model(preset_handle, model, scan_layers)

        return cls(
            model,
            sharding_strategy,
            model_name=model_name,
            precision=precision,
            scan_layers=scan_layers,
        )

    def _convert_text_input_to_model_input(
        self,
        prompts: Union[str | List[str]],
        max_length: int = 100,
        tokenizer: Optional[AutoTokenizer] = None,
        tokenizer_handle: Optional[str] = None,
    ):
        tokenizer = (
            initialize_tokenizer(tokenizer_handle) if tokenizer is None else tokenizer
        )
        tokens: Dict[str, np.ndarray] = tokenizer(
            prompts,
            max_length=max_length,
            padding="max_length",
            padding_side="right",
            truncation=True,
            return_tensors="np",
        )
        input_ids = tokens["input_ids"]
        attention_mask = tokens["attention_mask"]

        B, S = input_ids.shape
        positions = np.arange(S, dtype=np.int32)[None, :]
        positions = np.repeat(positions, B, axis=0)
        return {
            "tokens": input_ids,
            "segment_ids": attention_mask,
            "positions": positions,
        }

    def _generate(
        self,
        inputs,
        **kwargs,
    ):
        """Note: Currently running inference MaxText models is slow
        since it does not perform kv caching. This should improve with integration
        with MaxEngine.
        """
        return super()._generate(
            inputs, **kwargs, tokens_key="tokens", padding_mask_key="segment_ids"
        )

    def save_in_hf_format(
        self, output_dir: str, dtype: str = "auto", parallel_threads=8
    ):
        """Save the model in HuggingFace format, including the model configuration file (`config.json`),
            the model weights file (`model.safetensors` for models smaller than
            `DEFAULT_MAX_SHARD_SIZE` and `model-x-of-x.safetensors` for larger models),
            and the safe tensors index file (`model.safetensors.index.json`).

        Args:
            output_dir (str): Directory path where the model should be saved.
                Directory could be local or a Google cloud storage path, and
                will be created if it doesn't exist.
            dtype (str, optional): Data type for saved weights. Defaults to "auto".
            parallel_threads (int, optional): Number of parallel threads to use for saving.
                Defaults to 8. Make sure the local system has at least
                `parallel_threads * DEFAULT_MAX_SHARD_SIZE` free disk space,
                as each thread will maintain a local cache of size `DEFAULT_MAX_SHARD_SIZE`.
        """
        save_maxtext_model_in_hf_format(
            self, output_dir, dtype=dtype, parallel_threads=parallel_threads
        )
