import os
from functools import partial
from multiprocessing import cpu_count

from datasets import load_dataset
from transformers import AutoTokenizer


def _prepare_sample(batch, gpt2_tokenizer):
    batch_input_ids = list()
    batch_attention_mask = list()

    for feature_name, feature_values in batch.items():
        if feature_name != "turns":
            continue

        for turns in feature_values:
            last_idx = len(turns) - 1

            for i in range(1, last_idx):
                context = turns[:i]
                context = " ".join(context).strip()

                response = turns[i].strip()

                input_ids, attention_mask, _ = _prepare_dialogue(
                    gpt2_tokenizer, context, response
                )

                batch_input_ids.append(input_ids)
                batch_attention_mask.append(attention_mask)

    return {
        "input_ids": batch_input_ids,
        "attention_mask": batch_attention_mask,
    }


def _prepare_dialogue(gpt2_tokenizer, context, response):
    context = gpt2_tokenizer.encode(context)
    response = gpt2_tokenizer.encode(response)
    total_len = int(len(context) + len(response))

    if total_len >= gpt2_tokenizer.model_max_length:
        num_tokens_to_remove = total_len - gpt2_tokenizer.model_max_length

        context, response, _ = gpt2_tokenizer.truncate_sequences(
            context, response, num_tokens_to_remove=num_tokens_to_remove
        )

        context[-1] = gpt2_tokenizer.eos_token_id
        input_ids = context + response
    else:
        input_ids = context + [gpt2_tokenizer.eos_token_id] + response

    assert len(input_ids) <= gpt2_tokenizer.model_max_length

    attention_mask = [1] * len(input_ids)
    return input_ids, attention_mask, response


def prepare_dataset():
    gpt2_tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
    print(gpt2_tokenizer)

    dstc8_datset = load_dataset("roskoN/dstc8-reddit-corpus", keep_in_memory=False)

    _prepare_sample_partial = partial(_prepare_sample, gpt2_tokenizer=gpt2_tokenizer)

    num_proc = cpu_count()

    for split_name, dataset_split in list(dstc8_datset.items()):
        path = f"./data/encoded_{split_name}"

        if os.path.exists(path):
            print(f"{path} exists, skipping...")
            continue

        print(f"Processing {split_name}")
        encoded_dataset_split = dataset_split.map(
            function=_prepare_sample_partial,
            batched=True,
            num_proc=num_proc,
            remove_columns=dataset_split.column_names,
            # batch_size=4,
            # writer_batch_size=4,
            keep_in_memory=False,
        )
        print(encoded_dataset_split)

        encoded_dataset_split.save_to_disk(path)