# app.py
import gradio as gr
from groq import Groq
import os
import json

# Initialize Groq client
try:
    GROQ_API_KEY = os.environ["GROQ_API_KEY"]
    print("API Key:", GROQ_API_KEY)  # Debug print
    client = Groq(api_key=GROQ_API_KEY)
except KeyError as e:
    raise ValueError("GROQ_API_KEY not found in environment variables. Please set it in the Space settings under 'Variables'.")

# Define valid models (only the two specified models)
valid_models = [
    "qwen-2.5-32b",
    "qwen-2.5-coder-32b"
]

# Function to generate tutor output (lesson, question, feedback)
def generate_tutor_output(subject, difficulty, student_input, model):
    if not subject or not difficulty or not student_input:
        return json.dumps({
            "lesson": "Error: Please fill in all fields (subject, difficulty, and input).",
            "question": "No question available",
            "feedback": "No feedback available"
        })

    if model not in valid_models:
        return json.dumps({
            "lesson": f"Error: Invalid model selected: {model}. Please choose a valid model.",
            "question": "No question available",
            "feedback": "No feedback available"
        })

    prompt = f"""
    You are an expert tutor in {subject} at the {difficulty} level. 
    The student has provided the following input: "{student_input}"
    
    Please generate:
    1. A brief, engaging lesson on the topic (2-3 paragraphs)
    2. A thought-provoking question to check understanding
    3. Constructive feedback on the student's input
    
    Format your response as a JSON object with keys: "lesson", "question", "feedback"
    """

    try:
        print(f"Calling Groq API with model: {model}, subject: {subject}, difficulty: {difficulty}, input: {student_input}")
        completion = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": f"You are the world's best AI tutor, renowned for your ability to explain complex concepts in an engaging, clear, and memorable way with examples suitable for {difficulty} level students. Your expertise in {subject} is unparalleled, and you're adept at tailoring your teaching to {difficulty} level students. Your goal is to not just impart knowledge, but to inspire a love for learning and critical thinking.",
                },
                {
                    "role": "user",
                    "content": prompt,
                }
            ],
            model=model,
            max_tokens=1000,
        )
        response = completion.choices[0].message.content
        print(f"Groq API Response: {response}")  # Debug print
        return response
    except Exception as e:
        print(f"Groq API Error: {str(e)}")
        return json.dumps({
            "lesson": f"Error: Could not generate lesson. API error: {str(e)}",
            "question": "No question available",
            "feedback": "No feedback available due to API error"
        })

# Inline CSS to match the original styling
custom_css = """
body {
    font-family: 'Poppins', sans-serif;
}
h1 {
    font-size: 24px;
    margin-bottom: 20px;
}
label {
    font-weight: 600;
    margin-bottom: 10px;
    display: block;
}
select, input[type="text"], textarea {
    width: 100%;
    padding: 10px;
    margin-bottom: 15px;
    border: 1px solid #ddd;
    border-radius: 5px;
    font-size: 14px;
    font-family: 'Poppins', sans-serif;
}
button {
    background-color: #4CAF50;
    color: white;
    padding: 12px 20px;
    border: none;
    border-radius: 5px;
    font-size: 14px;
    font-weight: bold;
    cursor: pointer;
    transition: background-color 0.3s;
}
button:hover {
    background-color: #45a049;
}
.markdown-output {
    margin-top: 20px;
    padding: 15px;
    background-color: #f1f1f1;
    border-radius: 5px;
    min-height: 100px;
    white-space: pre-wrap;
}
.floating-button {
    position: fixed;
    bottom: 30px;
    left: 20px;
    background-color: #dbdad7;
    color: black;
    padding: 15px 25px;
    border: none;
    border-radius: 10px;
    font-size: 13px;
    font-weight: bold;
    cursor: pointer;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
    transition: background-color 0.3s, box-shadow 0.3s;
    text-decoration: none;
    font-family: 'Poppins', sans-serif;
    display: flex;
    align-items: center;
    justify-content: center;
}
.floating-button:hover {
    background-color: #dbdad7;
    box-shadow: 0 6px 12px rgba(0, 0, 0, 0.3);
}
.floating-button i {
    margin-right: 10px;
}
"""

# Gradio interface
with gr.Blocks(
    title="AI Tutor with Text and Visuals",
    css=custom_css
) as demo:
    gr.Markdown("# 🎓 Learn & Explore")

    # Section for generating Text-based output
    with gr.Row():
        with gr.Column(scale=2):
            subject = gr.Dropdown(
                ["Math", "Science", "History", "Geography", "Economics"], 
                label="Subject", 
                info="Choose the subject of your lesson"
            )
            difficulty = gr.Dropdown(
                ["Beginner", "Intermediate", "Advanced"], 
                label="Difficulty Level", 
                info="Select your difficulty level"
            )
            model_select = gr.Dropdown(
                valid_models,
                label="AI Model",
                value="qwen-2.5-32b",
                info="Select the AI model to use"
            )
            student_input = gr.Textbox(
                placeholder="Type your query here...", 
                label="Your Input", 
                info="Enter the topic you want to learn"
            )
            submit_button = gr.Button("Generate Lesson and Question", variant="primary")
        
        with gr.Column(scale=3):
            lesson_output = gr.Markdown(label="Lesson", elem_classes="markdown-output")
            question_output = gr.Markdown(label="Comprehension Question", elem_classes="markdown-output")
            feedback_output = gr.Markdown(label="Feedback", elem_classes="markdown-output")

    # Section for Visual output using iframe (unchanged)
    with gr.Row():
        with gr.Column(scale=2):
            gr.Markdown("## Image Generation")
            gr.Markdown("Use the embedded tool below to generate images.")
            submit_button_visual = gr.Button("Open Visual Tool", variant="primary")

        with gr.Column(scale=3):
            visual_output = gr.HTML(label="Image Generation Tool")

    gr.HTML("""
        <a href="https://discord.com/invite/EM4uHQrakm" class="floating-button" target="_blank">
            <i class="fab fa-discord"></i> Join Discord
        </a>
    """)

    gr.Markdown("""
    ### How to Use
    1. **Text Section**: Select a subject, difficulty level, and model, type your query, and click 'Generate Lesson and Question' to get your personalized lesson, question, and feedback.
    2. **Visual Section**: Click 'Open Visual Tool' to load the image generation tool, then use it to generate images.
    3. Review the AI-generated content to enhance your learning experience!

    *Example*: Try "Explain the water cycle" for a Beginner Science lesson, and use the image tool to generate "a photo of a rainforest ecosystem".
    """)

    def process_output(output):
        print(f"Raw API Output: {output}")  # Debug print
        try:
            parsed = json.loads(output)
            return parsed["lesson"], parsed["question"], parsed["feedback"]
        except Exception as e:
            print(f"JSON Parsing Error: {str(e)}")
            return "Error parsing output", "No question available", "No feedback available"

    def load_visual_tool():
        return """
        <iframe style="width: 100%; height: 500px;" allowfullscreen="true" frameborder="0" src="https://artificialguybr-realvisxl-free-demo.hf.space"></iframe>
        """

    # Generate Text-based Output
    submit_button.click(
        fn=lambda s, d, i, m: process_output(generate_tutor_output(s, d, i, m)),
        inputs=[subject, difficulty, student_input, model_select],
        outputs=[lesson_output, question_output, feedback_output]
    )

    # Load Visual Tool (iframe)
    submit_button_visual.click(
        fn=load_visual_tool,
        inputs=[],
        outputs=[visual_output]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)