import csv
import glob
import json


def filter_all_cases_checked_result():
    for i in glob.glob("../cases*_st_sqlite.csv"):
        yield i


def load_din_sql():
    results = {}

    with open("/home/jiajun/Downloads/spider/Few-shot-NL2SQL-with-prompting/data/output_dev_gpt35.json") as f:
        for row in json.loads(f.read()):
            question = row["question"]
            database = row["db_id"]
            sql = row["sql"]

            results[database + question] = sql

    return results


def load_generated_sql():
    results = {}

    for i in filter_all_cases_checked_result():
        with open(i) as f:
            for row in csv.DictReader(f):
                question = row["question"]
                database = row.get("\ufeffdatabase") or row.get("database")
                assert database is not None
                generated_sql = row["generated_sql"]

                results[database + question] = generated_sql

    return results


def gen_gold_sql():
    generated_sqls = load_generated_sql()
    # din_sqls = load_din_sql()

    with open("./database/combine.json") as f:
        cases = json.loads(f.read())

    with open("./gen_dev_gold.sql", "w") as f1:
        with open("./gen_gen_gold.sql", "w") as f2:
            for i in cases:
                dbname = i["db_id"]
                sql = i["query"]
                question = i["question"]
                sql = sql.replace("\t", " ").replace("\n", " ").rstrip(';')

                key = dbname + question
                generated_sql = generated_sqls.get(key)
                if not generated_sql:
                    continue
                generated_sql = generated_sql.replace("\t", " ").replace("\n", " ").rstrip(';')

                f1.write("{}\t{}\n".format(sql, dbname))
                f2.write("{}\t{}\n".format(generated_sql, dbname))


if __name__ == "__main__":
    gen_gold_sql()
    # load_din_sql()