import streamlit as st
import json
from gradio_client import Client
import time

st.set_page_config(page_title="AI Video Ad Generator", page_icon="🎬", layout="wide")

# Main title
st.title("🎬 AI Video Advertisement Generator")
st.markdown("Generate professional video ads from JSON using free community GPU models")

# Tabs
tab1, tab2 = st.tabs(["🎥 Generate Video", "📖 Guide"])

with tab1:
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.subheader("Input Configuration")
        
        # JSON Input
        json_input = st.text_area(
            "Ad Specification (JSON)",
            value="""{
  "product": "premium laptop",
  "brand_style": "modern tech",
  "visual_style": "cinematic commercial",
  "camera_movement": "smooth 360 rotation",
  "lighting": "dramatic studio backlight",
  "background": "gradient dark to light",
  "mood": "premium luxury",
  "key_features": ["ultra-thin", "metallic finish", "glowing edges"],
  "duration": "5 seconds"
}""",
            height=300
        )
        
        # Model selection
        with st.expander("⚙️ Model Settings"):
            model_space = st.selectbox(
                "Model Space",
                [
                    "THUDM/CogVideoX-5B",
                    "THUDM/CogVideoX-2B"
                ],
                help="CogVideoX models via community Spaces"
            )
        
        # Generate button
        generate_btn = st.button("🎬 Generate Video Ad", type="primary", use_container_width=True)
    
    with col2:
        st.subheader("Generated Video")
        video_placeholder = st.empty()
        status_placeholder = st.empty()

    # Generation logic
    if generate_btn:
        try:
            # Parse JSON
            ad_config = json.loads(json_input)
            
            # Build cinematic prompt from JSON
            prompt = f"""Professional commercial advertisement video showcasing {ad_config.get('product', 'product')}, 
{ad_config.get('visual_style', 'cinematic')} style, {ad_config.get('camera_movement', 'smooth camera movement')}, 
{ad_config.get('lighting', 'professional lighting')}, {ad_config.get('background', 'modern background')}, 
{ad_config.get('mood', 'premium')} aesthetic, product-focused hero shot, 
{', '.join(ad_config.get('key_features', []))}, commercial quality, 4K resolution, 
professional advertising photography, luxury brand style, high-end production value"""
            
            status_placeholder.info(f"🎨 Connecting to {model_space}...")
            
            with st.spinner("🎬 Generating video... This takes 60-120 seconds"):
                try:
                    # Try to connect to community Space
                    client = Client(model_space)
                    
                    # Generate video
                    result = client.predict(
                        prompt=prompt,
                        api_name="/predict"
                    )
                    
                    status_placeholder.success("✅ Video generated successfully!")
                    
                    # Display video
                    if isinstance(result, str):
                        video_placeholder.video(result)
                        
                        with open(result, "rb") as f:
                            video_bytes = f.read()
                        
                        st.download_button(
                            label="⬇️ Download Video",
                            data=video_bytes,
                            file_name=f"ad_{ad_config.get('product', 'video').replace(' ', '_')}.mp4",
                            mime="video/mp4",
                            use_container_width=True
                        )
                    
                    # Show generated prompt
                    with st.expander("📝 Generated Prompt"):
                        st.text(prompt)
                        
                except Exception as api_error:
                    status_placeholder.error(f"""
                    ⚠️ **Free GPU queue is full or Space unavailable**
                    
                    **Alternative Options:**
                    1. Try again in a few minutes (queue clears)
                    2. Use a paid service like Replicate or Fal.ai
                    3. Deploy your own GPU Space (requires payment)
                    
                    Technical error: {str(api_error)}
                    """)
                    
        except json.JSONDecodeError:
            st.error("❌ Invalid JSON format. Please check your input.")
        except Exception as e:
            st.error(f"❌ Error: {str(e)}")

with tab2:
    st.markdown("""
    ## 🎯 How It Works
    
    This app calls **community-hosted HuggingFace Spaces** with free GPU access.
    
    ### ⚠️ Important Limitations
    
    - **Queue wait times**: Popular Spaces have queues (2-5 minutes)
    - **Availability**: Spaces may go offline
    - **Generation time**: 60-120 seconds per video
    - **Video length**: 5-10 seconds maximum
    
    ### 🔄 If Generation Fails
    
    The free models have limited capacity. Options:
    
    1. **Wait and retry** - Queues clear quickly
    2. **Use paid APIs**:
       - Replicate.com ($0.01-0.05/video)
       - Fal.ai ($0.05-0.15/video)
       - Runway ML (paid plans)
    3. **Self-host** with your own GPU Space
    
    ### 💡 JSON Configuration
    
    ```json
    {
      "product": "What you're advertising",
      "visual_style": "cinematic/dynamic/elegant",
      "camera_movement": "rotation/zoom/pan",
      "lighting": "dramatic/soft/studio",
      "background": "gradient/solid/abstract",
      "mood": "premium/energetic/calm",
      "key_features": ["feature1", "feature2"]
    }
    ```
    
    ### 🚀 For Production Use
    
    Free models are **NOT suitable for production**. Consider:
    
    - **Replicate API** - Pay per generation, reliable
    - **Runway ML** - Professional video generation
    - **Luma Dream Machine** - High quality, paid
    - **Self-hosted GPU** - Full control, requires infrastructure
    """)

# Footer
st.divider()
st.markdown("""
<div style='text-align: center; color: #666;'>
    <p>⚠️ Using free community Spaces - expect wait times and availability issues</p>
</div>
""", unsafe_allow_html=True)