import csv
import os

def process_file(gold, pred, benchmark_set_path):
    # init glod list
    glist = []
    for l in gold:
        if len(l.strip()) != 0:
            glist.append(l.strip().split('\t')[0])

    # init pred list
    plist = []
    for l in pred:
        if len(l.strip()) != 0:
            plist.append(l.strip().split('\t')[0])

    assert len(plist) == len(glist), "number of sessions must equal, len(plist)={}, len(glist)={}".format(len(plist), len(glist))

    q_list = process_benchmark(benchmark_set_path)
    return zip(glist, plist, q_list["question_list"], q_list["schema_list"])

def process_benchmark(benchmark_set_path):
    # init question and schema
    with open(benchmark_set_path) as f:
        question_list = []
        schema_list = []
        gold_list = []
        prompt_list = []
        db_id_list = []
        reader = csv.DictReader(
            f, fieldnames=["db_id", "question", "sql", "prompt"])
        next(reader)
        for line in reader:
            question_list.append(line["question"])
            schema_list.append(line["prompt"].split(
                "# Question:")[0].strip('\n'))
            gold_list.append(line["sql"] + "\t" + line["db_id"])
            db_id_list.append(line["db_id"])
            prompt_list.append(line["prompt"])
    return {
        "db_id_list": db_id_list,
        "schema_list": schema_list,
        "question_list": question_list,
        "prompt_list": prompt_list,
        "gold_list": gold_list
    }

if __name__ == "__main__":
    dirname, _ = os.path.split(__file__)
    result = process_benchmark("./../../scripts/evaluation/benchmark.csv".format(dirname))
    for gold in result["question_list"]:
        print('-- gold --')
        print(gold)