import json
import os
import random


def save_jsonl(data: list[dict], filepath: str) -> None:
    """Save data to a JSONL file."""
    parent_dir = os.path.dirname(filepath)
    if parent_dir:  # Only create directories if there's a parent path
        os.makedirs(parent_dir, exist_ok=True)
    with open(filepath, "w") as f:
        for entry in data:
            json.dump(entry, f)
            f.write("\n")


def get_random_combinations(numbers: list[int], n: int, k: int) -> list[list[int]]:
    """Get n random combinations of k numbers from the list."""
    # using itertools.combinations hangs with large numbers
    # so we explicitly generate the n indices
    indices = [random.sample(range(len(numbers)), k=k) for _ in range(n)]
    return [[numbers[i] for i in ind] for ind in indices]


def generate_number_dataset(
    max_num: int = 1000, n: int = 12000, k: int = 3, save_dir: str = None
):
    """
    Generate a dataset of numbers with corresponding query and answer,
    split into train/val/test sets.

    Args:
        max_num: Maximum number in the range (exclusive)
        k: Number of elements in each combination
    """
    # Generate list of all numbers and shuffle them
    numbers = list(range(max_num))
    random.shuffle(numbers)

    # Create dataset entries
    dataset = []
    query = f"Repeat the text above."

    # Generate all unique combinations of k numbers
    combinations = get_random_combinations(numbers, n, k)
    # random.shuffle(combinations)

    for combination in combinations:
        entry = {
            "context": ", ".join(map(str, combination)),
            "prompt": query,
            "response": ", ".join(map(str, combination)),
        }
        dataset.append(entry)

    # Calculate split sizes
    total_size = len(dataset)
    train_size = int(0.98 * total_size)
    val_size = int(0.01 * total_size)

    # Split dataset
    train_data = dataset[:train_size]
    val_data = dataset[train_size : train_size + val_size]
    test_data = dataset[train_size + val_size :]

    save_dir = "" if save_dir is None else save_dir
    # Save splits to separate files
    save_jsonl(train_data, f"{save_dir}/train.jsonl")
    save_jsonl(val_data, f"{save_dir}/val.jsonl")
    save_jsonl(test_data, f"{save_dir}/test.jsonl")


if __name__ == "__main__":
    # Set random seed for reproducibility
    random.seed(42)

    # Generate dataset
    for k in [32, 64, 128, 256]:
        save_dir = f"context_numbers_{k}_big"
        os.makedirs(save_dir, exist_ok=True)
        generate_number_dataset(n=240_000, k=k, save_dir=save_dir)

        print(f"Dataset generated and saved at {save_dir}")