import gradio as gr
import json
import os
import google.generativeai as genai

# Function to get available JSON files in the working directory
def get_available_weeks():
    files = [f for f in os.listdir() if f.startswith('week-') and f.endswith('.json')]
    return files

# Function to load questions from the specified week file
def load_questions(week_file):
    try:
        with open(week_file, "r") as f:
            data = json.load(f)
        return data, None
    except FileNotFoundError:
        return None, f"File {week_file} not found."

# Flashcard UI function
def flashcard_ui(week_file, index):
    data, error = load_questions(week_file)
    if error:
        return f"Error: {error}", None
    question = data[index]["question"]
    total = len(data)
    return f"Question {index + 1}/{total}: {question}", ""

# Reveal answer function
def reveal_answer(week_file, index):
    data, error = load_questions(week_file)
    if error:
        return None
    answer = data[index]["answer"]
    return answer

# Function to handle navigation
def change_question(week_file, index, direction):
    data, _ = load_questions(week_file)
    total = len(data)
    index = (index + direction) % total
    return index, *flashcard_ui(week_file, index)

with open("aa2.txt", "r") as f:
  prompt = f.read()

genai.configure(api_key=os.environ["GEMINI_API_KEY"])
    
# Create the model configuration
generation_config = {
    "temperature": 0.5,
    "top_p": 0.98,
    "top_k": 64,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}

model = genai.GenerativeModel(
    model_name="gemini-1.5-flash-8b-exp-0827",
    generation_config=generation_config,
)

# AI Chatbot function
def chat_with_ai(user_input):
    chat_session = model.start_chat(
      history=[
        {
          "role": "user",
          "parts": [ prompt ],
        },
        {
          "role": "model",
          "parts": [
            "Sure, I can answer your question. \n",
          ],
        },
      ]
    )
    
    response = chat_session.send_message(user_input)
    return response.text

# Gradio interface
with gr.Blocks() as demo:
    available_weeks = get_available_weeks()
    
    # Splash page to select the week
    with gr.Row():
        week_file = gr.Dropdown(choices=available_weeks, label="Select CAM Week to Study")
        start_button = gr.Button("Start Studying")
    
    # Flashcard UI (hidden until week is selected)
    with gr.Column(visible=False) as flashcard_section:
        index = gr.State(0)
        question_output = gr.Textbox(label="Flashcard", interactive=False)
        answer_output = gr.Textbox(label="Answer", interactive=False)

        with gr.Row():
            prev_btn = gr.Button("Previous")
            reveal_btn = gr.Button("Reveal Answer")
            next_btn = gr.Button("Next")
    
    with gr.Column(visible=True):
        chat_input = gr.Textbox(label="Ask a question to the A.I.")
        with gr.Row():
            gr.Markdown(">Note: This A.I. has only been given the NBAA Management guide. Use the flashcards for specific week information.")
        chat_output = gr.Markdown(label="AI Response")
        chat_button = gr.Button("Ask AI")

    # Start button action to reveal the flashcard section
    start_button.click(lambda: gr.update(visible=True), inputs=[], outputs=flashcard_section)

    # Event handling for flashcards
    week_file.change(flashcard_ui, inputs=[week_file, index], outputs=[question_output, answer_output])
    reveal_btn.click(reveal_answer, inputs=[week_file, index], outputs=answer_output)
    prev_btn.click(change_question, inputs=[week_file, index, gr.Number(-1, visible=False)], outputs=[index, question_output, answer_output])
    next_btn.click(change_question, inputs=[week_file, index, gr.Number(1, visible=False)], outputs=[index, question_output, answer_output])

    # Chatbot button action
    chat_button.click(chat_with_ai, inputs=[chat_input], outputs=[chat_output])

# Launch the app
demo.launch()