import streamlit as st
import json
from huggingface_hub import InferenceClient
import time

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

# Initialize session state
if 'api_token' not in st.session_state:
    st.session_state.api_token = ""

# Sidebar - API Token
with st.sidebar:
    st.header("🔑 Configuration")
    api_token = st.text_input(
        "HuggingFace API Token",
        type="password",
        value=st.session_state.api_token,
        help="Get your token from https://huggingface.co/settings/tokens"
    )
    if api_token:
        st.session_state.api_token = api_token
    
    st.divider()
    st.markdown("### 📋 JSON Structure")
    st.code("""{
  "product": "laptop",
  "style": "cinematic",
  "mood": "premium",
  "duration": "short",
  "camera": "rotating",
  "lighting": "studio"
}""", language="json")

# Main title
st.title("🎬 AI Video Advertisement Generator")
st.markdown("Generate professional video ads from JSON specifications using state-of-the-art AI")

# 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
        )
        
        # Advanced settings
        with st.expander("⚙️ Advanced Settings"):
            model_choice = st.selectbox(
                "Model",
                ["tencent/HunyuanVideo", "THUDM/CogVideoX-5b", "genmo/mochi-1-preview"],
                help="HunyuanVideo: Best quality | CogVideoX: Longer videos | Mochi: Fastest"
            )
            
            resolution = st.selectbox("Resolution", ["720p", "1080p"], index=0)
            fps = st.slider("FPS", 24, 30, 24)
        
        # 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:
        if not st.session_state.api_token:
            st.error("⚠️ Please enter your HuggingFace API token in the sidebar")
        else:
            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"🎨 Generating with {model_choice}...")
                
                # Initialize client
                client = InferenceClient(token=st.session_state.api_token)
                
                # Progress simulation
                progress_bar = st.progress(0)
                for i in range(100):
                    time.sleep(0.3)
                    progress_bar.progress(i + 1)
                
                # Generate video
                with st.spinner("🎬 Creating your video ad... (this may take 30-60 seconds)"):
                    video_bytes = client.text_to_video(
                        prompt=prompt,
                        model=model_choice
                    )
                
                progress_bar.empty()
                status_placeholder.success("✅ Video generated successfully!")
                
                # Display video
                with col2:
                    video_placeholder.video(video_bytes)
                    
                    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 json.JSONDecodeError:
                st.error("❌ Invalid JSON format. Please check your input.")
            except Exception as e:
                st.error(f"❌ Error: {str(e)}")
                status_placeholder.empty()

with tab2:
    st.markdown("""
    ## 🎯 How to Use
    
    ### 1️⃣ Get API Token
    - Visit [HuggingFace Tokens](https://huggingface.co/settings/tokens)
    - Create a new token with "read" permissions
    - Paste it in the sidebar
    
    ### 2️⃣ Configure Your Ad
    Define your video ad using JSON with these properties:
    
    - **product**: What you're advertising
    - **brand_style**: Your brand aesthetic (modern, minimal, bold)
    - **visual_style**: Video style (cinematic, dynamic, elegant)
    - **camera_movement**: How camera moves (rotation, zoom, pan)
    - **lighting**: Lighting setup (dramatic, soft, studio)
    - **background**: Background style (gradient, solid, abstract)
    - **mood**: Overall feeling (premium, energetic, calm)
    - **key_features**: List of product highlights
    - **duration**: Video length preference
    
    ### 3️⃣ Generate
    Click "Generate Video Ad" and wait 30-60 seconds
    
    ### 🎨 Model Comparison
    
    - **HunyuanVideo**: Best quality, photorealistic, 5-6 seconds
    - **CogVideoX**: Good quality, longer duration, 10+ seconds
    - **Mochi**: Fastest generation, 3-5 seconds, lightweight
    
    ### 💡 Tips for Best Results
    
    - Be specific about visual style and camera movement
    - Include 3-5 key features maximum
    - Use cinematic/commercial terminology
    - Describe lighting and mood clearly
    - Keep product name concise
    
    ### ⚡ Example JSON Templates
    
    **Tech Product:**
    ```json
    {
      "product": "smartphone",
      "visual_style": "Apple-style commercial",
      "camera_movement": "slow orbit around device",
      "lighting": "soft gradient backlight",
      "mood": "minimalist premium"
    }
    ```
    
    **Fashion/Lifestyle:**
    ```json
    {
      "product": "luxury watch",
      "visual_style": "high-fashion editorial",
      "camera_movement": "close-up macro details",
      "lighting": "dramatic side lighting",
      "mood": "elegant timeless"
    }
    ```
    """)

# Footer
st.divider()
st.markdown("""
<div style='text-align: center; color: #666;'>
    <p>Powered by HuggingFace Inference API | Free GPU-accelerated video generation</p>
</div>
""", unsafe_allow_html=True)