import gradio as gr
import openai
import os
from typing import Generator

openai.api_key = os.getenv("OPENAI_API_KEY")

BASE_SYSTEM_MESSAGE = "You are a pilot assistant. You should provide concise information about the questions asked."
MAX_TOKENS = 1024

def get_system_message(aircraft: str) -> str:
    return f"{BASE_SYSTEM_MESSAGE} You are specifically knowledgeable about the {aircraft} aircraft."

def generate_response(message, history, aircraft):
    messages = [{"role": "system", "content": get_system_message(aircraft)}]
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

    response = openai.chat.completions.create(
        model="gpt-4o",
        messages=messages,
        max_tokens=MAX_TOKENS,
        stream=False,
    )
    return response.choices[0].message.content

def update_chat_history(history):
    return "\n".join(
        f"User: {msg[0]}\nAssistant: {msg[1]}" for msg in history
    )

def handle_submit(user_input, aircraft, history):
    if not user_input.strip():
        return "", history, "Please enter a question."
    reply = generate_response(user_input, history, aircraft)
    history.append((user_input, reply))
    updated_history = update_chat_history(history)
    return "", history, updated_history

with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("# Avia Agent", elem_id="title")
        restart_btn = gr.Button("Restart", scale=0.2, elem_classes=["small-button"])


    with gr.Row():
        aircraft = gr.Dropdown(
            choices=["C172", "Cirrus SR20/SR22"],
            value="C172",
            label="Select Aircraft"
        )

    with gr.Row():
        chat_history_display = gr.Textbox(
            label="Conversation History",
            lines=3,
            interactive=False
        )

    with gr.Row():
        user_input = gr.Textbox(
            label="Ask a question",
            placeholder="How to do steep turns?",
            lines=2,
        )
        submit_btn = gr.Button("Send", scale=0.1, elem_classes=["small-button"])


    conversation_state = gr.State([])

    submit_btn.click(
        handle_submit,
        inputs=[user_input, aircraft, conversation_state],
        outputs=[user_input, conversation_state, chat_history_display]
    )

    restart_btn.click(None, None, None, js="() => { window.location.reload() }")

if __name__ == "__main__":
    demo.launch()