app.py

import streamlit as st
import time
import concurrent.futures
import os
import pandas as pd
from videodb import connect
from editor_agent import (
    upload_and_index, 
    get_highlights, 
    create_trailer_stream, 
    analyze_virality, 
    generate_linkedin_post,
    get_video_from_id
)

# --- Caching Wrappers ---
@st.cache_resource(show_spinner=False)
def cached_upload(url):
    return upload_and_index(url)

@st.cache_resource(show_spinner=False)
def cached_fetch(video_id):
    return get_video_from_id(video_id)

# --- Helper for Parallel Execution ---
def process_trailer_pipeline(video_obj, transcript_text, focus_topic):
    highlights = get_highlights(transcript_text, focus=focus_topic)
    stream_link = None
    if highlights:
        stream_link = create_trailer_stream(video_obj, highlights)
    return highlights, stream_link

# 1. Page Config
st.set_page_config(
    page_title="AI Director", 
    page_icon="⚡", 
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- 🎨 CUSTOM CSS FOR "COOL" LOOK ---
st.markdown("""
    <style>
    /* Gradient Button */
    div.stButton > button:first-child {
        background: linear-gradient(45deg, #FF4B2B, #FF416C);
        color: white;
        border: none;
        padding: 0.6rem 1rem;
        border-radius: 10px;
        font-weight: bold;
        transition: transform 0.2s;
    }
    div.stButton > button:first-child:hover {
        transform: scale(1.02);
        box-shadow: 0px 4px 15px rgba(255, 65, 108, 0.4);
    }
    
    /* Metrics Styling */
    [data-testid="stMetricValue"] {
        font-size: 2rem;
        color: #FF416C;
    }
    
    /* Sidebar Styling */
    [data-testid="stSidebar"] {
        background-color: #111111;
        border-right: 1px solid #222;
    }
    
    /* Info Box Styling */
    .stAlert {
        background-color: #1E1E1E;
        border: 1px solid #333;
        color: #EEE;
    }
    </style>
""", unsafe_allow_html=True)

# Header
col_logo, col_title = st.columns([1, 6])
#with col_logo:
    #st.markdown("## ⚡")
with col_title:
    st.markdown("#                                                            AI Director")
    st.caption("Your autonomous AI Director that turns raw footage into viral highlights & social posts.")

# Initialize State
if "video_obj" not in st.session_state:
    st.session_state.video_obj = None
if "analysis_done" not in st.session_state:
    st.session_state.analysis_done = False 

# --- SIDEBAR CONFIG ---
with st.sidebar:
    st.header("🎛️ Studio Config")
    
    # 1. Source Selection
    st.subheader("1. Video Source")
    
    # Fetch existing videos for instant demo
    existing_videos = {}
    try:
        conn = connect(api_key=os.getenv("VIDEO_DB_API_KEY"))
        colls = conn.get_collections()
        if colls:
            videos = colls[0].get_videos()
            for v in reversed(videos):
                name_label = f"{v.name[:25]}..." if v.name else f"Untitled ({v.id[:8]})"
                existing_videos[name_label] = v.id
    except:
        pass

    source_mode = st.radio("Choose Source:", ["📂 Cloud Library", "🔗 New Upload"], label_visibility="collapsed")
    
    start_processing = False

    if source_mode == "📂 Cloud Library":
        if existing_videos:
            selected_name = st.selectbox("Select Video:", list(existing_videos.keys()))
            selected_id = existing_videos[selected_name]
            if st.button("🚀 Load Instant Demo"):
                with st.status("Fetching...", expanded=True) as status:
                    st.session_state.video_obj = cached_fetch(selected_id)
                    if st.session_state.video_obj:
                        status.update(label="Loaded!", state="complete", expanded=False)
                        st.session_state.analysis_done = False 
                        start_processing = True 
        else:
            st.warning("No videos found.")
            
    else:
        url = st.text_input("YouTube URL")
        if st.button("🚀 Upload & Process"):
            with st.status("Processing...", expanded=True) as status:
                st.write("1️⃣ Connecting to Cloud...")
                st.write("2️⃣ Indexing Video...")
                st.session_state.video_obj = cached_upload(url)
                if st.session_state.video_obj:
                    status.update(label="Ready!", state="complete", expanded=False)
                    st.session_state.analysis_done = False 
                    start_processing = True 

    st.divider()
    
    # 2. Director's Input
    st.subheader("2. Director's Vision")
    focus_topic = st.text_input("Focus Trailer On:", placeholder="e.g. Funny, Tech, Bloopers")
    if st.button("🔄 Re-Generate Analysis"):
        start_processing = True
        st.session_state.analysis_done = False

    st.divider()
    with st.expander("🛠️ Developer Tools"):
        dev_mode = st.toggle("Show Raw JSON")

# --- MAIN DASHBOARD ---
if st.session_state.video_obj:
    video = st.session_state.video_obj
    
    # We use Tabs for a cleaner layout
    tab1, tab2, tab3 = st.tabs(["🎬 The Studio", "📊 Viral Analytics", "🚀 Social Media"])
    
    # --- AUTO-START LOGIC ---
    if start_processing:
        st.session_state.analysis_done = False
    
    if not st.session_state.analysis_done:
        with st.status("🤖 AI Crew Working Parallel Tasks...", expanded=True) as status:
            try:
                transcript = video.get_transcript_text()
                
                # Using ThreadPool for parallel execution
                with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
                    future_map = {
                        executor.submit(analyze_virality, transcript): "virality",
                        executor.submit(generate_linkedin_post, transcript): "linkedin",
                        executor.submit(process_trailer_pipeline, video, transcript, focus_topic): "trailer"
                    }
                    
                    for future in concurrent.futures.as_completed(future_map):
                        task_name = future_map[future]
                        if task_name == "virality":
                            st.session_state.virality_data = future.result()
                        elif task_name == "linkedin":
                            st.session_state.linkedin_post = future.result()
                        elif task_name == "trailer":
                            st.session_state.highlights, st.session_state.stream_link = future.result()
                
                status.update(label="✅ Production Complete!", state="complete", expanded=False)
                st.session_state.analysis_done = True
                st.rerun() # Refresh to populate tabs
            except Exception as e:
                st.error(f"Error: {e}")

    # --- TAB 1: THE STUDIO (Video & Highlights) ---
    with tab1:
        col_video, col_details = st.columns([1.5, 1])
        
        with col_video:
            if "stream_link" in st.session_state and st.session_state.stream_link:
                st.subheader("🍿 AI-Generated Trailer")
                st.video(st.session_state.stream_link)
                st.caption(f"Stream Source: {st.session_state.stream_link}")
            else:
                st.info("Video will appear here after processing.")
                
        with col_details:
            st.subheader("💡 Director's Commentary")
            if "highlights" in st.session_state:
                for idx, clip in enumerate(st.session_state.highlights):
                    with st.container():
                        st.markdown(f"**Clip {idx+1} ({clip['start']}s - {clip['end']}s)**")
                        st.info(f"{clip['reason']}")
            else:
                st.markdown("*Waiting for analysis...*")

    # --- TAB 2: VIRAL ANALYTICS ---
    with tab2:
        if "virality_data" in st.session_state:
            data = st.session_state.virality_data
            
            # Score Cards
            m1, m2, m3 = st.columns(3)
            m1.metric("Viral Score", f"{data['score']}/100", delta="High Potential")
            
            with m2:
                st.markdown("**Core Hook**")
                if len(data['keywords']) > 0:
                    st.markdown(f":blue-background[{data['keywords'][0]}]")
            
            with m3:
                st.markdown("**Vibe**")
                if len(data['keywords']) > 1:
                    st.markdown(f":blue-background[{data['keywords'][1]}]")
            
            st.divider()
            
            # Simple Chart visualization
            st.subheader("📈 Audience Retention Forecast")
            chart_data = pd.DataFrame({
                'Time': ['Intro', 'Hook', 'Body', 'Climax', 'Outro'],
                'Engagement': [80, 90, 70, 95, 60] # Mock data for visuals
            })
            st.line_chart(chart_data.set_index('Time'))
            
        else:
            st.info("Analytics engine running...")

    # --- TAB 3: SOCIAL MEDIA ---
    with tab3:
        col_post, col_preview = st.columns(2)
        
        with col_post:
            st.subheader("📝 Draft Post")
            if "linkedin_post" in st.session_state:
                post_text = st.text_area("Edit your post:", value=st.session_state.linkedin_post, height=400)
                st.button("Copy to Clipboard")
            else:
                st.info("Copywriter agent is typing...")
        
        with col_preview:
            st.subheader("👀 Preview")
            st.markdown("""
            <div style="border:1px solid #ccc; padding:20px; border-radius:10px; background-color:white; color:black;">
                <div style="display:flex; align-items:center; margin-bottom:10px;">
                    <div style="width:40px; height:40px; background-color:#0077b5; border-radius:50%; margin-right:10px;"></div>
                    <div>
                        <div style="font-weight:bold;">You</div>
                        <div style="font-size:0.8em; color:gray;">AI Engineer • Just now</div>
                    </div>
                </div>
                <div style="font-size:0.9em;">
            """, unsafe_allow_html=True)
            
            if "linkedin_post" in st.session_state:
                st.markdown(st.session_state.linkedin_post.replace("\n", "<br>"), unsafe_allow_html=True)
            else:
                st.write("...")
                
            st.markdown("</div></div>", unsafe_allow_html=True)

else:
    # Empty State (Hero Section)
    st.info("👈 **Start Here:** Select a video from the Cloud Library or Upload a new one in the Sidebar.")