import math
import os
from typing import Any, Dict, List

import torch
import torch.distributed as dist
from datasets import Dataset

from veomni.utils import helper
from veomni.utils.helper import get_cache_dir


logger = helper.create_logger(__name__)


class DummyDataset:
    def __init__(self, size=100, num_shard=2, dataset_name: str = "test_dataset") -> None:
        self.size = size
        self.num_shard = num_shard

        self.save_path = get_cache_dir(f"./{dataset_name}")

        if not dist.is_initialized() or dist.get_rank() == 0:
            self.build_dummy_dataset()

        if dist.is_initialized():
            dist.barrier()

    def generate_data(self, index_list: List):
        for index in index_list:
            input_ids = [index + 1] * (index + 1)
            yield {"input_ids": input_ids, "attention_mask": [1] * len(input_ids), "labels": input_ids}

    def build_dummy_dataset(self):
        if not os.path.exists(self.save_path):
            os.makedirs(self.save_path)

        batch_len = math.ceil(self.size / self.num_shard)
        print(f"Total length: {self.size}, batch length: {batch_len}")

        index = 0
        for i in range(0, self.size, batch_len):
            print(f"Generating {index}th parquet file")
            ds = Dataset.from_generator(
                self.generate_data,
                gen_kwargs={"index_list": list(range(i + 1, i + batch_len + 1))},
                keep_in_memory=True,
                num_proc=1,
            )
            ds.to_parquet(os.path.join(self.save_path, f"{index}.parquet"))
            index += 1

    def clean_cache(self):
        if not dist.is_initialized() or dist.get_rank() == 0:
            if os.path.exists(self.save_path):
                os.system(f"rm -rf {self.save_path}")

    def __del__(self):
        self.clean_cache()


def process_dummy_example(
    example: Dict[str, Any],
    max_seq_len: int,
    source_name: str = None,
) -> List[Dict[str, "torch.Tensor"]]:
    tokenized_example = {}
    for k, v in example.items():
        if k == "ds_idx" or k == "source_name":
            continue
        else:
            tokenized_example[k] = torch.tensor(v[:max_seq_len], dtype=torch.long)
    return [tokenized_example]


class FakeModel:
    def __init__(self) -> None:
        pass

    def state_dict(self):
        return {}

    def load_state_dict(self, *args, **kwargs):
        pass


def compare_items(item, rank, group_size, group):
    item = item.to("cuda")
    item_list = [torch.empty_like(item) for _ in range(group_size)]

    dist.all_gather(item_list, item, group=group)

    for i in range(0, group_size):
        if not torch.equal(item, item_list[i]):
            logger.info(f"[rank{rank}]: group_rank {i} item is not equal to item {rank}")
            return False

    return True


def compare_global_batch(global_batch_list, global_batch_resume_list):
    for global_batch, global_batch_resume in zip(global_batch_list, global_batch_resume_list):
        for micro_batch, micro_batch_resume in zip(global_batch, global_batch_resume):
            for key in micro_batch.keys():
                if torch.is_tensor(micro_batch[key]):
                    assert torch.all(micro_batch[key] == micro_batch_resume[key])


def compare_metrics(metrics, metrics_resume):
    assert metrics["consume_tokens(M)"] == metrics_resume["consume_tokens(M)"]