import csv
import glob
import sqlite3



def load_all_cases():
    for filename in glob.glob('cases_*st_sqlite.csv'):
        yield filename


def bird_execute_sql(predicted_sql,ground_truth, db_path):
    try:
        conn = sqlite3.connect(db_path)
        # Connect to the database
        cursor = conn.cursor()
        cursor.execute(predicted_sql)
        predicted_res = cursor.fetchall()
        cursor.execute(ground_truth)
        ground_truth_res = cursor.fetchall()
        res = 0
        if set(predicted_res) == set(ground_truth_res):
            res = 1
        return res
    except Exception as e:
        return str(e)


for i in load_all_cases():
    new_filename = i + ".new.bird.csv"
    with open(new_filename, "w") as f:
        writer = csv.writer(f)
        writer.writerow(["database", "question", "sql", "generated_sql", "difficulty", "equal", "note", "time_cost", "clarified_task", "raw_generated_sql", "refine_note", "data_context_id", "bird_eval"])

        reader = csv.DictReader(open(i))
        for row in reader:
            # {'database': 'card_games', 'question': 'What is the type of card "Benalish Knight"?', 'sql': "SELECT DISTINCT T1.type FROM cards AS T1 INNER JOIN foreign_data AS T2 ON T1.uuid = T2.uuid WHERE T1.name = 'Benalish Knight';", 'generated_sql': "SELECT `type` FROM `cards` WHERE `name` = 'Benalish Knight';;", 'difficulty': 'simple', 'equal': 'human check', 'note': '', 'time_cost': '15', 'clarified_task': "Find the type of the card named 'Benalish Knight'.", 'raw_generated_sql': '', 'refine_note': '', 'data_context_id': '30651'}
            database = row["database"]
            question = row["question"]
            sql = row["sql"].removesuffix(";")
            generated_sql = row["generated_sql"].removesuffix(";")
            difficulty = row["difficulty"]
            equal = row["equal"]
            db_path = "spider/bird/dev_databases/{}/{}.sqlite".format(database, database)
            bird_eval = bird_execute_sql(sql, generated_sql, db_path)

            writer.writerow([database, question, sql, generated_sql, difficulty, equal, row["note"], row["time_cost"], row["clarified_task"], row["raw_generated_sql"], row["refine_note"], row["data_context_id"], bird_eval])