"""Evaluation utilities for the Arklex framework.

This module provides functionality for evaluating the performance of the Arklex
framework through conversation simulation and metrics extraction. It includes
utilities for simulating first-pass and second-pass conversations, extracting
task completion metrics, and generating labeled conversation data for analysis.
"""

import argparse
import json
import os
from typing import Any

from arklex.evaluation.chatgpt_utils import create_client
from arklex.evaluation.extract_conversation_info import extract_task_completion_metrics
from arklex.evaluation.simulate_first_pass_convos import simulate_conversations
from arklex.evaluation.simulate_second_pass_convos import get_labeled_convos
from arklex.utils.model_config import MODEL
from arklex.utils.model_provider_config import LLM_PROVIDERS


def evaluate(
    config: dict[str, Any],
) -> tuple[
    list[dict[str, Any]], list[dict[str, Any]], dict[str, Any], list[dict[str, Any]]
]:
    """Evaluate the performance of the Arklex framework based on the provided configuration.

    This function simulates conversations and extracts metrics based on the task specified in the config.
    It supports first-pass and second-pass evaluations, and returns the results of the evaluation.

    Args:
        config (Dict[str, Any]): Configuration dictionary containing task, model API, model parameters,
                                 synthetic data parameters, and other settings.

    Returns:
        Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]:
            A tuple containing the first-pass data, labeled conversations, goal metrics, and goals.
    """
    task: str = config["task"]
    model_api: str = config["model_api"]
    model_params: dict[str, Any] = config["model_params"]
    synthetic_data_params: dict[str, Any] = config["synthetic_data_params"]
    bot_goal: str | None = config.get("builder_objective")
    bot_goal = None if bot_goal == "" else bot_goal

    first_pass_data: list[dict[str, Any]]
    goals: list[dict[str, Any]]
    goal_metrics: dict[str, Any]
    labeled_convos: list[dict[str, Any]]

    if task == "first_pass" or task == "simulate_conv_only":
        # first pass
        first_pass_data, goals = simulate_conversations(
            model_api, model_params, synthetic_data_params, config
        )
        goal_metrics = extract_task_completion_metrics(
            first_pass_data, config["client"], bot_goal
        )
        data = first_pass_data

    # second pass
    if task == "all":
        labeled_convos = get_labeled_convos(
            first_pass_data, model_api, synthetic_data_params, model_params, config
        )
    else:
        labeled_convos = []
    return data, labeled_convos, goal_metrics, goals


if __name__ == "__main__":
    parser: argparse.ArgumentParser = argparse.ArgumentParser()
    parser.add_argument("--model_api", type=str)
    parser.add_argument("--model_params", type=dict, default={})
    parser.add_argument("--num_convos", type=int, default=5)
    parser.add_argument("--num_goals", type=int, default=5)
    parser.add_argument("--max_turns", type=int, default=5)
    parser.add_argument("--documents_dir", type=str)
    parser.add_argument("--config", type=str)
    parser.add_argument("--output_dir", type=str)
    parser.add_argument("--model", type=str, default=MODEL["model_type_or_path"])
    parser.add_argument(
        "--llm-provider", type=str, default=MODEL["llm_provider"], choices=LLM_PROVIDERS
    )
    parser.add_argument(
        "--customer_type", type=str, default=None, choices=["b2b", "b2c"]
    )
    parser.add_argument(
        "--task",
        type=str,
        default="first_pass",
        choices=["first_pass", "simulate_conv_only", "all"],
    )
    parser.add_argument(
        "--user_attributes", type=str, default="arklex/evaluation/user_attributes.json"
    )
    parser.add_argument("--custom_profile", action="store_true")
    parser.add_argument("--system_inputs", action="store_true")
    parser.add_argument("--data_file", type=str, default=None)
    args: argparse.Namespace = parser.parse_args()

    MODEL["model_type_or_path"] = args.model
    MODEL["llm_provider"] = args.llm_provider
    client: Any = create_client()

    assert args.model_api is not None, "Model api must be provided"
    assert args.config is not None, "Config file must be provided"
    assert args.user_attributes is not None, "User attribute file must be provided"
    if not args.output_dir:
        args.output_dir = os.path.join(args.documents_dir, "eval")

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir, exist_ok=True)

    with open(args.config) as f:
        config: dict[str, Any] = json.load(f)
    with open(args.user_attributes) as f:
        user_attributes: dict[str, Any] = json.load(f)
    # if args.testset:
    #     testset = json.load(open(args.testset))
    # else:
    #     testset = {}
    config["model_api"] = args.model_api
    config["documents_dir"] = args.documents_dir
    config["output_dir"] = args.output_dir
    config["model_params"] = args.model_params
    config["synthetic_data_params"] = {
        "num_convos": args.num_convos,
        "num_goals": args.num_goals,
        "max_turns": args.max_turns,
        "customer_type": args.customer_type,
        "data_file": args.data_file,
    }
    config["task"] = args.task
    config["user_attributes"] = user_attributes
    config["custom_profile"] = args.custom_profile
    config["system_inputs"] = args.system_inputs
    config["client"] = client

    first_pass_data: list[dict[str, Any]]
    final_convos: list[dict[str, Any]]
    goal_metrics: dict[str, Any]
    goals: list[dict[str, Any]]
    first_pass_data, final_convos, goal_metrics, goals = evaluate(config)

    with open(os.path.join(args.output_dir, "goals.json"), "w") as f:
        json.dump(goals, f, indent=4)

    with open(os.path.join(args.output_dir, "simulate_data.json"), "w") as f:
        json.dump(first_pass_data, f, indent=4)

    with open(os.path.join(args.output_dir, "labeled_data.json"), "w") as f:
        json.dump(final_convos, f, indent=4)

    with open(os.path.join(args.output_dir, "goal_completion.json"), "w") as f:
        json.dump(goal_metrics, f, indent=4)
