import asyncio
import os
import shutil
import subprocess
import tempfile
from typing import List
import logging

import streamlit as st
from moviepy import concatenate_videoclips, VideoFileClip
from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext
from pydantic_ai.models.gemini import GeminiModel
from pydantic_ai.providers.google_gla import GoogleGLAProvider
from config import api_key
import nest_asyncio
nest_asyncio.apply()
import re

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Configure your Gemini API key
gemini_llm = GeminiModel(
    'gemini-2.0-flash', provider=GoogleGLAProvider(api_key=api_key)
)


class ChapterDescription(BaseModel):
    """Describes a chapter in the video."""
    title: str = Field(description="Title of the chapter.")
    explanation: str = Field(description="Detailed explanation of the chapter's content, including how Manim should visualize it. Be very specific with Manim instructions, including animations, shapes, positions, colors, and timing. Include LaTeX for mathematical formulas. Specify scene transitions. Example: 'Create a number line. Animate a point moving along the number line to illustrate addition. Use Transform to show the equation changing. Transition to a new Scene.'")

class VideoOutline(BaseModel):
    """Describes the outline of the video."""
    title: str = Field(description="Title of the entire video.")
    chapters: List[ChapterDescription] = Field(description="List of chapters in the video.")

class ManimCode(BaseModel):
    """Describes the Manim code for a chapter."""
    code: str = Field(description="Complete Manim code for the chapter. Include all necessary imports. The code should create a single scene. Add comments to explain the code. Do not include any comments that are not valid Python comments. Ensure the code is runnable.")

outline_agent = Agent(
    model=gemini_llm,
    result_type=VideoOutline,
    system_prompt="""
    You are a video script writer. Your job is to create a clear and concise outline for an educational video explaining a concept.
    The video should have a title and a list of chapters (maximum 3). Each chapter should have a title and a detailed explanation.
    The explanation should be very specific about how the concept should be visualized using Manim. Include detailed instructions
    for animations, shapes, positions, colors, and timing. Use LaTeX for mathematical formulas. Specify scene transitions.
    Do not include code, only explanations.
    """
)

manim_agent = Agent(
    model=gemini_llm,
    result_type=ManimCode,
    system_prompt="""
    You are a Manim code generator. Your job is to create Manim code for a single chapter of a video, given a detailed explanation of the chapter's content and how it should be visualized.
    The code should be complete and runnable. Include all necessary imports. The code should create a single scene. Add comments to explain the code.
    Do not include any comments that are not valid Python comments. Ensure the code is runnable. Do not include any text outside of the code block.

    """
)

code_fixer_agent = Agent(
    model=gemini_llm,
    result_type=ManimCode,
    system_prompt="""
    You are a Manim code debugging expert. You will receive Manim code that failed to execute and the error message.
    Your task is to analyze the code and the error, identify the issue, and provide corrected, runnable Manim code.
    Ensure the corrected code addresses the error and still aims to achieve the visualization described in the original code.
    Include all necessary imports and ensure the code creates a single scene. Add comments to explain the changes you made.
    Do not include any comments that are not valid Python comments. Ensure the code is runnable. Do not include any text outside of the code block.
    """
)

def generate_manim_code(chapter_description: ChapterDescription) -> str:
    """Generates initial Manim code for a single chapter."""
    logging.info(f"Generating Manim code for chapter: {chapter_description.title}")
    result = manim_agent.run_sync(f"title: {chapter_description.title}. Explanation: {chapter_description.explanation}")
    return result.data.code

def fix_manim_code(error: str, current_code: str) -> str:
    """Attempts to fix the Manim code that resulted in an error."""
    logging.info(f"Attempting to fix Manim code due to error: {error}")
    result = code_fixer_agent.run_sync(f"Error: {error}\nCurrent Code: {current_code}")
    return result.data.code

def generate_video_outline(concept: str) -> VideoOutline:
    """Generates the video outline."""
    logging.info(f"Generating video outline for concept: {concept}")
    result = outline_agent.run_sync(concept)
    return result.data

def create_video_from_code(code: str, chapter_number: int) -> str:
    """Creates a video from Manim code and returns the video file path using subprocess.Popen."""
    temp_dir = None
    temp_file_name = None
    video_file_path = None

    try:
        # Create a temporary directory for Manim's output
        temp_dir = tempfile.mkdtemp()
        logging.info(f"Created temporary directory for Manim output: {temp_dir}")

        # Use tempfile to create a temporary Python file within the temporary directory
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, dir=temp_dir) as temp_file:
            temp_file.write(code)
            temp_file_name = temp_file.name
        logging.info(f"Created temporary Manim script: {temp_file_name}")

        process = None
        # Manim command needs to specify the output directory
        # The -o flag followed by a filename will output the video to that specific path relative to the current working directory,
        # or an absolute path if provided.
        # Alternatively, we can use -o with just the class name and specify the output directory using --output_directory
        # For simplicity, we'll try to explicitly specify the output path if the Manim version supports it,
        # or we'll ensure the manim command is executed from within the temporary directory.

        # We need to extract the class name to get Manim's default filename
        match = re.search(r"class\s+(\w+)\(Scene\):", code)
        if not match:
            raise ValueError(f"Could not extract class name from Manim code for chapter {chapter_number}")
        class_name = match.group(1)
        expected_output_filename = f"{class_name}.mp4"

        # Construct the command. We will run manim from the temporary directory.
        # Manim by default creates media/videos/<resolution>/<filename>.mp4
        # We want it in our temp_dir. The easiest way is to change the CWD.
        # So the manim command becomes `manim <script_name> -ql --media_dir <temp_dir>`
        # This will put the videos inside <temp_dir>/videos/<resolution>/
        command = [
            "manim",
            os.path.basename(temp_file_name), # Just the filename, as we're changing CWD
            "-ql",
            "--disable_caching",
            "--media_dir", # Specify the media directory directly
            temp_dir
        ]
        logging.info(f"Executing Manim command in '{temp_dir}': {' '.join(command)}")

        # Execute Manim from within the temporary directory
        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            cwd=temp_dir, # Crucial: set the current working directory for the subprocess
            shell=True,
            text=True
        )
        stdout, stderr = process.communicate(timeout=120)

        if process.returncode == 0:
            logging.info(f"Manim execution successful for chapter {chapter_number}.")
            logging.debug(f"Manim stdout:\n{stdout}")
            logging.debug(f"Manim stderr:\n{stderr}")
            # Manim 0.16.0+ changed output paths slightly.
            # It usually creates a structure like <media_dir>/videos/<quality>/<scene_name>.mp4
            # We explicitly set media_dir to temp_dir, so the video should be in:
            # temp_dir/videos/temp/480p15/<class_name>.mp4
            # Or if --output_file is used, it directly goes to temp_dir/<filename>.
            # With --media_dir, it will stick to its internal folder structure.
            manim_rendered_video_subdir = "videos/temp/480p15" # This is Manim's internal structure
            video_file_path_in_temp = os.path.join(temp_dir, manim_rendered_video_subdir, expected_output_filename)

            if not os.path.exists(video_file_path_in_temp):
                logging.error(f"Manim output file not found at expected path: {video_file_path_in_temp}")
                # As a fallback, try to find it directly in temp_dir if Manim changed its behavior or a custom flag was added
                fallback_path = os.path.join(temp_dir, expected_output_filename)
                if os.path.exists(fallback_path):
                    logging.info(f"Found video at fallback path: {fallback_path}")
                    video_file_path = fallback_path
                else:
                    raise FileNotFoundError(f"Manim output video file '{expected_output_filename}' not found at '{video_file_path_in_temp}' or '{fallback_path}' after successful execution. Check Manim's output configuration.")
            else:
                video_file_path = video_file_path_in_temp

        else:
            error_msg = f"Manim execution failed for chapter {chapter_number} with return code {process.returncode}:\nStdout:\n{stdout}\nStderr:\n{stderr}"
            logging.error(str(error_msg).split('\n')[-1])
            raise subprocess.CalledProcessError(process.returncode, command, output=stdout.encode(), stderr=stderr.encode())

    except subprocess.TimeoutExpired:
        logging.error(f"Manim process timed out for chapter {chapter_number}.")
        if process:
            process.kill()
        raise
    except FileNotFoundError:
        logging.error("Error: The 'manim' command was not found. Ensure Manim is installed and in your system's PATH.")
        raise
    except Exception as e:
        logging.error(f"An unexpected error occurred during video creation: {e}")
        raise
    finally:
        # Ensure the temporary file is deleted after Manim tries to use it
        if temp_file_name and os.path.exists(temp_file_name):
            os.remove(temp_file_name)
            logging.info(f"Deleted temporary Manim script: {temp_file_name}")
        # Clean up the entire temporary directory created for Manim output
        # It's important to do this AFTER the video file has been moved or processed
        if temp_dir and os.path.exists(temp_dir):
            try:
                # Need to wait until the video is processed/moved by the calling function.
                # For now, let's assume the video_file_path returned will be copied.
                # If we don't copy, we need to defer this cleanup.
                # Since the final combination happens in `generate_video`,
                # it's better to clean up the `temp_dir` there after all clips are loaded.
                # For now, we'll only clean up the script itself.
                pass # Shifting temp_dir cleanup to generate_video for safer handling
            except Exception as e:
                logging.error(f"Error cleaning up temporary directory {temp_dir}: {e}")

    # The video_file_path now points to a file within temp_dir.
    # The calling function (generate_video) will need to load this using moviepy,
    # and then we can clean up the temp_dir after all clips are loaded.
    return video_file_path

async def generate_video(concept: str):
    """Generates a video explanation for a given concept using Manim with error correction."""
    logging.info(f"Generating video for concept: {concept}")
    outline = generate_video_outline(concept)
    logging.info(f"Video outline: {outline}")

    video_files_paths = [] # Renamed for clarity, these are paths in temp dirs
    temp_dirs_to_clean = [] # To keep track of temporary directories to delete

    for i, chapter in enumerate(outline.chapters):
        logging.info(f"Processing chapter {i + 1}: {chapter.title}")
        manim_code = generate_manim_code(chapter)
        logging.debug(f"Generated Manim code for chapter {i + 1}:\n{manim_code}")

        success = False
        attempts = 0
        max_attempts = 2

        while attempts < max_attempts and not success:
            try:
                st.info(f"Attempting to render chapter {i + 1}: {chapter.title} (Attempt {attempts + 1}/{max_attempts})")
                # create_video_from_code will return the path within a temp dir
                video_file_path_for_chapter = create_video_from_code(manim_code, i + 1)
                video_files_paths.append(video_file_path_for_chapter)
                # We need to extract the base temp directory for cleanup later
                # It's usually the parent of 'videos/temp/480p15'
                # Find the root temp directory path from video_file_path_for_chapter
                # Example: /tmp/tmpxyz/videos/temp/480p15/MyScene.mp4 -> /tmp/tmpxyz
                # This assumes temp_dir is directly below the media/videos structure
                root_temp_dir = os.path.abspath(os.path.join(video_file_path_for_chapter, os.pardir, os.pardir, os.pardir))
                if root_temp_dir not in temp_dirs_to_clean:
                    temp_dirs_to_clean.append(root_temp_dir)

                logging.info(f"Video file created for chapter {i + 1}: {video_file_path_for_chapter}")
                success = True
            except subprocess.CalledProcessError as e:
                attempts += 1
                logging.error(f"Manim execution failed for chapter {i + 1} (Attempt {attempts}): {e}")
                st.warning(f"Manim rendering failed for chapter {i + 1}. Attempting to fix code...")
                manim_code = fix_manim_code(str(e), manim_code)
                logging.debug(f"Fixed Manim code (Attempt {attempts}):\n{manim_code}")
            except ValueError as e:
                logging.error(f"Error processing Manim code for chapter {i + 1}: {e}")
                st.error(f"Error processing chapter {i + 1}: {e}")
                return None
            except FileNotFoundError as e:
                logging.error(f"Manim or output file not found for chapter {i + 1}: {e}")
                st.error(f"Manim or its output file was not found for chapter {i + 1}: {e}. Please ensure Manim is installed and configured correctly.")
                return None
            except subprocess.TimeoutExpired:
                attempts += 1
                logging.error(f"Manim process timed out for chapter {i + 1}. Attempt {attempts}. Attempting to fix...")
                st.warning(f"Manim rendering timed out for chapter {i + 1}. Attempting to fix code...")
                manim_code = fix_manim_code(f"Manim process timed out.", manim_code)
                logging.debug(f"Fixed Manim code (Attempt {attempts}):\n{manim_code}")
            except Exception as e:
                logging.error(f"An unexpected error occurred during chapter {i + 1} processing: {e}")
                st.error(f"An unexpected error occurred for chapter {i + 1}: {e}")
                return None


        if not success:
            logging.error(f"Failed to generate video for chapter {i + 1} after {max_attempts} attempts. Skipping chapter.")
            st.warning(f"Failed to generate video for chapter {i + 1}. Skipping.")
            continue

    # Combine the video files
    final_video_path = None
    if video_files_paths:
        logging.info("Combining video files...")
        st.info("Combining all chapter videos...")
        clips = []
        try:
            for vf_path in video_files_paths:
                if os.path.exists(vf_path):
                    clips.append(VideoFileClip(vf_path))
                else:
                    logging.warning(f"Skipping non-existent video file for concatenation: {vf_path}")
                    st.warning(f"Skipping missing chapter video for final combination: {vf_path.split('/')[-1]}")

            if clips:
                final_video_path = f"final_video.mp4"
                final_clip = concatenate_videoclips(clips)
                final_clip.write_videofile(final_video_path, codec="libx264", audio_codec="aac")
                final_clip.close()
                logging.info(f"Final video created: {final_video_path}")
                st.success("Video generation complete!")
            else:
                logging.warning("No valid video files to combine.")
                st.warning("No valid video files were generated to combine.")
        except Exception as e:
            logging.error(f"Error combining video files: {e}")
            st.error(f"Error combining video files: {e}")
        finally:
            # Clean up all temporary directories after all clips are loaded and combined
            for clip in clips: # Ensure all clips are closed
                try:
                    clip.close()
                except Exception as close_e:
                    logging.warning(f"Error closing clip: {close_e}")

            for temp_dir_path in temp_dirs_to_clean:
                try:
                    if os.path.exists(temp_dir_path):
                        shutil.rmtree(temp_dir_path) # Recursively remove the temporary directory
                        logging.info(f"Deleted temporary directory: {temp_dir_path}")
                except Exception as e:
                    logging.error(f"Error deleting temporary directory {temp_dir_path}: {e}")