import os
import time
import json
import openai
import gradio as gr
from datetime import datetime
from openai.error import RateLimitError, APIConnectionError, Timeout, APIError, \
    ServiceUnavailableError
from huggingface_hub import hf_hub_download, HfApi


def get_main_data():
    """
    Initializes the key for the api and returns the parameters for the scores, name of the possible authors
    and prompts (the one for the conversation and another for the summary)
    """
    openai.api_key = os.environ.get('API_KEY')

    scores_parameters = [
        'Personalidad', 'Intereses', 'Lenguaje/Estilo', 'Autenticidad', 'Habilidad de conversación',
        'Marca/Producto', 'Identificación', 'Experiencia de uso', 'Recomendacion', 'Conversación organica'
    ]

    authors = ['Sofia', 'Eliza', 'Sindy', 'Carlos', 'Andres', 'Adriana', 'Carolina', 'Valeria']

    with open('prompt_conversation.txt', encoding='utf-8') as file:
        prompt_conversation = file.read()

    return scores_parameters, authors, prompt_conversation


def innit_bot(prompt: str):
    """
    Initialize the bot by adding the prompt from the txt file to the messages history
    """
    prompt.replace('HISTORY', '')
    message_history = [{"role": "system", "content": prompt}]

    return message_history


def make_visible():
    """
    Makes visible the returned elements
    """
    return (
        gr.Chatbot.update(visible=True),
        gr.Textbox.update(visible=True),
        gr.Row.update(visible=True))


def make_noninteractive():
    """
    Makes no interactive the returned elements
    """
    return gr.Dropdown.update(interactive=False)


def call_api(msg_history: gr.State, cost: gr.State):
    """
    Returns the API's response
    """
    response = openai.ChatCompletion.create(
        model="gpt-4",
        messages=msg_history,
        temperature=0.8
    )

    print("*" * 20)
    print(msg_history)
    print("*" * 20)

    tokens_input = response['usage']['prompt_tokens']
    tokens_output = response['usage']['completion_tokens']

    cost.append({'Model': 'gpt-4', 'Input': tokens_input, 'Output': tokens_output})

    return response


def handle_call(msg_history: gr.State, cost: gr.State):
    """
    Returns the response and waiting time of the AI. It also handles the possible errors
    """
    tries = 0
    max_tries = 3
    while True:
        try:
            start_time = time.time()
            response = call_api(msg_history, cost)
            end_time = time.time()
            break

        except (RateLimitError, APIError, Timeout, APIConnectionError, ServiceUnavailableError) as e:
            print(e)

            if tries == max_tries:
                response = "Despues de muchos intentos, no se pudo completar la comunicacion con OpenAI. " \
                           "Envia lo que tengas hasta el momento e inicia un chat nuevo dentro de unos minutos."
                raise gr.Error(response)

            tries += 1
            time.sleep(60)

    needed_time = end_time - start_time
    return response, needed_time


def get_template(chatbot_history: gr.Chatbot, previous_summary: gr.State):
    with open('prompt_summary.txt', encoding='utf-8') as file:
        template_summary = file.read()

    conversation = ''

    for i, msg in enumerate(chatbot_history):
        conversation += f'Usuario: {msg[0]} \n'
        conversation += f'Roomie: {msg[1]} \n'

    template_summary = template_summary.replace('CONVERSATION', conversation)

    return template_summary


def get_summary(chatbot_history: gr.Chatbot, previous_summary: gr.State, cost: gr.State):

    msg = get_template(chatbot_history, previous_summary)

    print(msg, end='\n\n')

    with open('prompt_summary_system.txt', encoding='utf-8') as file:
        system_prompt = file.read()

    calling = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": msg}
    ]
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=calling,
        temperature=0
    )

    tokens_input = response['usage']['prompt_tokens']
    tokens_output = response['usage']['completion_tokens']

    cost.append({'Model': 'gpt-3.5-turbo', 'Input': tokens_input, 'Output': tokens_output})

    return response["choices"][0]["message"]["content"]


def get_ai_answer(
        msg: str, msg_history: gr.State, num_interactions: gr.State, previous_summary: gr.State,
        cost: gr.State, chatbot_history: gr.Chatbot):
    """
    Returns the response given by the model, all the message history so far and the seconds
    the api took to retrieve such response. It also removes some messages in the message history
    so only the last n (keep) are used (costs are cheaper)
    """
    # Call GPT 3.5
    if num_interactions >= 2:
        previous_output = msg_history.pop()
        summary = get_summary(chatbot_history, previous_summary, cost)
        with open('prompt_conversation.txt', encoding='utf-8') as file:
            prompt_template = file.read()
        prompt_template = prompt_template.replace('HISTORY', summary)
        msg_history = [{"role": "system", "content": prompt_template}]
        msg_history.append(previous_output)
        print('RESUMEN DE GPT 3.5', summary, end='\n----------------------------------------------------------------\n')
    else:
        summary = ''

    # Call GPT 4
    msg_history.append({"role": "user", "content": msg})
    response, needed_time = handle_call(msg_history, cost)
    AI_response = response["choices"][0]["message"]["content"]
    msg_history.append({'role': 'assistant', 'content': AI_response})

    return AI_response, msg_history, needed_time, summary


def get_answer(
        msg: str, msg_history: gr.State, chatbot_history: gr.Chatbot,
        waiting_time: gr.State, num_interactions: gr.State, previous_summary: gr.State,
        cost: gr.State):
    """
    Cleans msg box, adds the new message to the message history,
    gets the answer from the bot and adds it to the chatbot history
    and gets the time needed to get such answer and saves it
    """
    # Get bot answer (output), messages history and waiting time
    AI_response, msg_history, needed_time, summary = get_ai_answer(
        msg, msg_history, num_interactions, previous_summary, cost, chatbot_history
    )

    # Save waiting time
    waiting_time.append(needed_time)

    # Save output in the chat
    chatbot_history.append((msg, AI_response))

    num_interactions += 1

    return "", msg_history, chatbot_history, waiting_time, num_interactions, summary, cost


def save_scores(
        author: gr.Dropdown, history: gr.Chatbot, waiting_time: gr.State, opinion: gr.Textbox,
        cost: gr.State, *score_values):
    """
    Saves the scores and chat's info into the json file
    """
    # Get the parameters for each score
    score_parameters, _, _ = get_main_data()

    # Get the score of each parameter
    scores = dict()
    for parameter, score in zip(score_parameters, score_values):

        # Check the score is a valid value if not, raise Error
        if score is None:
            raise gr.Error('Asegurese de haber seleccionado al menos 1 opcion en cada categoria')

        scores[parameter] = score

    # Get all the messages including their reaction
    chat = []
    for conversation in history:
        info = {
            'message': conversation[0],
            'answer': conversation[1],
            'waiting': waiting_time.pop(0)
        }
        chat.append(info)

    date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    with open('prompt_conversation.txt', encoding='utf-8') as file:
        prompt = file.read()

    # Save the info
    session = dict(
        prompt=prompt,
        temperature=0.8,
        scores=scores,
        opinion=opinion,
        chat=chat,
        cost=cost,
        author=author,
        model='gpt-4',
        date=date
    )

    # Open the file, add the new info and save it
    hf_hub_download(
        repo_id=os.environ.get('DATA'), repo_type='dataset', filename="data.json", token=os.environ.get('HUB_TOKEN'),
        local_dir="./"
    )

    with open('data.json', 'r') as infile:
        past_sessions = json.load(infile)

    # Add the new info
    past_sessions['sessions'].append(session)
    with open('data.json', 'w', encoding='utf-8') as outfile:
        json.dump(past_sessions, outfile, indent=4, ensure_ascii=False)

    # Save the updated file
    api = HfApi(token=os.environ.get('HUB_TOKEN'))
    api.upload_file(
        path_or_fileobj="data.json",
        path_in_repo="data.json",
        repo_id=os.environ.get('DATA'),
        repo_type='dataset'
    )

    # Return a confirmation message
    return 'Done'