import logging
import sys
import csv
import glob

sys.path.insert(0, "../")

from svc.eda_task import EDATask
from svc.context import set_session
from models.base import with_session


logging.basicConfig(level=logging.INFO)


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


def read_all_case(filename):
    with open(filename) as f:
        for row in csv.DictReader(f):
            yield row


def rerun():
    with with_session() as s:
        set_session(s)
        eda = EDATask(151339)

        for filename in filter_all_cases_checked_result():
            logging.info("handling....%s", filename)

            with open(filename + ".llm.csv", "w") as f:
                d = {'database': 'entrepreneur', 'question': 'What are the weights of entrepreneurs in descending order of money requested?', 'sql': 'SELECT T2.Weight FROM entrepreneur AS T1 JOIN people AS T2 ON T1.People_ID  =  T2.People_ID ORDER BY T1.Money_Requested DESC;', 'generated_sql': 'SELECT `people`.`weight` FROM `people` INNER JOIN `entrepreneur` ON `people`.`people_id` = `entrepreneur`.`people_id` ORDER BY `entrepreneur`.`money_requested` DESC;', 'difficulty': 'medium', 'equal': 'True', 'note': '', 'time_cost': '85', 'clarified_task': 'The task requires the weights of entrepreneurs to be listed in descending order of the amount of money they have requested for funding.', 'data_context_id': '150439', 'optimized_sql': "", 'resp_str': ""}
                writer = csv.DictWriter(f, d.keys())
                writer.writeheader()

                for row in read_all_case(filename):
                    # {'database': 'entrepreneur', 'question': 'What are the weights of entrepreneurs in descending order of money requested?', 'sql': 'SELECT T2.Weight FROM entrepreneur AS T1 JOIN people AS T2 ON T1.People_ID  =  T2.People_ID ORDER BY T1.Money_Requested DESC;', 'generated_sql': 'SELECT `people`.`weight` FROM `people` INNER JOIN `entrepreneur` ON `people`.`people_id` = `entrepreneur`.`people_id` ORDER BY `entrepreneur`.`money_requested` DESC;', 'difficulty': 'medium', 'equal': 'True', 'note': '', 'time_cost': '85', 'clarified_task': 'The task requires the weights of entrepreneurs to be listed in descending order of the amount of money they have requested for funding.', 'data_context_id': '150439'}
                    question = row["question"]
                    generated_sql = row["generated_sql"]
                    if "job failed" in generated_sql:
                        optimized_sql, resp_str = "optimized_sql", "resp_str"
                    else:
                        task = {"task": question, "is_test": True}
                        optimized_sql, resp_str = eda.sql_revision_with_llm("query_id", task, "SQLite", generated_sql, use_llm=True)
                    row["optimized_sql"] = optimized_sql
                    row["resp_str"] = resp_str
                    writer.writerow(row)


if __name__ == "__main__":
    rerun()