import os
import yaml
import json
import threading
import time
from langchain.llms import OpenAI
from typing import Tuple, Any
from aievaluation.process_file import process_benchmark
from aievaluation.prompt import get_prompt_with_retry
from concurrent.futures import ThreadPoolExecutor

UNKNOWN_FLAG = "unknown."
GEN_PROMPT_FAIL_FLAG = "generate prompt fail."
GEN_SQL_RESP_INVALID_FLAG = "generate sql response invalid."
GEN_SQL_EMPTY_FLAG = "generated sql is empty."

def execute(need_generate_prompt: bool, model_path: str, template_path: str, benchmark):
    result_list = []
    with open(model_path, 'r') as f:
        model_meta = yaml.safe_load(f)

    with open(template_path, 'r') as f:
        template = f.read()

    llm = OpenAI(temperature=model_meta['config']['temperature'], max_tokens=model_meta['config']['max_tokens'],
                 openai_api_key=model_meta['config']['key'], model_name=model_meta['version'],max_retries=1,request_timeout=120)

    workers = 10
    if workers > len(benchmark["prompt_list"]):
        workers = len(benchmark["prompt_list"])

    pool = ThreadPoolExecutor(max_workers=workers)
    future_list = []

    for i in range(len(benchmark["prompt_list"])):
        result_list.append(UNKNOWN_FLAG)
        prompt = benchmark["prompt_list"][i]
        question = benchmark["question_list"][i]
        index = i
        wait(0.2)
        future = pool.submit(execute_prompt, index, prompt, question, template,need_generate_prompt, llm, result_list)
        future_list.append(future)

    print("wait for future_list")
    for future in future_list:
        try:
            future.result(300)
        except Exception as e:
            print('future.result exception, set result as {}, {}'.format(UNKNOWN_FLAG, e))

    print("pool.shutdown()")
    pool.shutdown()
    return result_list

def execute_prompt(index, prompt, question, template, need_generate_prompt, llm, result_list):
    if need_generate_prompt is True:
        prompt, is_ok = get_prompt_with_retry(question)
        if (is_ok is False) or (prompt == ''):
            print("ai execute get prompt fail, skip this question, index {}, question {}, prompt {}, is_ok {}".format(index, question, prompt, is_ok))
            result_list[index] = GEN_PROMPT_FAIL_FLAG
            return

    print('ai execute index {}, question {}'.format(index, question))

    flag, result = evaluate_on(
        template, llm, prompt, 0)

    if flag == "exception":
        print("index {}, question {}, some error found. sleep and retry.".format(index, question))
    else:
        if need_generate_prompt is True:
            try:
                result_json = json.loads(result)
                result = result_json["sql"]
            except ValueError as e:
                print("ai execute result not json, index {}, question {}, result {}, execption {}".format(index, question, result, e))
                result_list[index] = GEN_SQL_RESP_INVALID_FLAG
                return

        print('ai execute result parse success, index {}, question {}, result {}'.format(index, question, result))
        if len(result.strip()) == 0:
            print('ai execute result parse success, but sql is empty, index {}, question {}, result {}'.format(index, question, result))
            result_list[index] = GEN_SQL_EMPTY_FLAG
        else:
            result_list[index] = result

def evaluate_on(template: str, llm: OpenAI, prompt: str, retry: int) -> Tuple[str, Any]:
    try:
        return ("result", llm(template.format(prompt=prompt)))
    except Exception as e:
        print(e)
        retry += 1
        return ("exception", retry)

def wait(sec: float):
    time.sleep(sec)


if __name__ == "__main__":
    dirname, _ = os.path.split(__file__)
    execution_result = execute("./../../scripts/evaluation/text-davinci-003.yaml".format(dirname), "./../../scripts/evaluation/template.txt".format(dirname), "./../../scripts/evaluation/benchmark.csv".format(dirname))