import json
import os

from Scripts.Utils import *


def get_dataset(input_file, output_file):
    with open(input_file, "r") as f:
        data = json.load(f)

    gen = []
    masked_info = []
    vals = ["task_id", "code", "test_list", "info"]#, "info_count"]

    for d in data:
        for ids, t in enumerate(d["text_options"]):
            temp = {}
            temp["subtask_id"] = ids
            for val in vals:
                if val in d.keys():
                    temp[val] = d[val]
            temp["maskinfo"] = t[0]
            temp["mask_key"] = t[1]
            temp["text"] = t[2]
            temp["orig_text"] = d["text"]
            if temp["maskinfo"] not in masked_info:
                gen.append(temp)
                masked_info.append(temp["maskinfo"])

    with open(output_file, "w") as f:
        json.dump(gen, f)
    return


def make_dataset(raw_args=None, iter=0, baseline=0):
    args = get_args(raw_args)
    args.benchmark = args.dataset.split("_")[0]
    folder = f"config-{args.sample}sample_{args.n_values}n_{args.threshold}threshold_{args.temp}temp"
    if not baseline:
        base_folder = os.path.join(
            "Responses", args.llm, args.info_prompt, "all_temp_variation", folder
        )
    else:
        base_folder = os.path.join(
            "Responses", args.llm, args.info_prompt, args.dataset, args.prompt, folder
        )
    input_file = os.path.join(
        base_folder, "min_spec_candidates", f"{args.dataset}_ambiguous_{iter+1}.json"
    )
    output_file = os.path.join(
        base_folder, "dataset", f"{args.dataset}_ambiguous_{iter+1}.json"
    )

    if not os.path.exists(input_file):
        return 0

    get_dataset(input_file, output_file)
    return len(json.load(open(output_file, "r")))


if __name__ == "__main__":
    # try:
    make_dataset(raw_args=None)
# except Exception as e:
#     import pdb
#     pdb.post_mortem()
