src/streamlit_app.py

import streamlit as st
import base64
from PIL import Image
import io
import time
import requests
import os

# Page config
st.set_page_config(
    page_title="Sachdeva Creation Agent",
    page_icon="👗",
    layout="wide"
)

# Custom CSS
st.markdown("""
<style>
    .stApp {
        background-color: #f9fafb;
    }
    .upload-box {
        border: 2px dashed #d1d5db;
        border-radius: 1rem;
        padding: 2rem;
        text-align: center;
        background-color: #f9fafb;
        transition: all 0.3s;
    }
    .upload-box:hover {
        border-color: #d97706;
        background-color: #fffbeb;
    }
    .big-title {
        font-size: 2.5rem;
        font-weight: 800;
        color: #92400e;
        text-align: center;
        margin-bottom: 0.5rem;
    }
    .subtitle {
        font-size: 1.25rem;
        color: #d97706;
        text-align: center;
        margin-bottom: 2rem;
    }
    .step-title {
        font-size: 1.5rem;
        font-weight: 700;
        color: #1f2937;
        margin-bottom: 0.5rem;
    }
    .hindi-text {
        color: #d97706;
        font-size: 1.1rem;
        font-weight: 500;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
if 'step' not in st.session_state:
    st.session_state.step = 'SHIRT'
if 'suit_data' not in st.session_state:
    st.session_state.suit_data = {}
if 'generated_image' not in st.session_state:
    st.session_state.generated_image = None
if 'video_bytes' not in st.session_state:
    st.session_state.video_bytes = None

def generate_suit_image_hf(suit_data, hf_api_key):
    """Generate suit image using Hugging Face Stable Diffusion"""
    
    # Try multiple models in order of reliability
    models = [
        "black-forest-labs/FLUX.1-schnell",  # Fast, free model
        "stabilityai/stable-diffusion-2-1",   # Stable, reliable
        "runwayml/stable-diffusion-v1-5"      # Backup option
    ]
    
    # Create detailed prompt
    fabric_parts = []
    if 'shirt' in suit_data:
        fabric_parts.append("embroidered shirt")
    if 'dupatta' in suit_data:
        fabric_parts.append("flowing dupatta")
    if 'salwar' in suit_data:
        fabric_parts.append("Patiala salwar")
    
    prompt = f"""Professional fashion photography of a beautiful Punjabi model wearing a complete Patiala suit with {', '.join(fabric_parts)}.
    The suit features intricate embroidery and traditional designs.
    Sunny daylight setting, elegant pose, high-end boutique background.
    Photorealistic, high quality, detailed fabric textures, vibrant colors.
    Full body shot, fashion magazine style."""
    
    payload = {
        "inputs": prompt,
        "parameters": {
            "negative_prompt": "blurry, low quality, distorted, deformed, ugly, bad anatomy",
            "num_inference_steps": 30,
            "guidance_scale": 7.5
        }
    }
    
    # Try each model
    for model_name in models:
        try:
            API_URL = f"https://api-inference.huggingface.co/models/{model_name}"
            headers = {"Authorization": f"Bearer {hf_api_key}"}
            
            st.info(f"🎨 Trying model: {model_name.split('/')[-1]}...")
            
            response = requests.post(API_URL, headers=headers, json=payload, timeout=60)
            
            if response.status_code == 200:
                img_bytes = response.content
                img_b64 = base64.b64encode(img_bytes).decode()
                
                st.success(f"✅ Successfully generated with {model_name.split('/')[-1]}!")
                return {
                    'url': f'data:image/png;base64,{img_b64}',
                    'prompt': prompt
                }
            elif response.status_code == 503:
                st.warning(f"⏳ Model {model_name.split('/')[-1]} is loading... Trying next model...")
                continue
            elif response.status_code == 403:
                st.warning(f"⚠️ No access to {model_name.split('/')[-1]}... Trying next model...")
                continue
            else:
                st.warning(f"⚠️ Error {response.status_code} with {model_name.split('/')[-1]}... Trying next model...")
                continue
                
        except Exception as e:
            st.warning(f"⚠️ Failed with {model_name.split('/')[-1]}: {str(e)}")
            continue
    
    # If all models fail
    st.error("""
    ❌ **All models failed to generate image**
    
    **Possible solutions:**
    
    1. **Verify your HF Token:**
       - Go to: https://huggingface.co/settings/tokens
       - Make sure token has "Read" permission
       - Copy the FULL token (starts with hf_)
       - Replace in Space Settings → Secrets → HF_TOKEN
    
    2. **Check your account:**
       - Make sure you're logged into HuggingFace
       - Some models may need you to accept their license
       - Visit: https://huggingface.co/black-forest-labs/FLUX.1-schnell
       - Click "Agree and access repository"
    
    3. **Wait and retry:**
       - Models may be loading (takes 20-30 seconds)
       - Click the button again after waiting
    
    4. **Restart your Space:**
       - Go to Space Settings
       - Click "Factory Reboot"
    """)
    return None

# VIDEO GENERATION - COMMENTED OUT FOR NOW
# def generate_video_hf(image_b64, prompt, hf_api_key):
#     """Generate video using Hugging Face Image-to-Video model"""
#     
#     API_URL = "https://api-inference.huggingface.co/models/ali-vilab/i2vgen-xl"
#     
#     headers = {"Authorization": f"Bearer {hf_api_key}"}
#     
#     # Decode base64
#     img_bytes = base64.b64decode(image_b64)
#     
#     try:
#         st.info("🎬 Creating your promotional video...")
#         
#         response = requests.post(API_URL, headers=headers, data=img_bytes)
#         
#         if response.status_code == 200:
#             return response.content
#         elif response.status_code == 503:
#             st.warning("⏳ Model is loading... Please wait and try again.")
#             return None
#         else:
#             st.error(f"❌ Error: {response.status_code} - {response.text}")
#             return None
#             
#     except Exception as e:
#         st.error(f"❌ Error generating video: {str(e)}")
#         return None

# Header
st.markdown('<div class="big-title">🌟 Sachdeva Creation 🌟</div>', unsafe_allow_html=True)
st.markdown('<div class="subtitle">AI-Powered Fashion Lookbook Generator / एआई-संचालित फैशन लुकबुक जनरेटर</div>', unsafe_allow_html=True)

# Sidebar for API keys
with st.sidebar:
    st.header("⚙️ Configuration")
    
    # Check if running on Hugging Face Spaces
    on_huggingface = os.getenv("SPACE_ID") is not None
    
    if on_huggingface:
        st.info("🤗 Running on Hugging Face Spaces")
        hf_api_key = os.getenv("HF_TOKEN", "")
        
        if not hf_api_key:
            st.error("⚠️ HF_TOKEN not found in secrets!")
            st.info("Add your Hugging Face token in Space Settings → Repository secrets")
        else:
            st.success("✅ HF_TOKEN found!")
            # Show first/last 4 chars for verification
            token_preview = f"{hf_api_key[:4]}...{hf_api_key[-4:]}"
            st.code(f"Token: {token_preview}", language="text")
    else:
        st.info("💻 Local Development Mode")
        hf_api_key = st.text_input("Hugging Face API Key", type="password", 
                                   help="Get your free token from huggingface.co/settings/tokens")
    
    st.markdown("---")
    st.markdown("### 📖 How to use:")
    st.markdown("""
    1. Upload shirt fabric photo
    2. Upload dupatta fabric photo
    3. Optionally upload salwar fabric
    4. Click **Generate Image** button
    5. Review the generated design
    6. Download your image
    """)
    # 7. Click **Generate Video** for promo (Coming Soon!)
    # 8. Download both image and video
    
    st.markdown("---")
    st.markdown("### 🔑 Get HuggingFace Token:")
    st.info("""
    1. Go to huggingface.co
    2. Sign up/Login
    3. Go to Settings → Access Tokens
    4. Create new token (Read access)
    5. Copy and add to Space secrets
    """)

# Main content based on step
if st.session_state.step == 'SHIRT':
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.markdown('<div class="step-title">Upload Shirt</div>', unsafe_allow_html=True)
        st.markdown('<div class="hindi-text">शर्ट अपलोड करें</div>', unsafe_allow_html=True)
        st.markdown("---")
        
        st.info("Please upload a clear photo of the shirt fabric focusing on design and embroidery.\n\nकृपया डिजाइन और कढ़ाई पर ध्यान केंद्रित करते हुए शर्ट के कपड़े की एक स्पष्ट फोटो अपलोड करें।")
        
        uploaded_file = st.file_uploader("Choose shirt image", type=['png', 'jpg', 'jpeg'], key='shirt_upload')
        
        if uploaded_file:
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Shirt Fabric", use_container_width=True)
            
            buffered = io.BytesIO()
            image.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode()
            st.session_state.suit_data['shirt'] = f'data:image/png;base64,{img_str}'
            
            if st.button("Next / अगला", type="primary", use_container_width=True):
                st.session_state.step = 'DUPATTA'
                st.rerun()

elif st.session_state.step == 'DUPATTA':
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.markdown('<div class="step-title">Upload Dupatta</div>', unsafe_allow_html=True)
        st.markdown('<div class="hindi-text">दुपट्टा अपलोड करें</div>', unsafe_allow_html=True)
        st.markdown("---")
        
        st.info("Please upload a clear photo of the dupatta fabric focusing on design and embroidery.\n\nकृपया डिजाइन और कढ़ाई पर ध्यान केंद्रित करते हुए दुपट्टा के कपड़े की एक स्पष्ट फोटो अपलोड करें।")
        
        uploaded_file = st.file_uploader("Choose dupatta image", type=['png', 'jpg', 'jpeg'], key='dupatta_upload')
        
        if uploaded_file:
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Dupatta Fabric", use_container_width=True)
            
            buffered = io.BytesIO()
            image.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode()
            st.session_state.suit_data['dupatta'] = f'data:image/png;base64,{img_str}'
            
            if st.button("Next / अगला", type="primary", use_container_width=True):
                st.session_state.step = 'SALWAR'
                st.rerun()

elif st.session_state.step == 'SALWAR':
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.markdown('<div class="step-title">Upload Salwar (Optional)</div>', unsafe_allow_html=True)
        st.markdown('<div class="hindi-text">सलवार अपलोड करें (वैकल्पिक)</div>', unsafe_allow_html=True)
        st.markdown("---")
        
        st.info("Please upload a clear photo of the salwar fabric focusing on design and embroidery.\n\nकृपया डिजाइन और कढ़ाई पर ध्यान केंद्रित करते हुए सलवार के कपड़े की एक स्पष्ट फोटो अपलोड करें।")
        
        uploaded_file = st.file_uploader("Choose salwar image (optional)", type=['png', 'jpg', 'jpeg'], key='salwar_upload')
        
        if uploaded_file:
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Salwar Fabric", use_container_width=True)
            
            buffered = io.BytesIO()
            image.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode()
            st.session_state.suit_data['salwar'] = f'data:image/png;base64,{img_str}'
        
        st.markdown("---")
        if st.button("Continue to Generation / जनरेशन के लिए जारी रखें", type="primary", use_container_width=True):
            st.session_state.step = 'REVIEW'
            st.rerun()

elif st.session_state.step == 'REVIEW':
    st.markdown('<div class="big-title">Generate Your Design</div>', unsafe_allow_html=True)
    st.markdown('<div class="subtitle">अपना डिज़ाइन बनाएं</div>', unsafe_allow_html=True)
    
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        # Show uploaded fabrics
        st.markdown("### Uploaded Fabrics:")
        fabric_cols = st.columns(len(st.session_state.suit_data))
        for idx, (key, img_data) in enumerate(st.session_state.suit_data.items()):
            with fabric_cols[idx]:
                img_b64 = img_data.split(',')[1]
                img_bytes = base64.b64decode(img_b64)
                image = Image.open(io.BytesIO(img_bytes))
                st.image(image, caption=key.title(), use_container_width=True)
        
        st.markdown("---")
        
        # Show generated image if exists
        if st.session_state.generated_image:
            st.markdown("### Generated Design:")
            img_b64 = st.session_state.generated_image['url'].split(',')[1]
            img_bytes = base64.b64decode(img_b64)
            image = Image.open(io.BytesIO(img_bytes))
            st.image(image, use_container_width=True)
            
            # Download image button
            st.download_button(
                "📥 Download Image / इमेज डाउनलोड करें",
                data=base64.b64decode(img_b64),
                file_name=f"sachdeva-creation-{int(time.time())}.png",
                mime="image/png",
                use_container_width=True
            )
            
            st.markdown("---")
            
            # VIDEO SECTION - COMMENTED OUT FOR NOW
            # if st.session_state.video_bytes:
            #     st.markdown("### Promotional Video:")
            #     st.video(st.session_state.video_bytes)
            #     
            #     st.download_button(
            #         "📥 Download Video / वीडियो डाउनलोड करें",
            #         data=st.session_state.video_bytes,
            #         file_name=f"sachdeva-creation-promo-{int(time.time())}.mp4",
            #         mime="video/mp4",
            #         use_container_width=True
            #     )
            # else:
            #     if st.button("🎬 Generate Promotional Video / प्रोमो वीडियो बनाएं", type="primary", use_container_width=True):
            #         if not hf_api_key:
            #             st.error("❌ HuggingFace API key is required!")
            #         else:
            #             with st.spinner("Creating video... This may take 1-2 minutes..."):
            #                 video_bytes = generate_video_hf(
            #                     img_b64, 
            #                     st.session_state.generated_image['prompt'],
            #                     hf_api_key
            #                 )
            #                 if video_bytes:
            #                     st.session_state.video_bytes = video_bytes
            #                     st.success("✅ Video generated successfully!")
            #                     st.rerun()
            
            st.info("🎬 Video generation feature coming soon!")
            
            st.markdown("---")
            
            # Action buttons
            col_a, col_b = st.columns(2)
            with col_a:
                if st.button("🔄 Regenerate Image / फिर से बनाएं", use_container_width=True):
                    st.session_state.generated_image = None
                    st.session_state.video_bytes = None
                    st.rerun()
            with col_b:
                if st.button("🆕 Start New Suit / नया सूट शुरू करें", use_container_width=True):
                    st.session_state.step = 'SHIRT'
                    st.session_state.suit_data = {}
                    st.session_state.generated_image = None
                    st.session_state.video_bytes = None
                    st.rerun()
        
        else:
            # Generate image button
            if st.button("✨ Generate Image / इमेज बनाएं", type="primary", use_container_width=True):
                if not hf_api_key:
                    st.error("❌ HuggingFace API key is required!")
                    st.info("Please add your HuggingFace token in the sidebar or Space settings.")
                else:
                    with st.spinner("Generating your design... Please wait..."):
                        generated_img = generate_suit_image_hf(st.session_state.suit_data, hf_api_key)
                        if generated_img:
                            st.session_state.generated_image = generated_img
                            st.success("✅ Image generated successfully!")
                            st.rerun()
            
            st.markdown("---")
            if st.button("← Go Back / वापस जाएं", use_container_width=True):
                st.session_state.step = 'SALWAR'
                st.rerun()

# Footer
st.markdown("---")
st.markdown("""
<div style='text-align: center; color: #6b7280; padding: 2rem;'>
    <p>Powered by Hugging Face AI 🤗</p>
    <p>© 2024 Sachdeva Creation - AI Fashion Innovation</p>
</div>
""", unsafe_allow_html=True)