import asyncio
import sys
from aievaluation.execution import UNKNOWN_FLAG, GEN_PROMPT_FAIL_FLAG, GEN_SQL_RESP_INVALID_FLAG, GEN_SQL_EMPTY_FLAG

GEN_MODEL_PROMPT_FAIL = "GEN_MODEL_PROMPT_FAIL"
GET_MODEL_RESULT_FAIL = "GET_MODEL_RESULT_FAIL"
MODEL_RESULT_INVALID = "MODEL_RESULT_INVALID"
MODEL_RESULT_HAS_NO_SQL = "MODEL_RESULT_HAS_NO_SQL"
MODEL_RESULT_SQL_EXEC_SUCCESS = "MODEL_RESULT_SQL_EXEC_SUCCESS"
MODEL_RESULT_SQL_EXEC_FAIL = "MODEL_RESULT_SQL_EXEC_FAIL"

total_results = [GEN_MODEL_PROMPT_FAIL, GET_MODEL_RESULT_FAIL, MODEL_RESULT_INVALID, MODEL_RESULT_HAS_NO_SQL, MODEL_RESULT_SQL_EXEC_SUCCESS, MODEL_RESULT_SQL_EXEC_FAIL]

sql_pred_gen_err_to_evaluation_err = {
    GEN_PROMPT_FAIL_FLAG: GEN_MODEL_PROMPT_FAIL,
    UNKNOWN_FLAG: GET_MODEL_RESULT_FAIL,
    GEN_SQL_RESP_INVALID_FLAG: MODEL_RESULT_INVALID,
    GEN_SQL_EMPTY_FLAG: MODEL_RESULT_HAS_NO_SQL
}

def evaluate(db_accessor, db_id_list, pred_list):
    result_statistic = {}
    for result in total_results:
        result_statistic[result] = 0

    result_list = []
    for index in range(len(pred_list)):
        print('executable evaluate index {}'.format(index))
        pred = pred_list[index]
        if pred in sql_pred_gen_err_to_evaluation_err:
            result = sql_pred_gen_err_to_evaluation_err[pred]
            result_list.append((result, pred))
            result_statistic[result] += 1
            print('executable evaluate exec pred on db skipped, db {} pred {}'.format(db, pred))
            continue

        db = db_accessor.get_default_db(db_id_list[index])
        flag, denotation = db_accessor.exec_on_db(db, pred)
        if flag == 'exception':
            print('executable evaluate exec pred on db exception, db {} pred {} exception {}'.format(db, pred, denotation))
            result = MODEL_RESULT_SQL_EXEC_FAIL
            result_list.append((result, denotation))
            result_statistic[result] += 1
        else:
            print('executable evaluate exec pred on db succeed, db {} pred {}'.format(db, pred))
            result = MODEL_RESULT_SQL_EXEC_SUCCESS
            result_list.append((result, denotation))
            result_statistic[result] += 1

    result_statistic['result'] = result_list

    calc_executable_score(result_statistic)

    return result_statistic

def get_result_score_name(result):
    return '{}_RATE'.format(result)

def calc_executable_score(result_statistic):
    if (result_statistic is None) or (('result' in result_statistic) is False):
        return
    
    total_result_cnt = len(result_statistic['result'])

    for result in total_results:
        result_cnt = 0
        if result in result_statistic:
            result_cnt = result_statistic[result]
        result_statistic[get_result_score_name(result)] = result_cnt/total_result_cnt