from pathlib import Path
from time import perf_counter
from typing import Type
from datetime import datetime

from openai import OpenAI

from .models.trial_context import TrialContext
from .models.trial_loop import TrialLoop
from .utils import log


def run_evaluation(team_name: str, fn: Type[TrialLoop], num_trials=10,
                   characters: list[str] = None):
    """
    Run a trial for each character in the alphabet for a given team.
    :param team_name: team name
    :param fn: trial loop function
    :param num_trials: number of trials to run (default 10)
    :param characters: characters to run trials for (default all characters)
    :return: None
    """

    if characters is None:
        characters = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R",
                      "S", "T", "U", "V", "W", "X", "Y", "Z"]

    team_path = Path(team_name)
    logging_path = team_path / "logging"
    output_path = team_path / "raw"

    Path.mkdir(team_path, exist_ok=True)
    Path.mkdir(logging_path, exist_ok=True)
    Path.mkdir(output_path, exist_ok=True)

    log_file = logging_path / f"{team_name}_raw_{datetime.now().strftime('%Y_%m_%d_%H_%M_%S_%Z%z')}.log"
    log_file.touch(exist_ok=True)

    for character in characters:
        log(f"Running trials for character {character} for team {team_name}")

        character_path = output_path / character
        Path.mkdir(character_path, exist_ok=True)

        for trial_number in range(num_trials):
            log_file.write_text(
                f"Running trial {trial_number} for character {character} for team {team_name}")
            ctx = TrialContext(team_name, character, trial_number, log_file)
            if (character_path / f"{trial_number}.txt").exists():
                log_file.write_text(
                    f"Trial {trial_number} for character {character} for team {team_name} already exists. Skipping.")
                continue

            run_trial(ctx, fn)


def run_trial(ctx: TrialContext, fn: Type[TrialLoop]):
    """
    Run a single trial.
    :param ctx: context containing trial information
    :param fn: trial loop function
    :return: None
    """
    output_path = ctx.get_output_path()
    log_file_path = ctx.get_log_file_path()

    try:
        final_response = fn.run(ctx, ctx.get_character())
    except (TimeoutError, ValueError) as e:
        log_file_path.write_text(f"Trial {ctx.get_trial_number()} failed with error: {e}")
        return

    with open(output_path / f"{ctx.get_trial_number()}.txt", "w") as f:
        log_file_path.write_text(f"Trial {ctx.get_trial_number()} succeeded")
        f.write(final_response)


def chat_with_chatgpt(ctx: TrialContext, messages: []):
    """
    Chat with ChatGPT.
    :param ctx: context containing trial information
    :param messages: history of messages
    :return: response
    """
    model = "gpt-3.5-turbo-1106"
    temperature = 1
    seed = 42
    max_time = 60
    token_limit = 17000
    log_file_path = ctx.get_log_file_path()

    current_time = perf_counter()
    elapsed_time = current_time - ctx.get_start_time()
    log_file_path.write_text(f"Elapsed time: {elapsed_time}")

    if elapsed_time > max_time:
        log_file_path.write_text(f"Time limit exceeded. {elapsed_time} > {max_time}")
        raise TimeoutError(f"Time limit exceeded. {elapsed_time} > {max_time}")

    client = OpenAI()
    chat_completion = client.chat.completions.create(
        messages=messages,
        model=model,
        temperature=temperature,
        seed=seed
    )

    response = chat_completion.choices[0].message.content
    ctx.add_prompt_token_count(chat_completion.usage.prompt_tokens)
    ctx.add_output_token_count(chat_completion.usage.completion_tokens)
    log_file_path.write_text(f"Response: {response}")
    log_file_path.write_text(f"Prompt token count: {chat_completion.usage.prompt_tokens}")
    log_file_path.write_text(
        f"Output token count: {chat_completion.usage.completion_tokens}")
    log_file_path.write_text(f"Total token count: {ctx.get_total_token_count()}")

    if ctx.get_total_token_count() > token_limit:
        log_file_path.write_text(
            f"Prompt token limit exceeded. {ctx.get_total_token_count()} > {token_limit}")
        raise ValueError(f"Prompt token limit exceeded. {ctx.get_total_token_count()} > {token_limit}")

    return response
