import gradio as gr
import uuid
import subprocess
import threading
import os
import time
from fastapi import FastAPI
from fastapi.responses import FileResponse
import asyncio


# A simple in-memory dictionary to track task status.
# For a production system, you'd use a database or Redis.
tasks = {}

def run_video_generation(task_id: str, topic: str, context: str):
    """
    This function runs the main generation script in a separate process.
    """
    tasks[task_id]['status'] = 'running'
    
    # Sanitize topic to create a valid directory name
    file_prefix = "".join(c if c.isalnum() else "_" for c in topic.lower())
    output_dir = os.path.join("output", file_prefix)
    
    command = [
        "python", "generate_video.py",
        "--model", "openai/o3-mini",      # Or get from request
        "--topic", topic,
        "--context", context,
        "--output_dir", "output",
        "--use_langfuse" # Assuming you have secrets set
    ]

    try:
        # Using subprocess to run the existing script
        process = subprocess.run(command, check=True, capture_output=True, text=True)
        
        # Assume the final video is named based on the topic
        # Note: The actual video path might differ. This is an assumption.
        # You may need to parse the stdout from generate_video.py to get the exact path.
        video_path = None
        for file in os.listdir(output_dir):
            if file.endswith("_combined.mp4"):
                video_path = os.path.join(output_dir, file)
                break

        if video_path and os.path.exists(video_path):
            tasks[task_id]['status'] = 'completed'
            tasks[task_id]['video_path'] = video_path
        else:
            tasks[task_id]['status'] = 'failed'
            tasks[task_id]['error'] = "Video file not found after generation."
            tasks[task_id]['stdout'] = process.stdout
            tasks[task_id]['stderr'] = process.stderr
            
    except subprocess.CalledProcessError as e:
        tasks[task_id]['status'] = 'failed'
        tasks[task_id]['error'] = str(e)
        tasks[task_id]['stdout'] = e.stdout
        tasks[task_id]['stderr'] = e.stderr
    except Exception as e:
        tasks[task_id]['status'] = 'failed'
        tasks[task_id]['error'] = str(e)

def start_generation_thread(topic: str, context: str):
    if not topic or not context:
        return "Topic and Context cannot be empty.", "", None

    task_id = str(uuid.uuid4())
    tasks[task_id] = {'status': 'queued'}

    # Use a background thread to run the time-consuming task
    thread = threading.Thread(
        target=run_video_generation,
        args=(task_id, topic, context)
    )
    thread.start()
    
    return f"Task started. Your Task ID is: {task_id}", task_id, None


def check_status(task_id: str):
    if not task_id:
        return "Please provide a Task ID.", None
        
    task = tasks.get(task_id)
    if not task:
        return "Task not found.", None
    
    status = task.get('status')
    if status == 'completed':
        video_path = task.get('video_path')
        return f"Status: {status}", video_path
    elif status == 'failed':
        error = task.get('error', 'Unknown error')
        stdout = task.get('stdout', '')
        stderr = task.get('stderr', '')
        return f"Status: {status}\nError: {error}\nOutput: {stdout}\nStderr: {stderr}", None
    
    return f"Status: {status}", None

# We need a lightweight FastAPI app in the background to serve the video files.
# Gradio can't serve files directly from arbitrary paths in a secure way.
fastapi_app = FastAPI()

@fastapi_app.get("/videos/{task_id}")
def get_video(task_id: str):
    """
    Serves the final generated video file.
    """
    task = tasks.get(task_id)
    if not task or task.get('status') != 'completed':
        return {"error": "Task not completed or not found"}
    
    video_path = task.get('video_path')
    if not os.path.exists(video_path):
        return {"error": "Video file not found."}
        
    return FileResponse(video_path, media_type="video/mp4", filename=os.path.basename(video_path))


# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Theorem-Explain-Agent Video Generation")
    gr.Markdown("Start a video generation task and check its status.")

    with gr.Tab("Start Generation"):
        topic_input = gr.Textbox(label="Topic", placeholder="e.g., The Pythagorean Theorem")
        context_input = gr.Textbox(label="Context", placeholder="A short explanation of the theorem.")
        start_button = gr.Button("Generate Video")
        
        with gr.Column():
            task_id_output = gr.Textbox(label="Task ID", interactive=False)
            status_output_start = gr.Textbox(label="Status", interactive=False)
    
    with gr.Tab("Check Status"):
        task_id_input = gr.Textbox(label="Task ID", placeholder="Enter the Task ID you received.")
        check_button = gr.Button("Check Status")
        
        with gr.Column():
            status_output_check = gr.Textbox(label="Status", interactive=False)
            video_output = gr.Video(label="Generated Video")

    # Actions
    start_button.click(
        fn=start_generation_thread,
        inputs=[topic_input, context_input],
        outputs=[status_output_start, task_id_output, video_output] # Clear video on new task
    )
    
    check_button.click(
        fn=check_status,
        inputs=[task_id_input],
        outputs=[status_output_check, video_output]
    )

    gr.Markdown("### How to Use")
    gr.Markdown(
        "1.  Enter a `Topic` and `Context` in the 'Start Generation' tab and click 'Generate Video'.\n"
        "2.  Copy the `Task ID` that appears.\n"
        "3.  Go to the 'Check Status' tab, paste the `Task ID`, and click 'Check Status' periodically.\n"
        "4.  When the generation is complete, the video will appear."
    )

# To run both Gradio and FastAPI, we mount the FastAPI app into Gradio's internal FastAPI app.
app = gr.mount_ όπου(demo, fastapi_app, path="/")