with gr.Blocks(css="footer{display:none !important}") as demo:
    gr.Markdown(
        """
        # 👩‍💻 AP Computer Science Learning Assistant  
        Welcome! This is your friendly AI tutor for **AP Computer Science**.  
        Ask questions about **Java**, **algorithms**, **object-oriented programming**, or anything else you're curious about.  
        I’ll explain things step by step, give you follow-up questions to deepen your understanding, and even draw out some fun visualizations!
        """
    )

    with gr.Tab("AP Computer Science"):
        chatbot = gr.Chatbot(label="💬 AP CS Tutor Chat", bubble_full_width=False, height=400)
        user_input = gr.Textbox(
            label="❓ What's your computer science question?",
            placeholder="e.g., What is recursion? How do for-loops work in Java?",
            lines=2
        )

        with gr.Row():
            model = gr.Dropdown(["gpt-4o", "gpt-4o-mini"], value="gpt-4o", label="🤖 Model Version")
            max_tokens = gr.Slider(800, 4000, value=2000, label="🧠 Max Tokens")
            temperature = gr.Slider(0, 1, value=0.6, label="🎯 Creativity (Temperature)")
            top_p = gr.Slider(0, 1, value=0.95, label="🔍 Top-P (Focus)")

        state = gr.State([])

        def cs_predict(message, history, model, max_tokens, temperature, top_p):
            full_response = ""
            for chunk in predict(
                message, history, "Computer Science", model, max_tokens, temperature, top_p
            ):
                full_response = chunk

            history.append([message, full_response])

            # Generate a fun image or diagram based on the CS concept
            image = generate_image(
                f"An educational diagram illustrating: {message}. Style: clean, colorful, helpful for high school students.",
                size="1024x1024"
            )
            return history, "", image

        image_output = gr.Image(label="AI-Generated Visual Explanation", type="pil")

        user_input.submit(
            cs_predict,
            inputs=[user_input, state, model, max_tokens, temperature, top_p],
            outputs=[chatbot, user_input, image_output]
        )

demo.launch()