# app.py
import streamlit as st
from PIL import Image
import io
from face_swap import swap_faces

# --- Page config ---
st.set_page_config(page_title="🎭 Face Swap Demo", layout="centered")

# --- Hero section ---
st.markdown(
    """
    <style>
    .title {
        font-size: 2.5rem;
        font-weight: bold;
        text-align: center;
        margin-bottom: 0.5rem;
    }
    .subtitle {
        text-align: center;
        font-size: 1.2rem;
        color: #555;
        margin-bottom: 2rem;
    }
    </style>
    <div class="title">🎭 Face Swap Demo</div>
    <div class="subtitle">Upload two images and see the magic of AI face-swapping</div>
    """,
    unsafe_allow_html=True,
)

# --- Upload section in cards ---
col1, col2 = st.columns(2)
with col1:
    st.markdown("### 🧑 Source Face")
    src_file = st.file_uploader(
        "Upload Source Image", type=["jpg", "jpeg", "png"], key="src"
    )
with col2:
    st.markdown("### 🎯 Target Image")
    tgt_file = st.file_uploader(
        "Upload Target Image", type=["jpg", "jpeg", "png"], key="tgt"
    )

# --- Options ---
with st.sidebar:
    st.header("⚙️ Options")
    resize_max = st.number_input(
        "Resize images to max dimension (px)", value=800, min_value=200, max_value=2000
    )
    download_name = st.text_input("Download filename", value="face_swap_result.png")

def load_image(file):
    img = Image.open(file).convert("RGB")
    # resize if large
    w, h = img.size
    max_dim = max(w, h)
    if max_dim > resize_max:
        scale = resize_max / max_dim
        img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
    return img

# --- Run button ---
if st.button("🚀 Run Face Swap", use_container_width=True):
    if not src_file or not tgt_file:
        st.error("⚠️ Please upload both source and target images.")
    else:
        src_img = load_image(src_file)
        tgt_img = load_image(tgt_file)

        st.markdown("### 🔍 Preview Images")
        st.image([src_img, tgt_img], caption=["Source", "Target"], width=300)

        with st.spinner("Running face-swap... please wait ⏳"):
            try:
                result_pil = swap_faces(src_img, tgt_img)

                st.success("✨ Done! Check the result below")

                # Show side-by-side comparison
                col_a, col_b = st.columns(2)
                with col_a:
                    st.markdown("**Before**")
                    st.image(tgt_img, use_column_width=True)
                with col_b:
                    st.markdown("**After (Swapped)**")
                    st.image(result_pil, use_column_width=True)

                # Download button
                buf = io.BytesIO()
                result_pil.save(buf, format="PNG")
                st.download_button(
                    "⬇️ Download result",
                    buf.getvalue(),
                    file_name=download_name,
                    mime="image/png",
                    use_container_width=True,
                )

            except Exception as e:
                st.error(f"❌ Face-swap failed: {e}")
                st.exception(e)

# --- Footer ---
st.markdown(
    """
    <hr>
    <div style="text-align:center; color: gray; font-size: 0.9rem;">
        Built with ❤️ using Streamlit and InsightFace
    </div>
    """,
    unsafe_allow_html=True,
)