import json
import os

import tqdm

from Scripts.config_experiment.components.regenNL_scorer import regen
from Scripts.Prompting import *
from Scripts.Utils import *

prompt_file = r"Scripts\config_experiment\prompts\info2NL.md"

with open(prompt_file, "r") as f:
    PROMPT = f.read()


def ambiguous_0(base_folder, args):
    dataset_folder = os.path.join(base_folder, args.dataset + "_ambiguous_0.json")
    original_folder = os.path.join("Responses", "original", args.dataset + ".json")
    seen=[]
    if os.path.exists(dataset_folder):
        with open(dataset_folder, "r") as f:
            subset = json.load(f)
        seen=[d["task_id"] for d in subset]

    data = json.load(open(original_folder, "r"))

    try:
        with open(dataset_folder, "r") as f:
            subset = json.load(f)
    except:
        subset = []
    for i, d in enumerate(tqdm.tqdm(data)):
        if i == args.sample: break
        if d['task_id'] in seen: continue
        seen.append(d['task_id'])
        d["subtask_id"] = 0
        # benchmark = 1
        args.info_prompt = "agg_comp_imp_fewshot"
        tries = 3
        while tries:
            try:
                info, _ = get_info(
                    d["text"], args, d["test_list"], 0
                )
                break
            except Exception as e:
                print(e)
                tries -= 1
                if tries == 0:
                    print(f"Failed to get info")
                    brk
        d["info"] = info

        # benchmark = 2
        args.info_prompt = "agg_comp_imp_fewshot_examples"
        tries = 3
        while tries:
            try:
                info, _ = get_info(
                    d["text"], args, d["test_list"], 0
                )
                break
            except Exception as e:
                print(e)
                tries -= 1
                if tries == 0:
                    print(f"Failed to get info")
                    brk
        d["info_examples"] = info

        subset.append(d)

        with open(dataset_folder, "w") as f:
            json.dump(subset, f)

    regen(PROMPT, dataset_folder, args.llm)
    return


def init_k(base_folder, args, iteration):
    # initialize empty lst of json as files in folder
    dataset_folder = os.path.join(
        base_folder, str(args.benchmark)+ "_" + args.dataset + f"_ambiguous_{iteration}.json"
    )
    if os.path.exists(dataset_folder):
        print(f"Responses already exists at {dataset_folder}")
        return

    temp = []
    with open(dataset_folder, "w") as f:
        json.dump(temp, f)
    return


def cache_init(raw_args=None):
    args = get_args(raw_args)
    args.dataset = "mbpp-updated_corrected_50"
    base_folder = os.path.join("Responses", "cache")

    # initialize all folders
    create_folder_if_not_exists(base_folder)
    
    ambiguous_0(base_folder, args)
    
    # args.benchmark=1 
    # for i in range(1, args.max_iter + 1):
    #     init_k(base_folder, args=args, iteration=i)
    # args.benchmark=2
    # for i in range(1, args.max_iter + 1):
    #     init_k(base_folder, args=args, iteration=i)

    return


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