import os
import json
import re
import gradio as gr
import asyncio
import logging
import subprocess
from serpapi import GoogleSearch
from pydantic import BaseModel
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.conditions import HandoffTermination, TextMentionTermination
from autogen_agentchat.teams import Swarm
from autogen_agentchat.ui import Console
from autogen_agentchat.messages import TextMessage, HandoffMessage, StructuredMessage
from autogen_ext.models.anthropic import AnthropicChatCompletionClient
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_ext.models.ollama import OllamaChatCompletionClient
from markdown_pdf import MarkdownPdf, Section
import traceback
import soundfile as sf
import tempfile
from pydub import AudioSegment

# Set up logging
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.FileHandler("lecture_generation.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Set up environment
OUTPUT_DIR = "outputs"
os.makedirs(OUTPUT_DIR, exist_ok=True)
XTTS_MODEL_DIR = "XTTS-v2"

# Define Pydantic model for slide data
class Slide(BaseModel):
    title: str
    content: str

class SlidesOutput(BaseModel):
    slides: list[Slide]

# Define search_web tool using SerpApi
def search_web(query: str, serpapi_key: str) -> str:
    try:
        params = {
            "q": query,
            "engine": "google",
            "api_key": serpapi_key,
            "num": 5
        }
        search = GoogleSearch(params)
        results = search.get_dict()
        
        if "error" in results:
            logger.error("SerpApi error: %s", results["error"])
            return f"Error during search: {results['error']}"
        
        if "organic_results" not in results or not results["organic_results"]:
            logger.info("No search results found for query: %s", query)
            return f"No results found for query: {query}"
        
        formatted_results = []
        for item in results["organic_results"][:5]:
            title = item.get("title", "No title")
            snippet = item.get("snippet", "No snippet")
            link = item.get("link", "No link")
            formatted_results.append(f"Title: {title}\nSnippet: {snippet}\nLink: {link}\n")
        
        formatted_output = "\n".join(formatted_results)
        logger.info("Successfully retrieved search results for query: %s", query)
        return f"Search results for {query}:\n{formatted_output}"
    
    except Exception as e:
        logger.error("Unexpected error during search: %s", str(e))
        return f"Unexpected error during search: {str(e)}"

# Define helper function for progress HTML
def html_with_progress(label, progress):
    return f"""
    <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
        <div style="width: 100%; background-color: #FFFFFF; border-radius: 10px; overflow: hidden; margin-bottom: 20px;">
            <div style="width: {progress}%; height: 30px; background-color: #4CAF50; border-radius: 10px;"></div>
        </div>
        <h2 style="font-style: italic; color: #555;">{label}</h2>
    </div>
    """

# Function to get model client based on selected service
def get_model_client(service, api_key):
    if service == "OpenAI-gpt-4o-2024-08-06":
        return OpenAIChatCompletionClient(model="gpt-4o-2024-08-06", api_key=api_key)
    elif service == "Anthropic-claude-3-sonnet-20240229":
        return AnthropicChatCompletionClient(model="claude-3-sonnet-20240229", api_key=api_key)
    elif service == "Google-gemini-1.5-flash":
        return OpenAIChatCompletionClient(model="gemini-1.5-flash", api_key=api_key)
    elif service == "Ollama-llama3.2":
        return OllamaChatCompletionClient(model="llama3.2")
    else:
        raise ValueError("Invalid service")

# Helper function to clean script text
def clean_script_text(script):
    if not script or not isinstance(script, str):
        logger.error("Invalid script input: %s", script)
        return None
    script = re.sub(r"\*\*Slide \d+:.*?\*\*", "", script)
    script = re.sub(r"\*\*|\*|\s*-\s*", "", script)
    script = re.sub(r"\[.*?\]", "", script)
    script = re.sub(r"Title:.*?\n|Content:.*?\n", "", script)
    script = re.sub(r"-", " ", script)
    script = script.replace("humanlike", "human like").replace("problemsolving", "problem solving")
    script = re.sub(r"\s+", " ", script).strip()
    script = re.sub(r"[^\w\s.,!?']", "", script)
    if len(script) < 10 or len(script) > 500:
        logger.error("Cleaned script length invalid (%d characters): %s", len(script), script)
        return None
    sentences = re.split(r"[.!?]+", script)
    sentences = [s.strip() for s in sentences if s.strip()]
    if len(sentences) < 1 or len(sentences) > 5:
        logger.error("Cleaned script has invalid sentence count (%d): %s", len(sentences), script)
        return None
    return script

# Helper function to validate and convert speaker audio (MP3 or WAV)
async def validate_and_convert_speaker_audio(speaker_audio):
    if not os.path.exists(speaker_audio):
        logger.error("Speaker audio file does not exist: %s", speaker_audio)
        return None
    
    try:
        # Check file extension
        ext = os.path.splitext(speaker_audio)[1].lower()
        if ext == ".mp3":
            logger.info("Converting MP3 to WAV: %s", speaker_audio)
            audio = AudioSegment.from_mp3(speaker_audio)
            # Convert to mono, 22050 Hz
            audio = audio.set_channels(1).set_frame_rate(22050)
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
                audio.export(temp_file.name, format="wav")
                speaker_wav = temp_file.name
        elif ext == ".wav":
            speaker_wav = speaker_audio
        else:
            logger.error("Unsupported audio format: %s", ext)
            return None
        
        # Validate WAV file
        data, samplerate = sf.read(speaker_wav)
        if samplerate < 16000 or samplerate > 48000:
            logger.error("Invalid sample rate for %s: %d Hz", speaker_wav, samplerate)
            return None
        if len(data) < 16000:
            logger.error("Speaker audio too short: %d frames", len(data))
            return None
        if data.ndim == 2:
            logger.info("Converting stereo WAV to mono: %s", speaker_wav)
            data = data.mean(axis=1)
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
                sf.write(temp_file.name, data, samplerate)
                speaker_wav = temp_file.name
        
        logger.info("Validated speaker audio: %s", speaker_wav)
        return speaker_wav
    
    except Exception as e:
        logger.error("Failed to validate or convert speaker audio %s: %s", speaker_audio, str(e))
        return None

# Helper function to generate audio using XTTS-v2 CLI
def generate_xtts_audio(text, speaker_wav, output_path):
    try:
        cmd = [
            "tts",
            "--model_name", "tts_models/multilingual/multi-dataset/xtts_v2",
            "--encoder_path", os.path.join(XTTS_MODEL_DIR, "model_se.pth.tar"),
            "--encoder_config", os.path.join(XTTS_MODEL_DIR, "config_se.json"),
            "--speaker_wav", speaker_wav,
            "--text", text,
            "--out_path", output_path,
            "--language_idx", "en"
        ]
        logger.debug("Executing tts command: %s", " ".join(cmd))
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True
        )
        logger.info("tts command succeeded for %s: %s", output_path, result.stdout)
        return True
    except subprocess.CalledProcessError as e:
        logger.error("tts command failed for %s: %s\n%s", output_path, e.stderr, e.stdout)
        return False
    except Exception as e:
        logger.error("Unexpected error running tts for %s: %s", output_path, str(e))
        return False

# Helper function to extract JSON from messages
def extract_json_from_message(message):
    if isinstance(message, TextMessage):
        content = message.content
        logger.debug("Extracting JSON from TextMessage: %s", content)
        if not isinstance(content, str):
            logger.warning("TextMessage content is not a string: %s", content)
            return None
        # Try standard JSON block
        pattern = r"```json\s*(.*?)\s*```"
        match = re.search(pattern, content, re.DOTALL)
        if match:
            try:
                return json.loads(match.group(1))
            except json.JSONDecodeError as e:
                logger.error("Failed to parse JSON from TextMessage: %s, Content: %s", e, content)
        # Fallback: Try raw JSON array
        json_pattern = r"\[\s*\{.*?\}\s*\]"
        match = re.search(json_pattern, content, re.DOTALL)
        if match:
            try:
                return json.loads(match.group(0))
            except json.JSONDecodeError as e:
                logger.error("Failed to parse fallback JSON from TextMessage: %s, Content: %s", e, content)
        # Fallback: Try any JSON-like structure
        try:
            parsed = json.loads(content)
            if isinstance(parsed, (list, dict)):
                logger.info("Parsed JSON from raw content: %s", parsed)
                return parsed
        except json.JSONDecodeError:
            pass
        logger.warning("No JSON found in TextMessage content: %s", content)
        return None
    
    elif isinstance(message, StructuredMessage):
        content = message.content
        logger.debug("Extracting JSON from StructuredMessage: %s", content)
        try:
            if isinstance(content, BaseModel):
                content_dict = content.dict()
                return content_dict.get("slides", content_dict)
            return content
        except Exception as e:
            logger.error("Failed to extract JSON from StructuredMessage: %s, Content: %s", e, content)
            return None
    
    elif isinstance(message, HandoffMessage):
        logger.debug("Extracting JSON from HandoffMessage context")
        for ctx_msg in message.context:
            if hasattr(ctx_msg, "content"):
                content = ctx_msg.content
                logger.debug("Handoff context message content: %s", content)
                if isinstance(content, str):
                    pattern = r"```json\s*(.*?)\s*```"
                    match = re.search(pattern, content, re.DOTALL)
                    if match:
                        try:
                            return json.loads(match.group(1))
                        except json.JSONDecodeError as e:
                            logger.error("Failed to parse JSON from HandoffMessage context: %s, Content: %s", e, content)
                    json_pattern = r"\[\s*\{.*?\}\s*\]"
                    match = re.search(json_pattern, content, re.DOTALL)
                    if match:
                        try:
                            return json.loads(match.group(0))
                        except json.JSONDecodeError as e:
                            logger.error("Failed to parse fallback JSON from HandoffMessage context: %s, Content: %s", e, content)
                    try:
                        parsed = json.loads(content)
                        if isinstance(parsed, (list, dict)):
                            logger.info("Parsed JSON from raw HandoffMessage context: %s", parsed)
                            return parsed
                    except json.JSONDecodeError:
                        pass
                elif isinstance(content, dict):
                    return content.get("slides", content)
        logger.warning("No JSON found in HandoffMessage context")
        return None
    
    logger.warning("Unsupported message type for JSON extraction: %s", type(message))
    return None

# Function to generate Markdown and convert to PDF (landscape, centered)
def generate_slides_pdf(slides):
    pdf = MarkdownPdf(page_size="A4", orientation="landscape")
    for slide in slides:
        content_lines = slide['content'].replace('\n', '\n\n')
        markdown_content = f"""
<div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; text-align: center; padding: 20px;">
# {slide['title']}

*Prof. AI Feynman*  
*Princeton University, April 26th, 2025*

{content_lines}
</div>

---
"""
        pdf.add_section(Section(markdown_content, toc=False))
    
    pdf_file = os.path.join(OUTPUT_DIR, "slides.pdf")
    pdf.save(pdf_file)
    
    logger.info("Generated PDF slides (landscape): %s", pdf_file)
    return pdf_file

# Async function to update audio preview
async def update_audio_preview(audio_file):
    if audio_file:
        logger.info("Updating audio preview for file: %s", audio_file)
        return audio_file
    return None

# Async function to generate lecture materials and audio
async def on_generate(api_service, api_key, serpapi_key, title, topic, instructions, lecture_type, speaker_audio, num_slides):
    if not serpapi_key:
        yield f"""
        <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
            <h2 style="color: #d9534f;">SerpApi key required</h2>
            <p style="margin-top: 20px;">Please provide a valid SerpApi key and try again.</p>
        </div>
        """
        return

    model_client = get_model_client(api_service, api_key)
    
    research_agent = AssistantAgent(
        name="research_agent",
        model_client=model_client,
        handoffs=["slide_agent"],
        system_message="You are a Research Agent. Use the search_web tool to gather information on the topic and keywords from the initial message. Summarize the findings concisely in a single message, then use the handoff_to_slide_agent tool to pass the task to the Slide Agent. Do not produce any other output.",
        tools=[search_web]
    )
    slide_agent = AssistantAgent(
        name="slide_agent",
        model_client=model_client,
        handoffs=["script_agent"],
        system_message=f"""
You are a Slide Agent. Using the research from the conversation history and the number of content slides ({num_slides}) specified in the initial message, generate exactly {num_slides} content slides, then add a quiz slide, an assignment slide, and a thank-you slide. Output ONLY a JSON array wrapped in ```json ... ``` in a TextMessage, where each slide is an object with 'title' and 'content' keys. Do not include any explanatory text, comments, or other messages. Ensure the JSON is valid and complete before proceeding. After outputting the JSON, use the handoff_to_script_agent tool to pass the task to the Script Agent.
Example output:
```json
[
    {{"title": "Slide 1", "content": "Content for slide 1"}},
    {{"title": "Quiz", "content": "Quiz questions"}},
    {{"title": "Assignment", "content": "Assignment details"}},
    {{"title": "Thank You", "content": "Thank you message"}}
]
```""",
        output_content_type=None,
        reflect_on_tool_use=False
    )
    script_agent = AssistantAgent(
        name="script_agent",
        model_client=model_client,
        handoffs=["feynman_agent"],
        system_message=f"""
You are a Script Agent. Access the JSON array of {num_slides + 3} slides from the conversation history. Generate a narration script (1-2 sentences) for each slide, summarizing its content. Output ONLY a JSON array wrapped in ```json ... ``` with exactly {num_slides + 3} strings, one script per slide, in the same order. Ensure the JSON is valid and complete. After outputting, use the handoff_to_feynman_agent tool. If scripts cannot be generated, retry once.
Example for 1 content slide:
```json
[
    "This slide covers the main topic.",
    "Answer these quiz questions.",
    "Complete this assignment.",
    "Thank you for attending."
]
```""",
        output_content_type=None,
        reflect_on_tool_use=False
    )
    feynman_agent = AssistantAgent(
        name="feynman_agent",
        model_client=model_client,
        handoffs=[],
        system_message="You are Agent Feynman. Review the slides and scripts from the conversation history to ensure coherence and completeness. Output a confirmation message summarizing the number of slides and scripts received. If slides or scripts are missing or invalid, report the issue clearly. Use 'TERMINATE' to signal completion. Example: 'Received 5 slides and 5 scripts. Lecture is coherent. TERMINATE'"
    )
    
    swarm = Swarm(
        participants=[research_agent, slide_agent, script_agent, feynman_agent],
        termination_condition=HandoffTermination(target="user") | TextMentionTermination("TERMINATE")
    )
    
    progress = 0
    label = "Research: in progress..."
    yield html_with_progress(label, progress)
    await asyncio.sleep(0.1)
    
    initial_message = f"""
    Lecture Title: {title}
    Topic: {topic}
    Additional Instructions: {instructions}
    Audience: {lecture_type}
    Number of Content Slides: {num_slides}
    Please start by researching the topic.
    """
    logger.info("Starting lecture generation for topic: %s", topic)
    
    slides = None
    scripts = None
    error_html = """
    <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
        <h2 style="color: #d9534f;">Failed to generate lecture materials</h2>
        <p style="margin-top: 20px;">Please try again with different parameters or a different model.</p>
    </div>
    """
    
    try:
        logger.info("Research Agent starting...")
        task_result = await Console(swarm.run_stream(task=initial_message))
        logger.info("Swarm execution completed")
        
        script_retry_count = 0
        max_script_retries = 2
        
        for message in task_result.messages:
            source = getattr(message, 'source', getattr(message, 'sender', None))
            logger.debug("Processing message from %s, type: %s, content: %s", source, type(message), message.to_text() if hasattr(message, 'to_text') else str(message))
            
            if isinstance(message, HandoffMessage):
                logger.info("Handoff from %s to %s, Context: %s", source, message.target, message.context)
                if source == "research_agent" and message.target == "slide_agent":
                    progress = 25
                    label = "Slides: generating..."
                    yield html_with_progress(label, progress)
                    await asyncio.sleep(0.1)
                elif source == "slide_agent" and message.target == "script_agent":
                    if slides is None:
                        logger.warning("Slide Agent handoff without slides JSON")
                        extracted_json = extract_json_from_message(message)
                        if extracted_json:
                            slides = extracted_json
                            logger.info("Extracted slides JSON from HandoffMessage context: %s", slides)
                    if slides is None:
                        label = "Slides: failed to generate..."
                        yield html_with_progress(label, progress)
                        await asyncio.sleep(0.1)
                    progress = 50
                    label = "Scripts: generating..."
                    yield html_with_progress(label, progress)
                    await asyncio.sleep(0.1)
                elif source == "script_agent" and message.target == "feynman_agent":
                    if scripts is None:
                        logger.warning("Script Agent handoff without scripts JSON")
                        extracted_json = extract_json_from_message(message)
                        if extracted_json:
                            scripts = extracted_json
                            logger.info("Extracted scripts JSON from HandoffMessage context: %s", scripts)
                    progress = 75
                    label = "Review: in progress..."
                    yield html_with_progress(label, progress)
                    await asyncio.sleep(0.1)
            
            elif source == "research_agent" and isinstance(message, TextMessage) and "handoff_to_slide_agent" in message.content:
                logger.info("Research Agent completed research")
                progress = 25
                label = "Slides: generating..."
                yield html_with_progress(label, progress)
                await asyncio.sleep(0.1)
            
            elif source == "slide_agent" and isinstance(message, (TextMessage, StructuredMessage)):
                logger.debug("Slide Agent message received: %s", message.to_text())
                extracted_json = extract_json_from_message(message)
                if extracted_json:
                    slides = extracted_json
                    logger.info("Slide Agent generated %d slides: %s", len(slides), slides)
                    # Save slide content to individual files
                    for i, slide in enumerate(slides):
                        content_file = os.path.join(OUTPUT_DIR, f"slide_{i+1}_content.txt")
                        try:
                            with open(content_file, "w", encoding="utf-8") as f:
                                f.write(slide["content"])
                            logger.info("Saved slide content to %s: %s", content_file, slide["content"])
                        except Exception as e:
                            logger.error("Error saving slide content to %s: %s", content_file, str(e))
                    progress = 50
                    label = "Scripts: generating..."
                    yield html_with_progress(label, progress)
                    await asyncio.sleep(0.1)
                else:
                    logger.warning("No JSON extracted from slide_agent message: %s", message.to_text())
            
            elif source == "script_agent" and isinstance(message, (TextMessage, StructuredMessage)):
                logger.debug("Script Agent message received: %s", message.to_text())
                extracted_json = extract_json_from_message(message)
                if extracted_json:
                    scripts = extracted_json
                    logger.info("Script Agent generated scripts for %d slides: %s", len(scripts), scripts)
                    # Save raw scripts to individual files
                    for i, script in enumerate(scripts):
                        script_file = os.path.join(OUTPUT_DIR, f"slide_{i+1}_raw_script.txt")
                        try:
                            with open(script_file, "w", encoding="utf-8") as f:
                                f.write(script)
                            logger.info("Saved raw script to %s: %s", script_file, script)
                        except Exception as e:
                            logger.error("Error saving raw script to %s: %s", script_file, str(e))
                    progress = 75
                    label = "Scripts generated and saved. Reviewing..."
                    yield html_with_progress(label, progress)
                    await asyncio.sleep(0.1)
                else:
                    logger.warning("No JSON extracted from script_agent message: %s", message.to_text())
                    if script_retry_count < max_script_retries:
                        script_retry_count += 1
                        logger.info("Retrying script generation (attempt %d/%d)", script_retry_count, max_script_retries)
                        # Re-prompt script agent
                        retry_message = TextMessage(
                            content="Please generate scripts for the slides as per your instructions.",
                            source="user",
                            recipient="script_agent"
                        )
                        task_result.messages.append(retry_message)
                        continue
            
            elif source == "feynman_agent" and isinstance(message, TextMessage) and "TERMINATE" in message.content:
                logger.info("Feynman Agent completed lecture review: %s", message.content)
                progress = 100
                label = "Lecture materials ready. Generating audio..."
                yield html_with_progress(label, progress)
                await asyncio.sleep(0.1)
        
        logger.info("Slides state: %s", "Generated" if slides else "None")
        logger.info("Scripts state: %s", "Generated" if scripts else "None")
        if not slides or not scripts:
            error_message = f"Failed to generate {'slides and scripts' if not slides and not scripts else 'slides' if not slides else 'scripts'}"
            error_message += f". Received {len(slides) if slides else 0} slides and {len(scripts) if scripts else 0} scripts."
            logger.error("%s", error_message)
            logger.debug("Dumping all messages for debugging:")
            for msg in task_result.messages:
                source = getattr(msg, 'source', getattr(msg, 'sender', None))
                logger.debug("Message from %s, type: %s, content: %s", source, type(msg), msg.to_text() if hasattr(msg, 'to_text') else str(msg))
            yield f"""
            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
                <h2 style="color: #d9534f;">{error_message}</h2>
                <p style="margin-top: 20px;">Please try again with a different model (e.g., Anthropic-claude-3-sonnet-20240229) or simplify the topic/instructions.</p>
            </div>
            """
            return
        
        if not isinstance(scripts, list) or not all(isinstance(s, str) for s in scripts):
            logger.error("Scripts are not a list of strings: %s", scripts)
            yield f"""
            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
                <h2 style="color: #d9534f;">Invalid script format</h2>
                <p style="margin-top: 20px;">Scripts must be a list of strings. Please try again.</p>
            </div>
            """
            return
        
        if len(slides) != len(scripts):
            logger.error("Mismatch between number of slides (%d) and scripts (%d)", len(slides), len(scripts))
            yield f"""
            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
                <h2 style="color: #d9534f;">Mismatch in slides and scripts</h2>
                <p style="margin-top: 20px;">Generated {len(slides)} slides but {len(scripts)} scripts. Please try again.</p>
            </div>
            """
            return
        
        # Generate PDF from slides
        pdf_file = generate_slides_pdf(slides)
        pdf_path = f"/{pdf_file}"
        
        audio_files = []
        speaker_audio = speaker_audio if speaker_audio else "feynman.mp3"
        validated_speaker_wav = await validate_and_convert_speaker_audio(speaker_audio)
        if not validated_speaker_wav:
            logger.error("Invalid speaker audio after conversion, skipping TTS")
            yield f"""
            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
                <h2 style="color: #d9534f;">Invalid speaker audio</h2>
                <p style="margin-top: 20px;">Please upload a valid MP3 or WAV audio file and try again.</p>
            </div>
            """
            return
        
        # Verify XTTS-v2 model files
        required_files = [
            os.path.join(XTTS_MODEL_DIR, "model_se.pth.tar"),
            os.path.join(XTTS_MODEL_DIR, "config_se.json")
        ]
        for f in required_files:
            if not os.path.exists(f):
                logger.error("Missing XTTS-v2 model file: %s", f)
                yield f"""
                <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
                    <h2 style="color: #d9534f;">Missing XTTS-v2 model files</h2>
                    <p style="margin-top: 20px;">Please ensure XTTS-v2 is downloaded to {XTTS_MODEL_DIR} and try again.</p>
                </div>
                """
                return
        
        # Process audio generation sequentially with retries
        for i, script in enumerate(scripts):
            cleaned_script = clean_script_text(script)
            audio_file = os.path.join(OUTPUT_DIR, f"slide_{i+1}.wav")
            script_file = os.path.join(OUTPUT_DIR, f"slide_{i+1}_script.txt")
            
            # Save cleaned script
            try:
                with open(script_file, "w", encoding="utf-8") as f:
                    f.write(cleaned_script or "")
                logger.info("Saved cleaned script to %s: %s", script_file, cleaned_script)
            except Exception as e:
                logger.error("Error saving cleaned script to %s: %s", script_file, str(e))
            
            if not cleaned_script:
                logger.error("Skipping audio for slide %d due to empty or invalid script", i + 1)
                audio_files.append(None)
                continue
            
            max_retries = 2
            for attempt in range(max_retries + 1):
                try:
                    current_text = cleaned_script
                    if attempt > 0:
                        sentences = re.split(r"[.!?]+", cleaned_script)
                        sentences = [s.strip() for s in sentences if s.strip()][:2]
                        current_text = ". ".join(sentences) + "."
                        logger.info("Retry %d for slide %d with simplified text: %s", attempt, i + 1, current_text)
                    
                    success = generate_xtts_audio(current_text, validated_speaker_wav, audio_file)
                    if not success:
                        raise RuntimeError("tts command failed")
                    
                    logger.info("Generated audio for slide %d: %s", i + 1, audio_file)
                    audio_files.append(audio_file)
                    yield f"""
                    <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
                        <h2 style="font-style: italic; color: #555;">Generated audio for slide {i + 1}/{len(scripts)}...</h2>
                    </div>
                    """
                    await asyncio.sleep(0.1)
                    break
                except Exception as e:
                    logger.error("Error generating audio for slide %d (attempt %d): %s\n%s", i + 1, attempt, str(e), traceback.format_exc())
                    if attempt == max_retries:
                        logger.error("Max retries reached for slide %d, skipping", i + 1)
                        audio_files.append(None)
                        break
        
        audio_files = [f"/{f}" if f else None for f in audio_files]
        
        slides_info = json.dumps({"slides": [
            {"title": slide["title"], "content": slide["content"]}
            for slide in slides
        ], "audioFiles": audio_files})
        
        html_output = f"""
        <div id="lecture-container" style="height: 700px; border: 1px solid #ddd; border-radius: 8px; display: flex; flex-direction: column; justify-content: space-between;">
            <div id="slide-content" style="flex: 1; overflow: hidden;">
                <iframe id="pdf-viewer" src="https://mozilla.github.io/pdf.js/web/viewer.html?file={pdf_path}" style="width: 100%; height: 100%; border: none;"></iframe>
            </div>
            <div style="padding: 20px;">
                <div id="progress-bar" style="width: 100%; height: 5px; background-color: #ddd; border-radius: 2px; margin-bottom: 10px;">
                    <div id="progress-fill" style="width: {(1/len(slides)*100)}%; height: 100%; background-color: #4CAF50; border-radius: 2px;"></div>
                </div>
                <div style="display: flex; justify-content: center; margin-bottom: 10px;">
                    <button onclick="prevSlide()" style="border-radius: 50%; width: 40px; height: 40px; margin: 0 5px; font-size: 1.2em; cursor: pointer;">⏮</button>
                    <button onclick="togglePlay()" style="border-radius: 50%; width: 40px; height: 40px; margin: 0 5px; font-size: 1.2em; cursor: pointer;">⏯</button>
                    <button onclick="nextSlide()" style="border-radius: 50%; width: 40px; height: 40px; margin: 0 5px; font-size: 1.2em; cursor: pointer;">⏭</button>
                </div>
                <p id="slide-counter" style="text-align: center;">Slide 1 of {len(slides)}</p>
            </div>
        </div>
        <script>
            const lectureData = {slides_info};
            let currentSlide = 0;
            const totalSlides = lectureData.slides.length;
            const slideCounter = document.getElementById('slide-counter');
            const progressFill = document.getElementById('progress-fill');
            let audioElements = [];
            let currentAudio = null;
            const pdfViewer = document.getElementById('pdf-viewer');

            for (let i = 0; i < totalSlides; i++) {{
                if (lectureData.audioFiles && lectureData.audioFiles[i]) {{
                    const audio = new Audio(lectureData.audioFiles[i]);
                    audioElements.push(audio);
                }} else {{
                    audioElements.push(null);
                }}
            }}

            function updateSlide() {{
                pdfViewer.src = `https://mozilla.github.io/pdf.js/web/viewer.html?file={pdf_path}#page=${{currentSlide + 1}}`;
                slideCounter.textContent = `Slide ${{currentSlide + 1}} of ${{totalSlides}}`;
                progressFill.style.width = `${{(currentSlide + 1) / totalSlides * 100}}%`;

                if (currentAudio) {{
                    currentAudio.pause();
                    currentAudio.currentTime = 0;
                }}

                if (audioElements[currentSlide]) {{
                    currentAudio = audioElements[currentSlide];
                    currentAudio.play().catch(e => console.error('Audio play failed:', e));
                }} else {{
                    currentAudio = null;
                }}
            }}

            function prevSlide() {{
                if (currentSlide > 0) {{
                    currentSlide--;
                    updateSlide();
                }}
            }}

            function nextSlide() {{
                if (currentSlide < totalSlides - 1) {{
                    currentSlide++;
                    updateSlide();
                }}
            }}

            function togglePlay() {{
                if (!audioElements[currentSlide]) return;
                if (currentAudio.paused) {{
                    currentAudio.play().catch(e => console.error('Audio play failed:', e));
                }} else {{
                    currentAudio.pause();
                }}
            }}

            audioElements.forEach((audio, index) => {{
                if (audio) {{
                    audio.addEventListener('ended', () => {{
                        if (index < totalSlides - 1) {{
                            nextSlide();
                        }}
                    }});
                }}
            }});
        </script>
        """
        logger.info("Lecture generation completed successfully")
        yield html_output
    
    except Exception as e:
        logger.error("Error during lecture generation: %s\n%s", str(e), traceback.format_exc())
        yield f"""
        <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
            <h2 style="color: #d9534f;">Error during lecture generation</h2>
            <p style="margin-top: 10px; font-size: 16px;">{str(e)}</p>
            <p style="margin-top: 20px;">Please try again or adjust your inputs.</p>
        </div>
        """
        return

# Gradio interface
with gr.Blocks(title="Agent Feynman") as demo:
    gr.Markdown("# <center>Learn Anything With Professor AI Feynman</center>")
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group():
                title = gr.Textbox(label="Lecture Title", placeholder="e.g. Introduction to AI")
                topic = gr.Textbox(label="Topic", placeholder="e.g. Artificial Intelligence")
                instructions = gr.Textbox(label="Additional Instructions", placeholder="e.g. Focus on recent advancements")
                lecture_type = gr.Dropdown(["Conference", "University", "High school"], label="Audience", value="University")
                api_service = gr.Dropdown(
                    choices=[
                        "OpenAI-gpt-4o-2024-08-06",
                        "Anthropic-claude-3-sonnet-20240229",
                        "Google-gemini-1.5-flash",
                        "Ollama-llama3.2"
                    ],
                    label="Model",
                    value="Google-gemini-1.5-flash"
                )
                api_key = gr.Textbox(label="Model Provider API Key", type="password", placeholder="Not required for Ollama")
                serpapi_key = gr.Textbox(label="SerpApi Key", type="password", placeholder="Enter your SerpApi key")
                num_slides = gr.Slider(1, 20, step=1, label="Number of Content Slides", value=3)
                speaker_audio = gr.Audio(label="Speaker sample audio (MP3 or WAV)", type="filepath", elem_id="speaker-audio")
                audio_preview = gr.Audio(label="Audio Preview", interactive=False)
                generate_btn = gr.Button("Generate Lecture")
        with gr.Column(scale=2):
            default_slide_html = """
            <div style="display: flex; flex-direction: column; justify-content: center; align-items: center; height: 100%; min-height: 700px; padding: 20px; text-align: center; border: 1px solid #ddd; border-radius: 8px;">
                <h2 style="font-style: italic; color: #555;">Waiting for lecture content...</h2>
                <p style="margin-top: 10px; font-size: 16px;">Please Generate lecture content via the form on the left first before lecture begins</p>
            </div>
            """
            slide_display = gr.HTML(label="Lecture Slides", value=default_slide_html)
    
    speaker_audio.change(
        fn=update_audio_preview,
        inputs=speaker_audio,
        outputs=audio_preview
    )
    
    generate_btn.click(
        fn=on_generate,
        inputs=[api_service, api_key, serpapi_key, title, topic, instructions, lecture_type, speaker_audio, num_slides],
        outputs=[slide_display]
    )

if __name__ == "__main__":
    demo.launch()