#!/usr/bin/env python
# coding=utf-8
"""This Python code defines a class Dataset with methods for initializing, loading,
and manipulating datasets from different backends such as Hugging Face and JSON.
 
The `Dataset` class includes methods for loading datasets from a dictionary and a Hugging
Face dataset, mapping datasets, and retrieving the backend dataset and arguments.
"""



# Importing necessary libraries and modules
import copy
import json
import logging

from cmath import e
from pathlib import Path
from typing import Optional, List

from datasets import load_dataset, load_from_disk
from datasets import Dataset as HFDataset
from tqdm import tqdm

from lmflow.args import DatasetArguments
from lmflow.utils.constants import (
    DATASET_DESCRIPTION_MAP,
    TEXT_ONLY_DATASET_DESCRIPTION,
    TEXT2TEXT_DATASET_DESCRIPTION,
    FLOAT_ONLY_DATASET_DESCRIPTION,
    INSTANCE_FIELDS_MAP,
)
from lmflow.utils.versioning import is_multimodal_available
from lmflow.utils.data_utils import get_dataset_type_fast, check_dataset_instances_key_fast

if is_multimodal_available():
    from .multi_modal_dataset import CustomMultiModalDataset


logger = logging.getLogger(__name__)


DATASET_TYPES = [
    "text_only",
    "text2text",
    "float_only",
    "image_text",
    "conversation",
    "paired_conversation",
    "paired_text_to_text",
    "text_to_textlist",
    "text_to_scored_textlist",
    # internal marker for already-tokenized datasets loaded from disk
    "pretokenized",
]

KEY_TYPE = "type"
KEY_INSTANCES = "instances"
KEY_SCORE = "score"

class Dataset:
    r"""
    Initializes the Dataset object with the given parameters.

    Parameters
    ------------
    data_args : DatasetArguments object.
        Contains the arguments required to load the dataset.

    backend : str,  default="huggingface"
        A string representing the dataset backend. Defaults to "huggingface".
    
    args : Optional.
        Positional arguments.
    
    kwargs : Optional.
        Keyword arguments.
    """
    def __init__(self, data_args: DatasetArguments=None, backend: str="huggingface", *args, **kwargs):
        self.data_args = data_args
        self.backend = backend
        self.backend_dataset = None
        self.type = None        # Original type of the dataset
        self.dataset_path = data_args.dataset_path

        # Support cache-only mode: try loading a dataset from disk if dataset_path is None
        if data_args.dataset_path is None:
            cache_dir = getattr(data_args, "dataset_cache_dir", None)
            if cache_dir is not None:
                try:
                    ds = load_from_disk(cache_dir)
                    self.backend_dataset = ds
                    # Heuristic: mark as pretokenized if tokenized columns exist
                    try:
                        cols = set(ds.column_names)
                        if {"input_ids", "attention_mask"}.issubset(cols):
                            self.type = "pretokenized"
                    except Exception:
                        pass
                    return
                except Exception as ex:
                    logger.warning(f"Failed to load dataset from disk at {cache_dir}: {ex}")
            return

        if backend == "huggingface":
            data_files = [
                x.absolute().as_posix()
                 for x in Path(self.dataset_path).glob("*.json")
            ]
            logger.info(f"Data files: \n{data_files}")
            
            # check if the dataset is in the correct format and get the dataset type (text_only, text2text, etc.)
            self._check_hf_json_format(data_files)
            # Load the dataset using the HuggingFace dataset library
            logger.info('Loading datasets')
            extensions = "json"
            raw_dataset = load_dataset(
                extensions,
                data_files=data_files,
                field=KEY_INSTANCES,
                split="train",
                cache_dir=data_args.dataset_cache_dir,
            )
            self.backend_dataset = raw_dataset
            self._check_instance_format()
        elif backend == "json":
            # TODO (@Jiachun)
            pass
        elif backend == "custom_multi_modal":
            # FIXME refactor the backend name
            if not is_multimodal_available():
                raise ValueError(
                    'Multimodal not available. Please install via `pip install -e ".[multimodal]"`'
                )
            raw_dataset = CustomMultiModalDataset(self.dataset_path, data_args)
            self.backend_dataset = raw_dataset
        elif backend == "webdataset_multi_modal":
            if not is_multimodal_available():
                raise ValueError(
                    'Multimodal not available. Please install via `pip install -e ".[multimodal]"`'
                )
            # Check if dataset_path is a VILA mixture name or a file path
            dataset_path_obj = Path(self.dataset_path) if self.dataset_path else None
            # Consider it a file path if:
            # 1. File exists, OR
            # 2. Has a file extension like .json or .wdsmeta (even if not exists yet)
            is_file_path = (
                dataset_path_obj is not None
                and (
                    dataset_path_obj.exists()
                    or dataset_path_obj.suffix in [".json", ".wdsmeta"]
                    or str(dataset_path_obj).endswith(".wdsmeta.json")
                )
            )
            
            print(f"[Dataset] dataset_path={self.dataset_path}, is_file_path={is_file_path}", flush=True)
            
            if is_file_path:
                # Peek at the JSON to determine its format:
                # - list  → LLaVA/ShareGPT4V conversation format → CustomMultiModalDataset
                # - dict with 'shardlist' → WebDataset metadata → CustomWebDatasetMultiModalDataset
                _is_llava_list = False
                try:
                    with open(self.dataset_path, 'r') as _f:
                        _meta_peek = json.load(_f)
                    if isinstance(_meta_peek, list):
                        _is_llava_list = True
                        print(
                            f"[Dataset] JSON is a list (LLaVA/ShareGPT4V format), "
                            f"routing to CustomMultiModalDataset instead of CustomWebDatasetMultiModalDataset",
                            flush=True,
                        )
                except Exception as _e:
                    print(f"[Dataset] WARNING: Could not peek at JSON format: {_e}", flush=True)

                if _is_llava_list:
                    raw_dataset = CustomMultiModalDataset(self.dataset_path, data_args)
                else:
                    # Standard .wdsmeta.json dict format
                    from .webdataset_multi_modal_dataset import CustomWebDatasetMultiModalDataset
                    raw_dataset = CustomWebDatasetMultiModalDataset(self.dataset_path, data_args)
                self.backend_dataset = raw_dataset
            else:
                # Mixture name: delegate to VILA's build_dataset
                # Set VILA_DATASETS env var BEFORE importing VILA modules
                import os
                try:
                    from llava.data.builder import build_dataset as vila_build_dataset
                    from llava.data.builder import register_datasets, register_mixtures
                    from llava.train.args import DataArguments as VILADataArguments
                    from llava.train.args import TrainingArguments as VILATrainingArguments
                    from llava.train.args import ModelArguments as VILAModelArguments
                    import llava.data.datasets_mixture as datasets_mixture
                except ImportError as e:
                    raise ImportError(
                        f"VILA mixture '{self.dataset_path}' requires VILA package. "
                        f"Ensure VILA is installed or in PYTHONPATH. Error: {e}"
                    )
                
                # Force re-register datasets with current VILA_DATASETS env var
                from llava.data import builder
                builder.DATASETS = register_datasets()
                builder.MIXTURES = register_mixtures()
                print(f"[Dataset] Registered {len(builder.DATASETS)} VILA datasets", flush=True)
                
                # Register VILA datasets and mixtures
                datasets_mixture.register_datasets_mixtures()
                
                # Create VILA-compatible args (minimal stub for now)
                # You may need to map LMFlow's data_args to VILA's DataArguments
                vila_data_args = VILADataArguments(
                    data_mixture=self.dataset_path,
                    is_multimodal=True,
                    lazy_preprocess=True,
                )
                # Create minimal training args without triggering distributed init
                # Use a simple namespace instead of full TrainingArguments to avoid device setup
                from types import SimpleNamespace
                vila_training_args = SimpleNamespace(
                    output_dir="output",
                    per_device_train_batch_size=1,
                    gradient_accumulation_steps=1,
                    world_size=1,
                    process_index=0,
                    local_rank=-1,
                    seq_parallel_size=1,
                )
                vila_model_args = VILAModelArguments(mm_use_bos_eos_tokens=False, unified_audio_encoder=False)
                
                # Get tokenizer and image_processor from data_args if available (VILA needs them)
                tokenizer = getattr(data_args, "tokenizer", None)
                if tokenizer is None:
                    raise ValueError(
                        "VILA mixture requires a tokenizer. "
                        "Please set data_args.tokenizer before initializing the dataset."
                    )
                
                # Attach image_processor to VILA data_args
                # Now VILA mm_utils.py is patched to support Qwen2.5-VL processor
                image_processor = getattr(data_args, "image_processor", None)
                if image_processor is None:
                    # Try to get from tokenizer
                    image_processor = getattr(tokenizer, "image_processor", None)
                
                if image_processor is None:
                    # Fallback: load default Qwen processor
                    from transformers import AutoProcessor
                    image_processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct")
                    print("[Dataset] Loaded default Qwen2.5-VL processor for VILA mixture", flush=True)
                
                vila_data_args.image_processor = image_processor
                
                vila_dataset = vila_build_dataset(
                    self.dataset_path,
                    vila_data_args,
                    vila_training_args,
                    tokenizer,
                    vila_model_args,
                )
                
                # Wrap VILA dataset to convert format from tokenized to messages
                # This allows using LMFlow's custom collators (e.g., MDM-aware collator)
                from lmflow.datasets.vila_dataset_adapter import VILADatasetAdapter
                adapted_dataset = VILADatasetAdapter(vila_dataset)
                self.backend_dataset = adapted_dataset
        else:
            raise NotImplementedError(f'Unsupported dataset backend "{backend}"')

    
    def __len__(self):
        return len(self.backend_dataset)
    

    def _check_instance_format(self):
        """
        Checks if data (instances) have required fields. 
        Raises messages with hints if not matched.
        """
        fields = self.backend_dataset.features
        correct_fields = INSTANCE_FIELDS_MAP[self.type]
        if not set(correct_fields).issubset(set(fields)):
            raise ValueError(
                f'data instance fields incorrect'
                f' {list(correct_fields)} are required.'
            )
            
    
    def _check_hf_json_format(self, data_files: List[str]):
        for single_file in tqdm(data_files, desc='Checking dataset keys'):
            # get type and check if it is consistent
            json_data_type = get_dataset_type_fast(single_file)
            if not json_data_type:
                raise ValueError(
                    f'"{KEY_TYPE}" must be provided to initialize a dataset,'
                    f' e.g.\n'
                    f'    {TEXT_ONLY_DATASET_DESCRIPTION}'
                )
            if self.type is None:
                self.type = json_data_type
            elif self.type != json_data_type:
                raise ValueError(
                    'All task files must have same data types. Previous'
                    f' files have type "{self.type}", but in file'
                    f' {single_file}, it has type "{self.type}".'
                )
            # check if instances key is provided
            key_instances_exists_flag = check_dataset_instances_key_fast(single_file, KEY_INSTANCES)
            if not key_instances_exists_flag:
                raise ValueError(
                    f'"{KEY_INSTANCES}" must be provided to initialize a'
                    f' dataset, e.g.\n'
                    f'    {TEXT_ONLY_DATASET_DESCRIPTION}'
                )


    def from_dict(self, dict_obj: dict, *args, **kwargs):
        r"""
        Create a Dataset object from a dictionary.

        Return a Dataset given a dict with format:
            {
                "type": TYPE,
                "instances": [
                    {
                        "key_1": VALUE_1.1,
                        "key_2": VALUE_1.2,
                        ...
                    },
                    {
                        "key_1": VALUE_2.1,
                        "key_2": VALUE_2.2,
                        ...
                    },
                    ...
                ]
            }

        Parameters
        -----------

        dict_obj : dict.
            A dictionary containing the dataset information.
        
        args : Optional.
            Positional arguments.
        
        kwargs : Optional.
            Keyword arguments.

        Returns
        ---------

        self : Dataset object.
        """
        if self.backend == "huggingface":
            if KEY_TYPE not in dict_obj:
                raise ValueError(
                    f'"{KEY_TYPE}" must be provided to initialize a dataset,'
                    f' e.g.\n'
                    f'    {TEXT_ONLY_DATASET_DESCRIPTION}'
                )
            if KEY_INSTANCES not in dict_obj:
                raise ValueError(
                    f'"{KEY_INSTANCES}" must be provided to initialize a'
                    f' dataset, e.g.\n'
                    f'    {TEXT_ONLY_DATASET_DESCRIPTION}'
                )

            self.type = dict_obj[KEY_TYPE]
            if not self.type in INSTANCE_FIELDS_MAP:
                raise ValueError(f'type "{self.type}" is not supported')

            correct_fields = INSTANCE_FIELDS_MAP[self.type]

            for i, instance in enumerate(dict_obj[KEY_INSTANCES]):
                fields = instance.keys()
                if not set(correct_fields).issubset(set(fields)):
                    raise ValueError(
                        f'data instance fields incorrect'
                        f' {list(correct_fields)} are required.'
                    )

            try:
                hf_dict = {}
                if len(dict_obj[KEY_INSTANCES]) > 0:
                    for key in dict_obj[KEY_INSTANCES][0].keys():
                        hf_dict[key] = [
                            instance[key] for instance in dict_obj[KEY_INSTANCES]
                        ]

                self.backend_dataset = HFDataset.from_dict(hf_dict, *args, **kwargs)
            except AttributeError as ex:
                raise ValueError(
                    f"Error occurs: {ex}. Failed to convert dict to"
                    f" \"{self.type}\" dataset," f" the standard format is as"
                    f" follows:\n"
                    f"    {DATASET_DESCRIPTION_MAP[self.type]}"
                )
            self._check_instance_format()

            return self
        elif self.backend == "dict":
            self.backend_dataset = dict_obj
            self.type = dict_obj[KEY_TYPE]
            return self
        else:
            raise NotImplementedError(
                f'Currently .from_dict is not supported for backend "{self.backend}"'
            )


    @classmethod
    def create_from_dict(cls, dict_obj, *args, **kwargs):
        r"""
        Returns
        --------

        Returns a Dataset object given a dict.
        """
        empty_data_args = DatasetArguments(dataset_path=None)
        dataset = Dataset(empty_data_args)
        return dataset.from_dict(dict_obj)


    def to_dict(self):
        r"""
        Returns
        ---------

        Return a dict represents the dataset:
            {
                "type": TYPE,
                "instances": [
                    {
                        "key_1": VALUE_1.1,
                        "key_2": VALUE_1.2,
                        ...
                    },
                    {
                        "key_1": VALUE_2.1,
                        "key_2": VALUE_2.2,
                        ...
                    },
                    ...
                ]
            }

        A python dict object represents the content of this dataset.
        """
        if self.backend == "huggingface":
            dict_obj = {}
            dict_obj[KEY_TYPE] = self.get_type()
            hf_dict = self.backend_dataset.to_dict()
            dict_obj[KEY_INSTANCES] = []

            first_key = None
            for key in hf_dict.keys():
                first_key = key
                break

            if first_key is not None:
                num_instances = len(hf_dict[first_key])
                dict_obj[KEY_INSTANCES] = [
                    {
                        key: hf_dict[key][i] for key in hf_dict.keys()
                    }
                    for i in range(num_instances)
                ]

            return dict_obj
        elif self.backend == "dict":
            dict_obj = self.backend_dataset
            return dict_obj
        else:
            raise NotImplementedError(
                f'Current .to_dict is not supported for backend "{self.backend}"'
            )


    def to_list(self):
        """Returns a list of instances."""
        if self.backend == "huggingface":
            instance_list = [self.backend_dataset.__getitem__(idx)
                             for idx in range(len(self.backend_dataset))]
            return instance_list
        elif self.backend == "dict":
            instance_list = copy.deepcopy(self.backend_dataset[KEY_INSTANCES])
            # TODO: should be a list of instances, instance should be huggingface datasets row format
            return instance_list
        else:
            raise NotImplementedError(
                f'Current .to_list is not supported for backend "{self.backend}"'
            )


    def map(self, *args, **kwargs):
        r"""
        Parameters
        ------------
        args : Optional.
            Positional arguments.
        
        kwargs : Optional.
            Keyword arguments.

        Returns
        ---------

        self : Dataset object.
        """
        # If the dataset uses Hugging Face as the backend, 
        # call the `map()` function of the Hugging Face backend dataset
        if self.backend == "huggingface":
            # Set the mapped dataset as the backend dataset of the current dataset
            mapped_backend_dataset = self.backend_dataset.map(*args, **kwargs)
            self.backend_dataset = mapped_backend_dataset
            return self
        else:
            # If the backend is not Hugging Face, raise a NotImplementedError
            raise NotImplementedError(
                f'Currently .map is not supported for backend "{self.backend}"'
            )


    def get_backend(self) -> Optional[str]:
        r"""
        Returns
        ---------

        self.backend
        """
        return self.backend


    def get_backend_dataset(self):
        r"""
        Returns
        ---------

        self.backend_dataset
        """
        return self.backend_dataset


    def get_fingerprint(self):
        r"""
        Returns
        ---------

        Fingerprint of the backend_dataset which controls the cache
        """
        return self.backend_dataset._fingerprint

    
    def get_data_args(self):
        r"""
        Returns
        ---------

        self.data_args
        """
        return self.data_args


    def get_type(self) -> str:
        r"""
        Returns
        ---------

        self.type
        """
        return self.type
    
    
    def save(
        self, 
        file_path: str, 
        format: str="json"
    ):
        r"""
        Save the dataset to a json file.

        Parameters
        ------------
        file_path : str.
            The path to the file where the dataset will be saved.
        """
        if format == "json":
            assert Path(file_path).suffix == ".json", "The file path must have a .json extension."
            with open(file_path, "w", encoding='utf-8') as fout:
                json.dump(self.to_dict(), fout, indent=4, ensure_ascii=False)
                
        else:
            logger.error(f"Unsupported format when saving the dataset: {format}.")
        
            
    def sample(self, n: int, seed: int=42):
        r"""
        Sample n instances from the dataset.

        Parameters
        ------------
        n : int.
            The number of instances to sample from the dataset.

        Returns
        ---------

        sample_dataset : Dataset object.
            A new dataset object containing the sampled instances.
        """
        if self.backend == "huggingface":
            sampled_dataset = self.backend_dataset.shuffle(seed=seed).select(range(n))
            output_dataset = self.create_from_dict(
                {
                    "type": self.get_type(),
                    "instances": [
                        {
                            col_name: sampled_dataset[col_name][i] for col_name in sampled_dataset.column_names
                        } for i in range(n)
                    ]
                }
            )
            return output_dataset
        else:
            raise NotImplementedError(
                f'Currently .sample is not supported for backend "{self.backend}"'
            )
            
            
    def train_test_split(self, test_size: float=0.2, shuffle: bool=True, seed: int=42):
        r"""
        Split the dataset into training and testing sets.

        Parameters
        ------------
        test_size : float, default=0.2.
            The proportion of the dataset that will be used for testing.

        Returns
        ---------

        train_dataset : Dataset object.
            A new dataset object containing the training instances.
        
        test_dataset : Dataset object.
            A new dataset object containing the testing instances.
        """
        if self.backend == "huggingface":
            splited = self.backend_dataset.train_test_split(
                test_size=test_size, shuffle=shuffle, seed=seed
            )
            train_dataset = self.create_from_dict(
                {
                    "type": self.get_type(),
                    "instances": [
                        {
                            col_name: splited["train"][col_name][i] for col_name in splited["train"].column_names
                        } for i in range(len(splited["train"]))
                    ]
                }
            )
            test_dataset = self.create_from_dict(
                {
                    "type": self.get_type(),
                    "instances": [
                        {
                            col_name: splited["test"][col_name][i] for col_name in splited["test"].column_names
                        } for i in range(len(splited["test"]))
                    ]
                }
            )
            return train_dataset, test_dataset
        else:
            raise NotImplementedError(
                f'Currently .train_test_split is not supported for backend "{self.backend}"'
            )
            
            
    def drop_instances(self, indices: list):
        r"""
        Drop instances from the dataset.

        Parameters
        ------------
        indices : list.
            A list of indices of the instances to drop from the dataset.
        """
        if self.backend == "huggingface":
            self.backend_dataset = self.backend_dataset.remove_indices(indices)
        else:
            raise NotImplementedError(
                f'Currently .drop_instances is not supported for backend "{self.backend}"'
            )
            
    
    def sanity_check(
        self, 
        drop_invalid: bool=True,
    ):
        r"""
        Perform a sanity check on the dataset.
        """
        if self.backend == "huggingface":
            self.hf_dataset_sanity_check(drop_invalid)
        else:
            raise NotImplementedError(
                f'Currently .sanity_check is not supported for backend "{self.backend}"'
            )
            
            
    def hf_dataset_sanity_check(
        self,
        drop_invalid: bool=True,
    ):
        r"""
        Perform a sanity check on the HuggingFace dataset.
        """
        if self.backend_dataset is None or len(self.backend_dataset) == 0:
            raise ValueError("Dataset is empty.")

        if self.type == 'text_to_textlist':
            num_output_per_instance = len(self.backend_dataset['output'][0])
            dataset_cache = self.backend_dataset.filter(lambda x: len(x['input'])!=0)
            dataset_cache = self.backend_dataset.filter(lambda x: len(x['output']) == num_output_per_instance)
            dataset_cache = self.backend_dataset.filter(lambda x: not all([len(output) == 0 for output in x['output']]))
            
            if len(dataset_cache) != len(self.backend_dataset):
                warning_info = (
                    f"Found {len(self.backend_dataset) - len(dataset_cache)} invalid instances "
                    "during hf_dataset_sanity_check, please check:\n"
                    "   1. length of input strings should not be empty\n"
                    "   2. length of output strings should not be all empty\n"
                    "   3. number of output strings should be consistent\n" # since we will use tensor reshape later
                )
                if drop_invalid:
                    self.backend_dataset = dataset_cache
                    logger.warning(warning_info+"Invalid instances are dropped.")
                else:
                    raise ValueError(warning_info)
        
        else:
            logger.warning(f"No sanity check for {self.type} dataset.")