import os
from typing import List, Dict
import gradio as gr
from openai import OpenAI
import json

# Initialize the OpenAI client
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Use GPT-4-Turbo model
MODEL_NAME = "gpt-4o-mini"

def read_prompt(file_path: str) -> str:
    with open(file_path, 'r') as file:
        return file.read().strip()

# Read the initial system prompt
SYSTEM_PROMPT = read_prompt('ToolSearchPrompt2.txt')

# Initialize message logs
message_logs: Dict[str, List[Dict[str, str]]] = {}

def save_message_logs():
    with open('message_logs.json', 'w') as f:
        json.dump(message_logs, f)

def load_message_logs():
    global message_logs
    try:
        with open('message_logs.json', 'r') as f:
            message_logs = json.load(f)
    except FileNotFoundError:
        message_logs = {}

load_message_logs()

def handle_chat(message: str, history: List[List[str]], user_id: str) -> List[List[str]]:
    if user_id not in message_logs:
        message_logs[user_id] = [{"role": "system", "content": SYSTEM_PROMPT}]
    
    message_logs[user_id].append({"role": "user", "content": message})
    
    response = client.chat.completions.create(
        model=MODEL_NAME,
        messages=message_logs[user_id],
        max_tokens=1600
    )
    
    assistant_message = response.choices[0].message.content
    message_logs[user_id].append({"role": "assistant", "content": assistant_message})
    
    history.append((message, assistant_message))
    save_message_logs()
    return history

def user_identity(name: str):
    return gr.update(value=name, interactive=False), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)

# Create the Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# AI Assistant")
    gr.Markdown("Enter your name to start chatting!")
    
    user_id = gr.Textbox(label="Your Name", placeholder="Enter your name here")
    submit_name = gr.Button("Submit Name")
    
    chatbot = gr.Chatbot(visible=False)
    msg = gr.Textbox(label="Message", placeholder="Type your message here...", visible=False)
    clear = gr.Button("Clear", visible=False)

    submit_name.click(
        user_identity,
        inputs=[user_id],
        outputs=[user_id, chatbot, msg, clear]
    )

    msg.submit(handle_chat, inputs=[msg, chatbot, user_id], outputs=[chatbot])
    msg.submit(lambda: "", inputs=None, outputs=[msg])
    clear.click(lambda: None, None, chatbot, queue=False)

# Launch the interface
if __name__ == "__main__":
    iface.launch()