app.py

"""
🎬 Viral Video Generator - Main Application
Generates viral short-form videos from any topic using:
- Web scraping for research
- Ollama cloud models for script writing
- edge-tts for voiceover
- Remotion + FFmpeg for video rendering
- Multi-source free stock images (Pexels, Openverse, Pixabay, Picsum)
"""
import gradio as gr
import os
import json
import time
import shutil
import traceback
from pathlib import Path

from scraper import research_topic
from llm import generate_viral_script, generate_title, OLLAMA_CLOUD_MODELS, DEFAULT_MODEL, get_model_name
from tts import generate_voiceover, generate_scene_audio_segments, combine_scene_audio, VOICE_OPTIONS
from media_downloader import download_assets_for_scenes, get_orientation_from_size
from video_render import (
    start_asset_server,
    get_video_dimensions,
    prepare_remotion_props,
    render_video,
    render_video_with_ffmpeg_fallback,
)

# =================== CONSTANTS ===================
WORK_DIR = "/home/user/app/tmp"
ASSETS_DIR = "/home/user/app/assets"
OUTPUT_DIR = "/home/user/app/output"

SIZE_OPTIONS = [
    "9:16 (Vertical/Reels)",
    "16:9 (Horizontal/YouTube)",
    "1:1 (Square/Posts)",
    "4:5 (Portrait/Instagram)",
    "4:3 (Standard)",
]

THEME_OPTIONS = ["white", "black"]

SPEED_OPTIONS = {
    "Normal": "+0%",
    "Slightly Faster": "+10%",
    "Fast": "+20%",
    "Slow": "-10%",
    "Very Slow": "-20%",
}

# =================== PIPELINE ===================

def validate_inputs(ollama_key, model, custom_model, topic):
    errors = []
    if not ollama_key or len(ollama_key) < 5:
        errors.append("❌ Ollama API key is required (get from ollama.com/settings)")
    if not model and not (custom_model and custom_model.strip()):
        errors.append("❌ Please select an AI model or type a custom model name")
    if not topic or len(topic) < 3:
        errors.append("❌ Please enter a topic (at least 3 characters)")
    return errors


def create_session_dir():
    session_id = f"session_{int(time.time())}_{os.getpid()}"
    session_dir = os.path.join(WORK_DIR, session_id)
    for sub in ["", "audio", "images", "output"]:
        os.makedirs(os.path.join(session_dir, sub), exist_ok=True)
    return session_dir


def generate_video_pipeline(
    ollama_key, ollama_base_url, pexels_key, pixabay_key,
    model, custom_model, topic, additional_notes,
    theme, size, channel_name, voice, speech_rate,
    num_scenes, target_duration, progress=gr.Progress(),
):
    errors = validate_inputs(ollama_key, model, custom_model, topic)
    if errors:
        raise gr.Error("\n".join(errors))

    resolved_model = get_model_name(model, custom_model)
    session_dir = create_session_dir()
    log_messages = []

    def log(msg):
        log_messages.append(msg)
        print(msg)

    try:
        width, height = get_video_dimensions(size)
        orientation = get_orientation_from_size(size)
        rate = SPEED_OPTIONS.get(speech_rate, "+0%")

        if not ollama_base_url:
            ollama_base_url = "https://ollama.com/v1"
        ollama_base_url = ollama_base_url.strip().rstrip("/")
        if not ollama_base_url.endswith("/v1"):
            ollama_base_url += "/v1"

        # STEP 1: Research
        progress(0.05, desc="🔍 Researching topic...")
        log(f"📋 Topic: {topic}")
        log(f"🤖 Model: {resolved_model}")
        log("🔍 Starting web research...")
        research = research_topic(topic, additional_notes, progress_callback=lambda msg: log(msg))
        log(f"✅ Research complete: {len(research['sources'])} sources, {len(research['key_facts'])} key facts")

        # STEP 2: Script
        progress(0.20, desc="✍️ Writing viral script...")
        log(f"✍️ Generating script with {resolved_model}...")
        script_data = generate_viral_script(
            topic=topic, research_text=research["full_text"], key_facts=research["key_facts"],
            model=resolved_model, api_key=ollama_key, additional_notes=additional_notes,
            theme=theme, num_scenes=int(num_scenes), target_duration=int(target_duration),
            base_url=ollama_base_url,
        )
        scenes = script_data["scenes"]
        log(f"✅ Script generated: {len(scenes)} scenes")
        for s in scenes:
            log(f"  Scene {s['scene_number']}: {s['hook']} ({s['duration']}s)")

        progress(0.30, desc="📝 Creating title...")
        title = generate_title(topic, resolved_model, ollama_key, ollama_base_url)
        log(f"📝 Title: {title}")

        # STEP 3: Voiceover
        progress(0.35, desc="🎤 Generating voiceover...")
        log(f"🎤 Generating voiceover with {voice}...")
        audio_dir = os.path.join(session_dir, "audio")
        scene_audios = generate_scene_audio_segments(scenes=scenes, voice_name=voice, output_dir=audio_dir, rate=rate)
        combined_audio_path = os.path.join(audio_dir, "full_narration.mp3")
        audio_result = combine_scene_audio(scene_audios=scene_audios, output_path=combined_audio_path, gap_ms=400)
        log(f"✅ Voiceover generated: {audio_result['duration_seconds']:.1f}s")
        scene_timestamps = audio_result.get("scene_timestamps", [])
        for ts in scene_timestamps:
            idx = ts["scene_index"]
            if idx < len(scenes):
                scenes[idx]["actual_duration_ms"] = ts["duration_ms"]

        # STEP 4: Images
        progress(0.45, desc="🖼️ Downloading stock images...")
        sources = []
        if pexels_key: sources.append("Pexels")
        sources.append("Openverse")
        if pixabay_key: sources.append("Pixabay")
        sources.append("Picsum")
        log(f"🖼️ Searching images from: {' → '.join(sources)}")
        images_dir = os.path.join(session_dir, "images")
        scenes = download_assets_for_scenes(
            scenes=scenes, pexels_key=pexels_key or "", save_dir=images_dir,
            orientation=orientation, pixabay_key=pixabay_key or "",
            progress_callback=lambda msg: log(msg),
        )
        images_found = sum(1 for s in scenes if s.get("localImagePath") and os.path.exists(s.get("localImagePath", "")))
        log(f"✅ Downloaded {images_found}/{len(scenes)} images")

        # STEP 5: Render
        progress(0.60, desc="🎥 Rendering video...")
        log("🎥 Starting video render...")
        output_path = os.path.join(session_dir, "output", "viral_video.mp4")
        fps = 30
        total_duration_sec = audio_result["duration_seconds"]
        total_frames = int(total_duration_sec * fps) + fps
        render_success = False
        bundle_path = "/home/user/app/dist/bundle"

        if os.path.exists(bundle_path):
            try:
                log("🎥 Attempting Remotion render...")
                asset_base_url = start_asset_server(session_dir)
                audio_url = f"{asset_base_url}/audio/full_narration.mp3"
                props = prepare_remotion_props(
                    title=title, scenes=scenes, audio_url=audio_url, theme=theme,
                    channel_name=channel_name or "", width=width, height=height,
                    duration_frames=total_frames, fps=fps,
                )
                render_video(props=props, output_path=output_path, bundle_path=bundle_path,
                            concurrency=1, timeout_ms=300000,
                            progress_callback=lambda msg: (log(msg), progress(0.75, desc=msg)))
                render_success = True
                log("✅ Remotion render complete!")
            except Exception as e:
                log(f"⚠️ Remotion render failed: {str(e)[:200]}")
                log("🔄 Falling back to FFmpeg renderer...")
        else:
            log("ℹ️ Remotion bundle not found, using FFmpeg renderer")

        if not render_success:
            progress(0.65, desc="🎥 Rendering with FFmpeg...")
            render_video_with_ffmpeg_fallback(
                scenes=scenes, audio_path=combined_audio_path, output_path=output_path,
                width=width, height=height, fps=fps, theme=theme, title=title,
                channel_name=channel_name or "", scene_timestamps=scene_timestamps,
                progress_callback=lambda msg: (log(msg), progress(0.80, desc=msg)),
            )
            log("✅ FFmpeg render complete!")

        progress(0.95, desc="✅ Finalizing...")

        script_display = f"# 🎬 {title}\n\n**Topic:** {topic}\n**Model:** {resolved_model}\n"
        script_display += f"**Duration:** {total_duration_sec:.1f}s | **Scenes:** {len(scenes)}\n"
        script_display += f"**Theme:** {theme} | **Size:** {size}\n\n---\n\n"
        for s in scenes:
            script_display += f"### Scene {s['scene_number']} ({s['duration']}s)\n"
            script_display += f"**🎯 Hook:** {s['hook']}\n\n**🎤 Narration:** {s['narration']}\n\n"
            script_display += f"**🖼️ Visual:** {s['visual']}\n\n---\n\n"

        research_display = f"# 📚 Research Summary\n\n**Sources found:** {len(research['sources'])}\n\n"
        for i, src in enumerate(research["sources"][:5]):
            research_display += f"{i+1}. {src['url']}\n   > {src['preview'][:150]}...\n\n"
        research_display += f"\n**Key Facts ({len(research['key_facts'])}):**\n\n"
        for fact in research["key_facts"][:10]:
            research_display += f"- {fact}\n"

        progress(1.0, desc="✅ Done!")
        return (output_path, script_display, research_display, "\n".join(log_messages), gr.update(visible=True))

    except gr.Error:
        raise
    except Exception as e:
        log(f"❌ Error: {str(e)}\n\n{traceback.format_exc()}")
        raise gr.Error(f"Pipeline failed: {str(e)}\n\nCheck the Logs tab for details.")


# =================== GRADIO UI ===================

CUSTOM_CSS = """
    .main-header { text-align:center; padding:20px; background:linear-gradient(135deg,#667eea 0%,#764ba2 100%); border-radius:15px; margin-bottom:20px; color:white; }
    .main-header h1 { color:white !important; font-size:2.5em !important; margin:0 !important; }
    .main-header p { color:rgba(255,255,255,0.9) !important; font-size:1.1em !important; }
    .step-header { font-size:1.2em; font-weight:600; color:#6366f1; border-bottom:2px solid #6366f1; padding-bottom:8px; margin-bottom:15px; }
    footer { display:none !important; }
"""

def build_ui():
    with gr.Blocks(title="🎬 Viral Video Generator") as demo:
        gr.HTML('<div class="main-header"><h1>🎬 Viral Video Generator</h1><p>Research → Script → Voiceover → Video — All AI-Powered</p></div>')

        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML('<div class="step-header">🔑 Step 1: API Keys</div>')
                with gr.Group():
                    ollama_key = gr.Textbox(label="Ollama API Key *", type="password", placeholder="Enter your Ollama API key...", info="Required. Get from ollama.com/settings → API Keys")
                    ollama_base_url = gr.Textbox(label="Ollama Base URL", value="https://ollama.com/v1", info="Default Ollama cloud. Change for custom endpoints.")

                gr.HTML('<p style="margin:8px 0 4px;font-weight:600;color:#6366f1;">🖼️ Image API Keys (optional — free sources used automatically)</p>')
                with gr.Group():
                    pexels_key = gr.Textbox(label="Pexels API Key (optional)", type="password", placeholder="Free at pexels.com/api", info="Optional. Without it, Openverse + Picsum used (no key needed).")
                    pixabay_key = gr.Textbox(label="Pixabay API Key (optional)", type="password", placeholder="Free at pixabay.com/api/docs", info="Optional. Extra source with videos support.")

                gr.HTML('<div class="step-header">📝 Step 2: Content</div>')
                with gr.Group():
                    model = gr.Dropdown(label="AI Model (Ollama Cloud)", choices=OLLAMA_CLOUD_MODELS, value=DEFAULT_MODEL, info="Top = 🟢 Free tier. Bottom = 🔵 may need Pro ($20/mo).")
                    custom_model = gr.Textbox(label="Custom Model Name (overrides dropdown)", placeholder="e.g., my-model:7b — type any model not in list", info="If filled, overrides dropdown selection.")
                    topic = gr.Textbox(label="Topic *", placeholder="e.g., 'How black holes are formed'", lines=2, info="What should the video be about?")
                    additional_notes = gr.Textbox(label="Additional Notes (optional)", placeholder="e.g., 'Focus on recent discoveries'", lines=2)

                gr.HTML('<div class="step-header">🎨 Step 3: Style & Format</div>')
                with gr.Group():
                    with gr.Row():
                        theme = gr.Radio(label="Theme", choices=THEME_OPTIONS, value="white")
                        size = gr.Dropdown(label="Video Size", choices=SIZE_OPTIONS, value="9:16 (Vertical/Reels)")
                    with gr.Row():
                        voice = gr.Dropdown(label="Voice", choices=list(VOICE_OPTIONS.keys()), value="Aria (US Female)")
                        speech_rate = gr.Dropdown(label="Speech Speed", choices=list(SPEED_OPTIONS.keys()), value="Normal")
                    channel_name = gr.Textbox(label="Channel Name (optional)", placeholder="e.g., TechExplained", info="Shown in intro and outro.")
                    with gr.Row():
                        num_scenes = gr.Slider(label="Number of Scenes", minimum=3, maximum=12, value=6, step=1)
                        target_duration = gr.Slider(label="Target Duration (seconds)", minimum=15, maximum=180, value=60, step=5)

                generate_btn = gr.Button("🚀 Generate Viral Video", variant="primary", size="lg")

            with gr.Column(scale=1):
                gr.HTML('<div class="step-header">🎬 Output</div>')
                with gr.Tabs():
                    with gr.Tab("📹 Video"):
                        video_output = gr.Video(label="Generated Video", height=500)
                        download_btn = gr.DownloadButton(label="⬇️ Download Video", visible=False, variant="secondary", size="lg")
                    with gr.Tab("📜 Script"):
                        script_output = gr.Markdown(value="*Script will appear here after generation...*")
                    with gr.Tab("📚 Research"):
                        research_output = gr.Markdown(value="*Research findings will appear here...*")
                    with gr.Tab("📋 Logs"):
                        logs_output = gr.Textbox(label="Pipeline Logs", lines=20, interactive=False, value="Waiting for generation to start...")

        gr.HTML("""<div style="text-align:center;padding:20px;opacity:0.7;font-size:0.9em;">
            <p>🖼️ <b>Image sources (auto-fallback):</b>
            <a href="https://www.pexels.com/api/" target="_blank">Pexels</a> →
            <a href="https://openverse.org" target="_blank">Openverse</a> (no key!) →
            <a href="https://pixabay.com/api/docs/" target="_blank">Pixabay</a> →
            <a href="https://picsum.photos" target="_blank">Picsum</a> (always works)</p>
            <p>⚡ <a href="https://ollama.com/search?c=cloud" target="_blank">Ollama Cloud</a> + Remotion + edge-tts + DuckDuckGo</p></div>""")

        def on_generate(ollama_key, ollama_base_url, pexels_key, pixabay_key, model, custom_model,
                        topic, additional_notes, theme, size, channel_name, voice, speech_rate,
                        num_scenes, target_duration, progress=gr.Progress()):
            results = generate_video_pipeline(ollama_key, ollama_base_url, pexels_key, pixabay_key,
                model, custom_model, topic, additional_notes, theme, size, channel_name,
                voice, speech_rate, num_scenes, target_duration, progress)
            video_path, script, research, logs, _ = results
            return (video_path, script, research, logs, gr.update(visible=True, value=video_path))

        generate_btn.click(
            fn=on_generate,
            inputs=[ollama_key, ollama_base_url, pexels_key, pixabay_key, model, custom_model,
                    topic, additional_notes, theme, size, channel_name, voice, speech_rate,
                    num_scenes, target_duration],
            outputs=[video_output, script_output, research_output, logs_output, download_btn],
            concurrency_limit=2,
        )
    return demo


if __name__ == "__main__":
    for d in [WORK_DIR, ASSETS_DIR, OUTPUT_DIR]:
        os.makedirs(d, exist_ok=True)
    demo = build_ui()
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, show_error=True, css=CUSTOM_CSS)