from typing import List, Tuple, Dict, Generator
from langchain.chat_models import ChatOpenAI
import gradio as gr
import json
import os
import openai

parent = os.path.dirname(__file__)
data_path = os.path.join(parent, 'data.json')

def create_formatted_messages( messages: List[dict]) -> List[Tuple[str, str]]:
    '''
    create formatted history
    =
    Create formatted history for gradio chatbot
    '''

    user_messages = []
    formatted_messages = []
    assistant_messages = []

    for message in messages:
        if message['role'] == 'user':
            user_messages.append(message['content'])
        elif message['role'] == 'assistant':
            assistant_messages.append(message['content'])

        if user_messages and assistant_messages:
            formatted_messages.append(
                (''.join(user_messages), ''.join(assistant_messages))
            )
            user_messages = []
            assistant_messages = []
   
    # append any remaining messages
    if user_messages:
        formatted_messages.append((''.join(user_messages), None))
    elif assistant_messages:
        formatted_messages.append((None, ''.join(assistant_messages)))

    return formatted_messages

class ChatGPT:

    def chat(self, message: str, messages: List[Dict[str, str]], message_length:int) \
        -> Generator[Tuple[List[Tuple[str, str]], List[Dict[str, str]]], None, None]:
        if messages == None:
            messages = []
            messages.append({'role': 'system', 'content': 'ChatDefense is available to assist you with your legal questions.'})

        messages.append({'role': 'user', 'content': message})
        # We have no content for the assistant's response yet but we will update this:
        messages.append({'role': 'assistant', 'content': ''})

        response_message = ''
        try:
            chat_generator = self.openai.client.create(
                messages=messages[-2*int(message_length):], 
                model=self.openai.model_name,
                stream=True, 
            )

            for chunk in chat_generator:
                if 'choices' in chunk:
                    for choice in chunk['choices']:
                        if 'delta' in choice and 'content' in choice['delta']:
                            new_token = choice['delta']['content']
                            # Add the latest token:
                            response_message += new_token
                            # Update the assistant's response in our model:
                            messages[-1]['content'] = response_message

                        if 'finish_reason' in choice and choice['finish_reason'] == 'stop':
                            break
                formatted_messages = create_formatted_messages(messages)
                # TODO database
                json.dump(messages, open(data_path, 'w'))
                yield '', messages, formatted_messages
    
        except Exception as e:
            print(str(e))


    def clear(self):
        json.dump([], open(data_path, 'w'))
        return [], []

    def init_agent(self, openai_api_key):
        self.openai = ChatOpenAI(openai_api_key=openai_api_key)
        return gr.update(visible = True)

chatgpt = ChatGPT()

with gr.Blocks() as iface:
    
    with gr.Row():
        openai_api_key_textbox = gr.Textbox(
            placeholder='Paste your OpenAI API key here to start ChatGPT(sk-...) and press Enter',
            label='OpenAI API key',
            type='password',
        )
        # larger lenght will generate a more accuarte response but exponetially larger token size
        message_length = gr.Number(value=5, label='History message length')

    data = json.load(open(data_path)) if os.path.exists(data_path) else {}
    chatbot = gr.Chatbot(value=create_formatted_messages(data)).style()
    
    with gr.Row(visible=False) as input_raws:
        with gr.Column(scale=.9):
            txt = gr.Textbox(placeholder='Enter message here', show_label=False).style(container=False)
        with gr.Column(scale=.1, min_width=0):
            clear = gr.Button(value='Clear')
        
        messages = gr.State(data or [])
        txt.submit(chatgpt.chat, [txt, messages, message_length], [txt, messages, chatbot])
        clear.click(chatgpt.clear, [], [chatbot, messages])

    openai_api_key_textbox.submit(chatgpt.init_agent, [openai_api_key_textbox], [input_raws])

iface.queue().launch()