from datasets import load_from_disk, Dataset, DatasetDict
from tqdm import tqdm
from eval_prompt import EVALUATION_PROMPT
from openai import OpenAI, APIError
import json
import os
from functools import partial
import multiprocessing
import asyncio
from single_agent_with_rag import _single_agent_answer_with_rag, _single_agent_answer_with_rag_cot
from multiagent import _multiagent_with_rag_cot

OPENAI_API_KEY = "sk-urFGAQRThR6pysea0aC93bD27fA34bA69811A9254aAaD8B2"
OPENAI_BASE_URL = "http://154.44.26.195:17935/v1"
MODEL_NAME = "chatgpt-4o-latest"
DATASET_PATH = "_backend/evaluate/eval_rag_dataset"
EVAL_RESULT_PATH = "_backend/evaluate/eval_rag_result"


def load_eval_rag_dataset(dataset_path: str) -> DatasetDict:
    """Loads the eval_rag_dataset from disk.

    Args:
        dataset_path (str): The path to the dataset.

    Returns:
        DatasetDict: The loaded dataset.
    """
    return load_from_disk(dataset_path)


def get_response_from_llm(messages: list[dict], tools: list = None, model: str = MODEL_NAME):
    client = OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL)
    try:
        if tools is None:
            response = client.chat.completions.create(
                model=model,
                messages=messages,
            )
        else:
            response = client.chat.completions.create(
                model=model,
                messages=messages,
                tools=tools
            )
        content = response.choices[0].message.content
        return content
    
    except APIError as e:
        print(e)
        return "apierror"

    except Exception as e:
        print(e)
        return "error"


def _single_model_answer(question: str, model: str):
    """Answers a question with a single model.

    Args:
        question (str): The question to answer.
        context (str): The context to answer the question in.
        model (str): The model to use.

    Returns:
        str: The answer.
    """
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": question},
    ]

    if model == "o1-2024-12-17" or model == "o3-mini":
        messages = [{"role": "user", "content": question}]

    return get_response_from_llm(messages=messages, model=model)


def single_model_answer(model: str):
    eval_dataset = load_eval_rag_dataset(DATASET_PATH)
    num_threads = multiprocessing.cpu_count()
    with multiprocessing.Pool(processes=num_threads) as pool:
        results = list(
            tqdm(
                pool.imap(
                    partial(_single_model_answer, model=model),
                    eval_dataset['question'],
                ),
                total=len(eval_dataset),
                desc=f"{model} Answering:",
            )
        )
    final_result = []
    for i, idx in enumerate(eval_dataset):
        final_result.append({"question": idx['question'], "answer": results[i], "source_doc": idx['source_doc']})
    
    os.makedirs(os.path.join(EVAL_RESULT_PATH, model), exist_ok=True)
    with open(f"{EVAL_RESULT_PATH}/{model}/single_model_answer.json", "w") as f:
        json.dump(final_result, f, indent=2)


def run_async_in_process(func, *args, **kwargs):
    return asyncio.run(func(*args, **kwargs))


def single_model_answer_with_rag(model: str):
    eval_dataset = load_eval_rag_dataset(DATASET_PATH)
    num_threads = multiprocessing.cpu_count()
    with multiprocessing.Pool(processes=num_threads) as pool:
        results = list(
            tqdm(
                pool.imap(
                    partial(run_async_in_process, _single_agent_answer_with_rag, model=model),
                    eval_dataset['question'],
                ),
                total=len(eval_dataset),
                desc=f"{model} Answering:",
            )
        )
    final_result = []
    for i, idx in enumerate(eval_dataset):
        final_result.append({"question": idx['question'], "answer": results[i], "source_doc": idx['source_doc']})
    
    os.makedirs(os.path.join(EVAL_RESULT_PATH, model), exist_ok=True)
    with open(f"{EVAL_RESULT_PATH}/{model}/single_model_answer_with_rag.json", "w") as f:
        json.dump(final_result, f, indent=2)


def single_model_answer_with_rag_cot(model: str):
    eval_dataset = load_eval_rag_dataset(DATASET_PATH)
    num_threads = multiprocessing.cpu_count()
    with multiprocessing.Pool(processes=num_threads) as pool:
        results = list(
            tqdm(
                pool.imap(
                    partial(run_async_in_process, _single_agent_answer_with_rag_cot, model=model),
                    eval_dataset['question'],
                ),
                total=len(eval_dataset),
                desc=f"{model} Answering:",
            )
        )
    final_result = []
    for i, idx in enumerate(eval_dataset):
        final_result.append({"question": idx['question'], "answer": results[i], "source_doc": idx['source_doc']})
    
    os.makedirs(os.path.join(EVAL_RESULT_PATH, model), exist_ok=True)
    with open(f"{EVAL_RESULT_PATH}/{model}/single_model_answer_with_rag_cot.json", "w") as f:
        json.dump(final_result, f, indent=2)


def multiagent_with_rag_cot(model: str):
    eval_dataset = load_eval_rag_dataset(DATASET_PATH)
    num_threads = 16 #multiprocessing.cpu_count()
    with multiprocessing.Pool(processes=num_threads) as pool:
        results = list(
            tqdm(
                pool.imap(
                    partial(run_async_in_process, _multiagent_with_rag_cot),
                    eval_dataset['question'],
                ),
                total=len(eval_dataset),
                desc=f"{model} Answering:",
            )
        )
    final_result = []
    for i, idx in enumerate(eval_dataset):
        final_result.append({"question": idx['question'], "answer": results[i], "source_doc": idx['source_doc']})
    
    # final_result = []
    # for idx in tqdm(eval_dataset):
    #     answer = asyncio.run(_multiagent_with_rag_cot(idx['question']))
    #     # answer = await _multiagent_with_rag_cot(idx['question'])
    #     final_result.append({"question": idx['question'], "answer": answer, "source_doc": idx['source_doc']})

    os.makedirs(os.path.join(EVAL_RESULT_PATH, model), exist_ok=True)
    with open(f"{EVAL_RESULT_PATH}/{model}/multiagent_with_rag_cot.json", "w") as f:
        json.dump(final_result, f, indent=2)


def _eval_rag_dataset(instruction: str, response: str, context: str, model: str):
    """Evaluates a response with a single model.

    Args:
        instruction (str): The instruction to evaluate the response with.
        response (str): The response to evaluate.
        context (str): The context to evaluate the response in.
        model (str): The model to use.

    Returns:
        str: The evaluation.
    """
    messages = [
        {"role": "system", "content": "You are a fair evaluator language model."},
        {"role": "user", "content": EVALUATION_PROMPT.format(instruction=instruction, response=response, context=context)},
    ]
    eval_result = get_response_from_llm(messages)
    feedback, score = [item.strip() for item in eval_result.split("[RESULT]")]


def eval_rag_dataset():
    eval_dataset = load_eval_rag_dataset(DATASET_PATH)
    for i in eval_dataset:
        print()


if __name__ == "__main__":
    # single_model_answer(model="chatgpt-4o-latest")
    # single_model_answer(model="o1-2024-12-17")
    single_model_answer(model="o3-mini")
    # single_model_answer_with_rag(model="gpt-4o-2024-08-06")
    # single_model_answer_with_rag_cot(model="gpt-4o-2024-08-06")
    # multiagent_with_rag_cot(model="gpt-4o-2024-08-06")
    