import gradio as gr
from groq import Groq
import os
import threading
import tempfile
import logging
from moviepy.editor import TextClip, concatenate_videoclips, AudioFileClip, ColorClip

# Set up logging for debugging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Disable proxies to avoid previous 'proxies' error
os.environ["HTTP_PROXY"] = ""
os.environ["HTTPS_PROXY"] = ""

# Initialize Groq client with error handling
try:
    client = Groq(api_key=os.environ.get("GROQ_API_KEY", ""))
    logger.info("Groq client initialized successfully with API key: %s", "set" if os.environ.get("GROQ_API_KEY") else "not set")
except Exception as e:
    logger.error("Failed to initialize Groq client: %s", str(e))
    raise

# Load Text-to-Image Models with error handling
try:
    model1 = gr.load("models/prithivMLmods/SD3.5-Turbo-Realism-2.0-LoRA", fallback=None)
    logger.info("Model 1 loaded successfully: SD3.5-Turbo-Realism-2.0-LoRA")
except Exception as e:
    logger.error("Failed to load Model 1: %s", str(e))
    model1 = None  # Fallback to None if loading fails

try:
    model2 = gr.load("models/Purz/face-projection", fallback=None)
    logger.info("Model 2 loaded successfully: face-projection")
except Exception as e:
    logger.error("Failed to load Model 2: %s", str(e))
    model2 = None  # Fallback to None if loading fails

# Stop event for threading (image generation)
stop_event = threading.Event()

# Function to generate tutor output (lesson, question, feedback)
def generate_tutor_output(subject, difficulty, student_input):
    if not all([subject, difficulty, student_input]):
        return '{"lesson": "Please fill in all fields.", "question": "", "feedback": ""}'
    
    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:
        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 and giving math examples. Your expertise in {subject} is unparalleled, and you're adept at tailoring your teaching to {difficulty} level students."
            }, {
                "role": "user",
                "content": prompt,
            }],
            model="mixtral-8x7b-32768",
            max_tokens=1000,
        )
        return completion.choices[0].message.content
    except Exception as e:
        logger.error("Error in generate_tutor_output: %s", str(e))
        return '{"lesson": "Error generating lesson.", "question": "", "feedback": ""}'

# Function to generate images based on model selection
def generate_images(text, selected_model):
    stop_event.clear()
    if not text:
        return ["No text provided."] * 3

    if selected_model == "Model 1 (Turbo Realism)":
        model = model1
    elif selected_model == "Model 2 (Face Projection)":
        model = model2
    else:
        return ["Invalid model selection."] * 3

    if model is None:
        return ["Selected model is not available."] * 3

    results = []
    for i in range(3):
        if stop_event.is_set():
            return ["Image generation stopped by user."] * 3
        modified_text = f"{text} variation {i+1}"
        try:
            result = model(modified_text)
            results.append(result)
        except Exception as e:
            logger.error("Error generating image %d: %s", i+1, str(e))
            results.append(None)
    return results

# Function to generate text-to-video with voice
def generate_text_to_video(text):
    if not text:
        return "No text provided for video generation."
    
    try:
        narration_prompt = f"Convert this text to a natural-sounding narration: {text}"
        narration_response = client.chat.completions.create(
            messages=[{
                "role": "system",
                "content": "You are an AI voice generator that produces natural, human-like speech."
            }, {
                "role": "user",
                "content": narration_prompt,
            }],
            model="mixtral-8x7b-32768",
            max_tokens=500,
        )
        narration_text = narration_response.choices[0].message.content

        with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as temp_audio:
            audio_duration = len(narration_text.split()) / 2  # Rough estimate: 2 words/sec
            audio = ColorClip(size=(100, 100), color=(0, 0, 0), duration=audio_duration).set_audio(None)
            audio.write_audiofile(temp_audio.name, fps=44100, logger=None)

        clips = []
        words = narration_text.split()
        chunk_size = 10
        for i in range(0, len(words), chunk_size):
            chunk = " ".join(words[i:i + chunk_size])
            clip = TextClip(chunk, fontsize=50, color='white', size=(1280, 720), bg_color='black')
            clip = clip.set_duration(audio_duration / (len(words) / chunk_size))
            clips.append(clip)

        final_video = concatenate_videoclips(clips)
        audio_clip = AudioFileClip(temp_audio.name)
        final_video = final_video.set_audio(audio_clip)

        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_video:
            final_video.write_videofile(temp_video.name, fps=24, logger=None)
            video_path = temp_video.name

        os.unlink(temp_audio.name)
        return video_path
    except Exception as e:
        logger.error("Error generating video: %s", str(e))
        return f"Error generating video: {str(e)}"

# Set up the Gradio interface
with gr.Blocks(title="AI Tutor with Visuals") as demo:
    gr.Markdown("# 🎓 Your AI Tutor with Visuals & Images")

    with gr.Row():
        with gr.Column(scale=2):
            subject = gr.Dropdown(
                ["Math", "Science", "History", "Literature", "Code", "AI"], 
                label="Subject", 
                info="Choose the subject of your lesson",
                value="Math"
            )
            difficulty = gr.Radio(
                ["Beginner", "Intermediate", "Advanced"], 
                label="Difficulty Level", 
                info="Select your proficiency level",
                value="Beginner"
            )
            student_input = gr.Textbox(
                placeholder="Type your query here...", 
                label="Your Input", 
                info="Enter the topic you want to learn"
            )
            submit_button_text = gr.Button("Generate Lesson & Question", variant="primary")
        
        with gr.Column(scale=3):
            lesson_output = gr.Markdown(label="Lesson")
            question_output = gr.Markdown(label="Comprehension Question")
            feedback_output = gr.Markdown(label="Feedback")

    with gr.Row():
        with gr.Column(scale=2):
            model_selector = gr.Radio(
                ["Model 1 (Turbo Realism)", "Model 2 (Face Projection)"],
                label="Select Image Generation Model",
                value="Model 1 (Turbo Realism)"
            )
            submit_button_visual = gr.Button("Generate Visuals", variant="primary")
            submit_button_video = gr.Button("Generate Video with Voice", variant="primary")
        
        with gr.Column(scale=3):
            output1 = gr.Image(label="Generated Image 1")
            output2 = gr.Image(label="Generated Image 2")
            output3 = gr.Image(label="Generated Image 3")
            video_output = gr.Video(label="Generated Video with Voice")

    gr.Markdown("""
    ### How to Use
    1. **Text Section**: Select a subject and difficulty, type your query, and click 'Generate Lesson & Question'.
    2. **Visual Section**: Select the model, then click 'Generate Visuals' for 3 images or 'Generate Video with Voice' for a narrated video.
    3. Review the AI-generated content to enhance your learning experience!
    """)

    def process_output_text(subject, difficulty, student_input):
        try:
            tutor_output = generate_tutor_output(subject, difficulty, student_input)
            parsed = eval(tutor_output)  # Use json.loads in production
            return parsed["lesson"], parsed["question"], parsed["feedback"]
        except Exception as e:
            logger.error("Error parsing tutor output: %s", str(e))
            return "Error parsing output", "No question available", "No feedback available"

    def process_output_visual(text, selected_model):
        try:
            images = generate_images(text, selected_model)
            return images[0], images[1], images[2]
        except Exception as e:
            logger.error("Error in process_output_visual: %s", str(e))
            return None, None, None

    def process_output_video(text):
        try:
            video_path = generate_text_to_video(text)
            return video_path
        except Exception as e:
            logger.error("Error in process_output_video: %s", str(e))
            return None

    submit_button_text.click(
        fn=process_output_text,
        inputs=[subject, difficulty, student_input],
        outputs=[lesson_output, question_output, feedback_output]
    )
    submit_button_visual.click(
        fn=process_output_visual,
        inputs=[student_input, model_selector],
        outputs=[output1, output2, output3]
    )
    submit_button_video.click(
        fn=process_output_video,
        inputs=[student_input],
        outputs=[video_output]
    )

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