import streamlit as st
import os
from vton import virtual_try_on
import requests
from PIL import Image
from io import BytesIO

# Set page config for a wider layout and custom title
st.set_page_config(
    page_title="Virtual Try-On App",
    page_icon="👔",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better styling with Renesistech branding
st.markdown("""
<style>
    /* Global Styles */
    .stApp {
        max-width: 1400px;
        margin: 0 auto;
        background-color: #ffffff;
        font-family: 'Inter', sans-serif;
    }
    
    /* Header Styles */
    .renesistech-header {
        background: linear-gradient(to right, #f8f9fa, #ffffff);
        padding: 2rem;
        margin-bottom: 2rem;
        box-shadow: 0 4px 6px rgba(0,0,0,0.05);
        text-align: center;
        border-bottom: 1px solid #eaeaea;
    }
    
    /* Main Content Styles */
    .stButton > button {
        background-color: #007bff;
        color: white;
        padding: 0.75rem 2rem;
        border-radius: 8px;
        border: none;
        font-weight: 600;
        transition: all 0.3s ease;
    }
    
    .stButton > button:hover {
        background-color: #0056b3;
        transform: translateY(-2px);
    }
    
    /* Upload Box Styles */
    .upload-box {
        border: 2px dashed #e0e0e0;
        border-radius: 12px;
        padding: 2rem;
        text-align: center;
        margin: 1rem 0;
        background-color: #fafafa;
        transition: all 0.3s ease;
    }
    
    .upload-box:hover {
        border-color: #007bff;
        background-color: #f8f9fa;
    }
    
    /* Output Box Styles */
    .output-box {
        border: 1px solid #eaeaea;
        border-radius: 12px;
        padding: 2rem;
        margin: 2rem 0;
        background-color: #ffffff;
        box-shadow: 0 2px 4px rgba(0,0,0,0.05);
    }
    
    /* Footer Styles */
    .renesistech-footer {
        background: linear-gradient(to right, #f8f9fa, #ffffff);
        padding: 2rem;
        text-align: center;
        border-top: 1px solid #eaeaea;
        margin-top: 3rem;
    }
    
    /* Typography Improvements */
    h1, h2, h3, .stMarkdown p {
        color: #2c3e50;
        letter-spacing: -0.5px;
    }
    
    .stTextArea textarea {
        border-radius: 8px;
        border: 1px solid #eaeaea;
        padding: 1rem;
        font-size: 1rem;
    }
    
    .stTextArea textarea:focus {
        border-color: #007bff;
        box-shadow: 0 0 0 2px rgba(0,123,255,0.25);
    }
    
    /* Image Display Improvements */
    .stImage img {
        border-radius: 12px;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    }
</style>
""", unsafe_allow_html=True)

# Add Renesistech header
st.markdown("""
<div class='renesistech-header'>
    <h1 style='color: #333; margin: 0;'>Renesistech</h1>
    <p style='color: #666; margin: 10px 0 0 0;'>Virtual Try-On Solution</p>
</div>
""", unsafe_allow_html=True)

# Title and description
st.title("✨ Virtual Try-On Application")
st.markdown("""Upload a garment image and a person's photo to see how the garment would look on them!""")

# Create two columns for input images
col1, col2 = st.columns(2)

# Garment image upload
with col1:
    st.subheader("Garment Image")
    garment_file = st.file_uploader("Upload garment image", type=["jpg", "jpeg", "png", "webp"], key="garment_upload")
    garment_url = None
    if garment_file:
        st.image(garment_file, caption="Uploaded Garment", use_container_width=True)
        # Save uploaded file to temporary file
        garment_img = Image.open(garment_file)
        garment_buffer = BytesIO()
        garment_img.save(garment_buffer, format="PNG")
        garment_url = "temp_garment.png"
        with open(garment_url, "wb") as f:
            f.write(garment_buffer.getvalue())

# Person image upload
with col2:
    st.subheader("Person Image")
    person_file = st.file_uploader("Upload person image", type=["jpg", "jpeg", "png", "webp"], key="person_upload")
    person_url = None
    if person_file:
        st.image(person_file, caption="Uploaded Person", use_container_width=True)
        # Save uploaded file to temporary file
        person_img = Image.open(person_file)
        person_buffer = BytesIO()
        person_img.save(person_buffer, format="PNG")
        person_url = "temp_person.png"
        with open(person_url, "wb") as f:
            f.write(person_buffer.getvalue())

# Garment description input
st.subheader("Garment Description")
garment_desc = st.text_area(
    "Describe the garment (e.g., color, style, type)",
    height=100,
    placeholder="Example: A cute pink top with floral pattern"
)

# Check for API token
if "REPLICATE_API_TOKEN" not in os.environ:
    st.warning(
        "⚠️ REPLICATE_API_TOKEN is not set. Please set it using: "
        "export REPLICATE_API_TOKEN='your_token_here'"
    )

# Process button
if st.button("Generate Try-On", type="primary"):
    if not (garment_url and person_url and garment_desc):
        st.error("Please provide all required inputs!")
    else:
        try:
            with st.spinner("🔄 Processing virtual try-on..."):
                # Read the image data from the temporary files
                with open(garment_url, 'rb') as f:
                    garment_data = f.read()
                with open(person_url, 'rb') as f:
                    person_data = f.read()
                    
                output_path = virtual_try_on(garment_data, person_data, garment_desc)
                
                # Display result in a styled box
                st.markdown("<div class='output-box'>", unsafe_allow_html=True)
                st.subheader("🎉 Try-On Result")
                st.image(output_path, caption="Virtual Try-On Result", use_container_width=True)
                st.markdown("</div>", unsafe_allow_html=True)
                
                # Cleanup temporary files
                if os.path.exists("temp_garment.png"):
                    os.remove("temp_garment.png")
                if os.path.exists("temp_person.png"):
                    os.remove("temp_person.png")
                    
        except Exception as e:
            st.error(f"Error occurred: {str(e)}")

# Renesistech Footer
st.markdown("""
<div class='renesistech-footer'>
    <p style='color: #333; font-weight: bold; margin: 0;'>© 2025 Renesistech</p>
    <p style='color: #666; margin: 5px 0 0 0;'>Innovating Fashion Technology</p>
</div>
""", unsafe_allow_html=True)