import os
import re
import uuid
from time import gmtime, sleep, strftime

import openai
from rich import print

from fgn.utils.file_operations import open_file, save_to_project_folder

openai.api_key = os.environ["OPENAI_API_KEY"]


def save_completion(prompt):
    zulu = strftime("%Y-%m-%dT%H:%M:%SZ", gmtime())
    save_to_project_folder(
        f"data/completion/completion_{uuid.uuid4()}_{zulu}.txt", prompt
    )


def save_embedding(prompt):
    zulu = strftime("%Y-%m-%dT%H:%M:%SZ", gmtime())
    save_to_project_folder(
        f"data/embeddings/embedding_{uuid.uuid4()}_{zulu}.txt", prompt
    )


def gpt3_completion(
    prompt,
    engine="text-davinci-003",
    temp=1.0,
    top_p=1.0,
    tokens=400,
    freq_pen=0.0,
    pres_pen=0.0,
    stop=None,
):
    if stop is None:
        stop = ["<<STOP>>"]
    max_retry = 3
    retry = 0
    prompt = prompt.encode(encoding="ASCII", errors="ignore").decode()
    while True:
        try:
            response = openai.Completion.create(
                engine=engine,
                prompt=prompt,
                temperature=temp,
                max_tokens=tokens,
                top_p=top_p,
                frequency_penalty=freq_pen,
                presence_penalty=pres_pen,
                stop=stop,
            )
            text = response["choices"][0]["text"].strip()
            save_completion(text)
            return text
        except Exception as oops:
            retry += 1
            if retry >= max_retry:
                return "GPT3 error: %s" % oops
            print("Error communicating with OpenAI:", oops)
            sleep(1)


def gpt_embedding(text, model="text-embedding-ada-002"):
    text = text.replace("\n", " ")
    emb = openai.Embedding.create(input=[text], model=model)["data"][0]["embedding"]
    save_embedding(emb)
    return emb


def gpt_chat_completion(
    messages, model, max_retry=5, backoff_factor=2, initial_wait=0.1
):
    """
    Sends chat inputs to OpenAI API and receives the chatbot response.

    Args:
        messages (List[Dict]): A list of messages with role and content.
        model (str, optional): The GPT model to be used.
        max_retry (int, optional): The maximum number of retries.
        backoff_factor (int, optional): The factor to exponentially increase the waiting time.
        initial_wait (float, optional): The initial waiting time between retries in seconds.

    Returns:
        str: The content of the message generated by the chatbot.
    """

    # Initialize retry attempts
    retry = 0

    # Run the loop for retry attempts
    while retry <= max_retry:
        try:
            response = openai.ChatCompletion.create(model="gpt-4", messages=messages)
            text = response["choices"][0]["message"]["content"].strip()
            save_completion(text)
            return text
        except Exception as oops:
            # If the error is due to maximum context length, return the error message so that the user can handle it
            if "maximum context length" in str(oops):
                return f"GPT error: {oops}"

            retry += 1  # Increment the retry attempts

            # If reached the maximum retry attempts, return the error message
            if retry > max_retry:
                return f"GPT error: {oops}"

            # Calculate the waiting time for exponential backoff
            wait_time = initial_wait * (backoff_factor ** (retry - 1))

            # Print the error and wait before retrying
            print(
                f"Error communicating with OpenAI (attempt {retry}/{max_retry}): {oops}"
            )
            sleep(wait_time)


def gpt4_completion(prompt, model="gpt-4"):
    completion = gpt_chat_completion([{"role": "user", "content": prompt}], model=model)
    save_completion(completion)
    return completion


def generate_filename(
    prompt, prefix="", suffix="", extension="md", max_chars=60, time=False
):
    """
    Generates a filename based on the given prompt.
    :param prompt: the prompt to generate the filename from
    :param prefix: the prefix to add to the filename
    :param suffix: the suffix to add to the filename
    :param extension: the extension to add to the filename
    :param max_chars: the maximum number of characters in the filename (default: 60)
    :param time: whether to add the current time to the filename (default: False)
    :return: the generated filename
    """
    prompt = prompt[:300]

    file_name = gpt3_completion(
        f"Carefully read the following text and generate a highly relevant and unique filename for it. The filename "
        f"should be in all lowercase letters and consist of at most {max_chars} characters, separated by underscores. "
        f"Exclude any file extensions from the filename. Pay close attention to the content "
        f"and context of the text to avoid hallucinations or mistakes. The text is: '{prompt}'\n\nfile name:",
        temp=0,
        tokens=max_chars * 10,
    )

    # Remove all non-underscore and non-alphanumeric characters, and truncate to max_chars
    file_name = re.sub(r"[^a-zA-Z0-9_]", "", file_name)
    file_name = file_name[:max_chars]

    if prefix:
        file_name = f"{prefix}_{file_name}"

    if suffix:
        file_name = f"{file_name}_{suffix}"

    if time:
        file_name = f"{file_name}_{strftime('%Y-%m-%d_%H-%M-%S', gmtime())}"

    if extension:
        file_name = f"{file_name}.{extension}"

    # print('Generated filename:', file_name)

    return file_name


def generate_output_file(prompt, extension="md", max_chars=60, time=True):
    return generate_filename(
        prompt, extension=extension, max_chars=max_chars, time=time
    )
