import logging
import glob
import csv
from models.base import with_session

from svc.eda_task import EDATask
from utils.local import set_session


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def get_all_cases_file():
    for i in glob.glob('cases_new_*st_sqlite.csv'):
        yield i


def revise_csv_file(filename):
    new_filename = filename + ".revise.csv"

    with open(new_filename, "w") as f:
        writer = csv.writer(f)
        writer.writerow(["database", "question", "sql", "generated_sql", "difficulty", "optimized_sql", "spider_eval"])

        with open(filename, 'r') as f:
            reader = csv.DictReader(f)
            # {'database': 'inn_1', 'question': 'Find the name of rooms whose base price is between 120 and 150.', 'sql': 'SELECT roomname FROM rooms WHERE baseprice BETWEEN 120 AND 150;', 'generated_sql': 'SELECT `roomname` FROM `rooms` WHERE `baseprice` BETWEEN 120 AND 150;', 'difficulty': 'easy', 'equal': 'False', 'note': '', 'time_cost': 'N/A', 'clarified_task': '20', 'raw_generated_sql': 'Find the name of rooms whose base price is between 120 and 150.', 'refine_note': '', 'data_context_id': '44140'}
            for row in reader:
                question = row["question"]
                generated_sql = row["generated_sql"]
                database = row["database"]

                optimized_sql = revise_sql(database, question, generated_sql)
                writer.writerow([database, question, row["sql"], generated_sql, row["difficulty"], optimized_sql, ""])


def revise_sql(database, question, generated_sql):
    try:
        with with_session() as session:
            set_session(session)
            eda = EDATask(2, "")
            optimized_sql_1 = eda.revise_sql_for_spider_eval("", {"task": question}, "SQLite", generated_sql, ai_model="claude-3")
            optimized_sql_2 = eda.revise_sql_for_ignore_case("", {"task": question}, "SQLite", optimized_sql_1)
            optimized_sql_3 = eda.revise_sql_for_refine_SQL_condition("", {"task": question}, "SQLite", optimized_sql_2, ai_model="claude-3")

            return optimized_sql_3
    except Exception as e:
        logging.info("error in revise_sql, database, question, generated_sql, error: %s, %s, %s, %s", database, question, generated_sql, e)
        return generated_sql


if __name__ == "__main__":
    for i in get_all_cases_file():
        revise_csv_file(i)
        logging.info("revise %s done" % i)