import os
import gradio as gr
from openai import AzureOpenAI
from typing import List, Tuple, Generator
from dotenv import load_dotenv
import time
load_dotenv()

# Initialize Azure OpenAI client
client = AzureOpenAI(
    api_key=os.getenv("AZURE_OPENAI_KEY"),
    api_version="2024-02-01",
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT")
)

# Use the deployment name for your model
deployment_name = 'gpt-4'


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


# Read the prompt from file
PROMPT = read_prompt('ToolSearchPrompt2.txt')


def handle_chat(message: str, history: List[Tuple[str, str]]) -> Generator[str, None, None]:
    # Prepare the messages for the API call
    messages = [{"role": "system", "content": PROMPT}]
    for human, ai in history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": ai})
    messages.append({"role": "user", "content": message})

    # Call the Azure OpenAI API with streaming enabled
    stream = client.chat.completions.create(
        model=deployment_name,
        messages=messages,
        max_tokens=1600,
        stream=True
    )

    # Stream the response
    full_response = ""
    for chunk in stream:
        if chunk.choices and chunk.choices[0].delta.content:
            content = chunk.choices[0].delta.content
            full_response += content
            yield full_response
            time.sleep(0.1)


# Create the Gradio interface
iface = gr.ChatInterface(
    fn=handle_chat,
    title="Teacher's Assistant Chatbot",
    description="Ask me anything about teaching methods, lesson planning, or educational resources!",
    examples=[
        "How can I make my math lessons more engaging?",
        "What are some effective strategies for classroom management?",
        "Can you suggest some online resources for teaching science?",
    ],
    retry_btn="Regenerate",
    undo_btn="Undo",
    clear_btn="Clear",
)

# Launch the interface
iface.launch()