import os
import threading
import time

from langchain.llms import OpenAI
from typing import Tuple, Any
from aievaluation.process_file import process_file
from aievaluation.template import prompt_template
from concurrent.futures import ThreadPoolExecutor
from kubernetes import client, config
import base64

UNKNOWN_FLAG = "unknown."

def evaluate(gold, pred, benchmark_set_path):
    result_list = []
    # init openai
    # get api key
    config.load_kube_config()
    v1 = client.CoreV1Api()
    secret = v1.read_namespaced_secret("openai-secret", "evaluation").data
    key = str(base64.b64decode(secret["openaikey"]), 'utf-8')

    llm = OpenAI(temperature=0, max_tokens=1024,
                 openai_api_key=key, model_name="gpt-3.5-turbo", max_retries=1, request_timeout=120)
    question_set = process_file(gold, pred, benchmark_set_path)

    gold_list, pred_list, question_list, schema_list = zip(*question_set)

    workers = 10
    if workers > len(gold_list):
        workers = len(gold_list)

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

    for i in range(len(gold_list)):
        result_list.append(UNKNOWN_FLAG)
        index = i
        future = pool.submit(evaluate_ai, index, question_list[index], schema_list[index], gold_list[index], pred_list[index], 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))
            continue

    print("pool.shutdown()")

    pool.shutdown()
    return {
        "result": result_list
    }

def evaluate_ai(index, question, schema, gold, pred, llm, result_list):
    flag, result = evaluate_on(
        llm, question, schema, gold, pred, 0)
    if flag == "exception":
        # if any error found. we should sleep 10 sec.
        print("ai evaluate some error found")
        return
    else:
        print('ai evaluate result {}'.format(result))
        result_list[index] = result
        return


def evaluate_on(llm: OpenAI, question: str, schema: str, gold: str, pred: str, retry: int) -> Tuple[str, Any]:
    try:
        return ("result", llm(prompt_template.format(question=question, schema=schema, gold_answer=gold, pred_answer=pred)))
    except Exception as e:
        print('ai evaluate exception {}'.format(e))
        retry += 1
        return ("exception", retry)


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


if __name__ == "__main__":
    import csv
    dirname, _ = os.path.split(__file__)
    result = evaluate("./../../scripts/evaluation/gold.txt".format(dirname), "./../../scripts/evaluation/pred.txt".format(
        dirname), "./../../scripts/evaluation/benchmark.csv".format(dirname))
    print(result)