import openai
import gradio as gr
from gradio.components import Audio, Textbox
import os
import re
import tiktoken
from transformers import GPT2Tokenizer
import whisper
import pandas as pd
from datetime import datetime, timezone, timedelta
import notion_df
import concurrent.futures
import nltk
from nltk.tokenize import sent_tokenize
nltk.download('punkt')


# # Define the tokenizer and model
# tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
# model = openai.api_key = os.environ["OPENAI_API_KEY"]

# # Define the initial message and messages list
# initmessage = 'You are a USMLE Tutor. Respond with ALWAYS layered "bullet points" (listing rather than sentences) to all input with a fun mneumonics to memorize that list. But you can answer up to 1200 words if the user requests longer response.'
# initial_message = {"role": "system", "content": 'You are a USMLE Tutor. Respond with ALWAYS layered "bullet points" (listing rather than sentences) to all input with a fun mneumonics to memorize that list. But you can answer up to 1200 words if the user requests longer response.'}
# messages = [initial_message]
# messages_rev = [initial_message]

# # Define the answer counter
# answer_count = 0

# # Define the Notion API key
# API_KEY = os.environ["API_KEY"]

import openai
import gradio as gr
from gradio.components import Audio, Textbox
import os
import re
import tiktoken
from transformers import GPT2Tokenizer
import whisper
import pandas as pd
from datetime import datetime, timezone, timedelta
import notion_df
import concurrent.futures
import nltk
from nltk.tokenize import sent_tokenize
nltk.download('punkt')
import spacy
from spacy import displacy
from gradio import Markdown
import threading


# Define the tokenizer and model

# openai.api_type = "azure"
# openai.api_base = "https://yena.openai.azure.com/"
# openai.api_version = "2022-12-01"


tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
model = openai.api_key = os.environ["OPENAI_API_KEY"]

# Define the initial message and messages list
initmessage = 'You are a MCAT Tutor. Respond with ALWAYS layered "bullet points" (listing rather than sentences) to all input with a fun mneumonics to memorize that list. But you can answer up to 1200 words if the user requests longer response.'
initial_message = {"role": "system", "content": 'You are a MCAT Tutor. Pay especially attention to "testable" or "exam," or any related terms in the input and highlight them as "EXAM TOPIC." Respond ALWAYS quiz me with high yield and relevant qustions on the input and the answers layed out with layered "bullet points" (listing rather than sentences) to all input with a fun mneumonics to memorize that list. Expand on each point with great detail lists not sentence.'}

messages = [initial_message]
messages_rev = [initial_message]

# Define the answer counter
answer_count = 0

# Define the Notion API key
API_KEY = os.environ["API_KEY"]

# Define the answer counter
answer_count = 0

nlp = spacy.load("en_core_web_sm")
def process_nlp(system_message):
    # Colorize the system message text
    colorized_text = colorize_text(system_message['content'])
    return colorized_text

def train(text):
    now_et = datetime.now(timezone(timedelta(hours=-4)))
    published_date = now_et.strftime('%m-%d-%y %H:%M')
    df = pd.DataFrame([text])
    notion_df.upload(df, 'https://www.notion.so/US-62e861a0b35f43da8ef9a7789512b8c2?pvs=4', title=str(published_date), api_key=API_KEY)

def colorize_text(text):
    colorized_text = ""
    lines = text.split("\n")

    for line in lines:
        doc = nlp(line)
        for token in doc:
            if token.ent_type_:
                colorized_text += f'**{token.text_with_ws}**'
            elif token.pos_ == 'NOUN':
                colorized_text += f'<span style="color: #FF3300; background-color: transparent;">{token.text_with_ws}</span>'
            elif token.pos_ == 'VERB':
                colorized_text += f'<span style="color: #FFFF00; background-color: transparent;">{token.text_with_ws}</span>'
            elif token.pos_ == 'ADJ':
                colorized_text += f'<span style="color: #00CC00; background-color: transparent;">{token.text_with_ws}</span>'
            elif token.pos_ == 'ADV':
                colorized_text += f'<span style="color: #FF6600; background-color: transparent;">{token.text_with_ws}</span>'
            elif token.is_digit:
                colorized_text += f'<span style="color: #9900CC; background-color: transparent;">{token.text_with_ws}</span>'
            elif token.is_punct:
                colorized_text += f'<span style="color: #8B4513; background-color: transparent;">{token.text_with_ws}</span>'
            elif token.is_quote:
                colorized_text += f'<span style="color: #008080; background-color: transparent;">{token.text_with_ws}</span>'
            else:
                colorized_text += token.text_with_ws
        colorized_text += "<br>"

    return colorized_text


def colorize_and_update(system_message, submit_update):
    colorized_system_message = colorize_text(system_message['content'])
    submit_update(None, colorized_system_message)  # Pass the colorized_system_message as the second output

def update_text_output(system_message, submit_update):
    submit_update(system_message['content'], None)


def transcribe(audio, text, submit_update=None):

    global messages
    global answer_count
    transcript = {'text': ''} 
    input_text = []

    # Check if the first word of the first line is "COLORIZE"
    if text and text.split("\n")[0].split(" ")[0].strip().upper() == "COLORIZE":
        train(text)
        colorized_input = colorize_text(text)
        return text, colorized_input
    
    # Transcribe the audio if provided
    if audio is not None:
        audio_file = open(audio, "rb")
        transcript = openai.Audio.transcribe("whisper-1", audio_file, language="en")

    # Tokenize the text input
    if text is not None:
        # Split the input text into sentences
        sentences = re.split("(?<=[.!?]) +", text)
    
        # Initialize a list to store the tokens
        input_tokens = []
    
        # Add each sentence to the input_tokens list
        for sentence in sentences:
            # Tokenize the sentence using the GPT-2 tokenizer
            sentence_tokens = tokenizer.encode(sentence)
            # Check if adding the sentence would exceed the token limit
            if len(input_tokens) + len(sentence_tokens) < 1440:
                # Add the sentence tokens to the input_tokens list
                input_tokens.extend(sentence_tokens)
            else:
                # If adding the sentence would exceed the token limit, truncate it
                sentence_tokens = sentence_tokens[:1440-len(input_tokens)]
                input_tokens.extend(sentence_tokens)
                break
        # Decode the input tokens into text
        input_text = tokenizer.decode(input_tokens)
    
        # Add the input text to the messages list
    messages.append({"role": "user", "content": transcript["text"]+input_text})


    # Check if the accumulated tokens have exceeded 2096
    num_tokens = sum(len(tokenizer.encode(message["content"])) for message in messages)
    if num_tokens > 2096:
        # Concatenate the chat history
        chat_transcript = "\n\n".join([f"[ANSWER {answer_count}]{message['role']}: {message['content']}" for message in messages if message['role'] != 'system'])

        # Append the number of tokens used to the end of the chat transcriptd
        chat_transcript += f"\n\nNumber of tokens used: {num_tokens}\n\n"

        # Get the current time in Eastern Time (ET)
        now_et = datetime.now(timezone(timedelta(hours=-4)))
        # Format the time as string (YY-MM-DD HH:MM)
        published_date = now_et.strftime('%m-%d-%y %H:%M')

        # Upload the chat transcript to Notion
        df = pd.DataFrame([chat_transcript])
        notion_df.upload(df, 'https://www.notion.so/YENA-be569d0a40c940e7b6e0679318215790?pvs=4', title=str(published_date+'back_up'), api_key=API_KEY)

        # Reset the messages list and answer counter
        messages = [initial_message]
        messages.append({"role": "user", "content": initmessage})
        answer_count = 0
        # Add the input text to the messages list
        messages.append({"role": "user", "content": input_text})
    else:
        # Increment the answer counter
        answer_count += 1

    # Generate the system message using the OpenAI API
    with concurrent.futures.ThreadPoolExecutor() as executor:
        prompt = [{"text": f"{message['role']}: {message['content']}\n\n"} for message in messages]
        system_message = openai.ChatCompletion.create(
            model="gpt-4",
            messages=messages,
            max_tokens=2000
        )["choices"][0]["message"]

    # Immediately update the text output
    if submit_update:  # Check if submit_update is not None
        update_text_output(system_message, submit_update)

    # Add the system message to the messages list
    messages.append(system_message)

    # Add the system message to the beginning of the messages list
    messages_rev.insert(0, system_message)
    # Add the input text to the messages list
    messages_rev.insert(0, {"role": "user", "content": input_text + transcript["text"]})

    # Start a separate thread to process the colorization and update the Gradio interface
    if submit_update:  # Check if submit_update is not None
        colorize_thread = threading.Thread(target=colorize_and_update, args=(system_message, submit_update))
        colorize_thread.start()

    # Return the system message immediately

    chat_transcript = system_message['content']

    # with open("./MSK_PS_conversation_history.txt", "a") as f:
    #     f.write(chat_transcript)
    
    # Get the current time in Eastern Time (ET)
    now_et = datetime.now(timezone(timedelta(hours=-4)))
    # Format the time as string (YY-MM-DD HH:MM)
    published_date = now_et.strftime('%m-%d-%y %H:%M')

    # Upload the chat transcript to Notion
    df = pd.DataFrame([chat_transcript])
    notion_df.upload(df, 'https://www.notion.so/YENA-be569d0a40c940e7b6e0679318215790?pvs=4', title=str(published_date+'back_up'), api_key=API_KEY)
    
    return system_message['content'], colorize_text(system_message['content'])



# Define the input and output components for Gradio
audio_input = Audio(source="microphone", type="filepath", label="Record your message")
text_input = Textbox(label="Type your message", max_length=4096)
# Define the input and output components for Gradio
output_text = Textbox(label="Text Output")
output_html = Markdown()

# Define the Gradio interface
iface = gr.Interface(
    fn=transcribe,
    inputs=[audio_input, text_input],
    outputs=[output_text, output_html],  # Add both output components
    title="Hold On, Pain Ends (HOPE)",
    description="Talk to Your Tutor MCAT HOPE. If you want to colorize your note, type COLORIZE in the first line of your input.",
    theme="compact",
    layout="vertical",
    allow_flagging=False
)


# Run the Gradio interface
iface.launch()