src/pages/05_Passport_Image_Generator.py

import streamlit as st
import requests
import os
import boto3
import uuid
from dotenv import load_dotenv
from PIL import Image
import logging
from botocore.exceptions import NoCredentialsError, ClientError
from typing import Optional, IO

# --- 1. CONFIGURATION ---
st.set_page_config(page_title="Passport Photo Generator", layout="wide", page_icon="📸")
load_dotenv()

# Load Config from .env
API_BASE_URL = os.getenv("API_BASE_URL", "https://tripgen.spun.global")
AWS_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY_ID")
AWS_SECRET_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
AWS_REGION = os.getenv("AWS_REGION", "ap-southeast-3")
S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")

# --- 2. S3 HELPER FUNCTIONS ---
def get_s3_client():
    if not AWS_ACCESS_KEY or not AWS_SECRET_KEY:
        st.error("❌ AWS Credentials missing in .env file.")
        return None
    return boto3.client(
        's3',
        aws_access_key_id=AWS_ACCESS_KEY,
        aws_secret_access_key=AWS_SECRET_KEY,
        region_name=AWS_REGION
    )

# Ensure your s3_client is initialized globally or passed in
s3_client = get_s3_client()

def upload_buffer_and_get_presigned_url(
    file_obj: IO, 
    bucket: str, 
    s3_key: str, 
    content_type: str = 'image/jpeg'
) -> Optional[str]:
    """
    Uploads an in-memory file object to S3 and returns its presigned URL.
    """
    try:
        logging.info(f"Uploading data to s3://{bucket}/{s3_key}...")
        
        # Reset file pointer to the beginning to ensure full upload
        file_obj.seek(0)
        
        # Use upload_fileobj for memory streams (BytesIO, SpooledTemporaryFile)
        s3_client.upload_fileobj(
            file_obj, 
            bucket, 
            s3_key,
            ExtraArgs={'ContentType': content_type} # Crucial for browser preview
        )
        
        logging.info("Generating presigned URL for access...")
        
        # Generate a presigned URL valid for 1 hour (3600 seconds)
        presigned_url = s3_client.generate_presigned_url(
            'get_object',
            Params={'Bucket': bucket, 'Key': s3_key},
            ExpiresIn=3600 
        )
        
        logging.info("Upload and signing successful.")
        return presigned_url

    except NoCredentialsError:
        logging.error("AWS credentials not found. Configure them to upload to S3.")
        return None
    except ClientError as e:
        logging.error(f"S3 Client Error: {e}")
        return None
    except Exception as e:
        logging.error(f"Unexpected error during S3 upload: {e}")
        return None
    
@st.cache_data(ttl=3600) 
def fetch_visa_metadata():
    """Fetches supported visa types and specs from the backend."""
    try:
        url = f"{API_BASE_URL}/photo-metadata"
        response = requests.get(url, timeout=5)
        if response.status_code == 200:
            return response.json()
        else:
            st.error(f"Failed to load visa options (Status: {response.status_code})")
            return {}
    except Exception as e:
        st.error(f"Could not connect to metadata endpoint: {e}")
        return {}

# --- 3. MAIN UI ---
st.title("📸 AI Passport & Visa Photo Generator")
st.markdown("Upload a raw photo, choose a background, and get a compliant passport photo.")

# Map visa to photo qualification
visa_options = fetch_visa_metadata()

# Sidebar for Settings
with st.sidebar:
    st.header("Settings")
    
    if visa_options:
        visa_name = st.selectbox(
            "Select Visa Type",
            options=list(visa_options.keys()),
            index=0
        )
        
        # Get details for selected visa
        selected_meta = visa_options[visa_name]
        visa = visa_name
        bg_color = selected_meta.get("background_color", "white")
        height = selected_meta.get("height", 4)
        width = selected_meta.get("width", 3)
        st.markdown(f"**Background Color:** {bg_color.capitalize()}")
    else:
        st.warning("No visa options available.")
        visa_name = None   

# Main Input Area
uploaded_file = st.file_uploader("Upload your photo (JPG/PNG)", type=["jpg", "jpeg", "png"])

# --- 4. LAYOUT: INPUT VS RESPONSE ---
col1, col2 = st.columns(2)

# Global variables to hold state
input_url = None

with col1:
    st.subheader("1. Input Image")
    if uploaded_file:
        # Display the uploaded image directly from memory
        st.image(uploaded_file, caption="Original Photo", use_container_width=True)
    else:
        st.info("Waiting for upload...")

with col2:
    st.subheader("2. Response Image")
    
    # Only show button if file is uploaded
    if uploaded_file:
        generate_btn = st.button("🚀 Process Photo", type="primary", use_container_width=True)
        
        if generate_btn:
            # A. UPLOAD TO S3
            with st.status("Processing...", expanded=True) as status:
                status.write("Uploading input to S3...")
                # Reset file pointer to beginning before upload
                uploaded_file.seek(0)
                input_url = upload_buffer_and_get_presigned_url(
                    file_obj=uploaded_file,
                    bucket=S3_BUCKET_NAME,
                    s3_key=f"inputs/{uploaded_file.name}",
                    content_type=uploaded_file.type
                )

                
                if input_url:
                    status.write(f"Upload complete. {input_url}")
                    status.write(f"✅ Uploaded Input")
                    
                    # B. CALL API
                    status.write("🤖 Sending to AI Engine...")
                    
                    api_url = f"{API_BASE_URL}/generate-visa-photo"
                    payload = {
                        "raw_photo": input_url,
                        "bg_color_name": bg_color,
                        "height": height,
                        "width": width
                    }
                    
                    try:
                        response = requests.post(api_url, json=payload, timeout=60)
                        
                        if response.status_code == 200:
                            # --- CHANGED: Handle JSON response with URL ---
                            try:
                                result_data = response.json()
                                result_url = result_data.get('url')
                                
                                if result_url:
                                    status.write("📥 Fetching image from S3 Result...")
                                    
                                    # Download the result image for display
                                    img_response = requests.get(result_url)
                                    
                                    if img_response.status_code == 200:
                                        image_bytes = img_response.content

                                        status.update(label="✅ Processing Complete!", state="complete", expanded=False)

                                        st.success("Generation Successful!")


                                        # Display Image
                                        st.image(image_bytes, caption=f"Processed ({visa})", use_container_width=True)

                                        # Display URL
                                        st.markdown(f"**Result URL:** [{result_url}]({result_url})")

                                        # Download Button
                                        st.download_button(
                                            label="⬇️ Download HD Image",
                                            data=image_bytes,
                                            file_name=f"passport_{visa.replace(' ', '_')}.jpg",
                                            mime="image/jpeg",
                                            type="primary"
                                        )
                                    else:
                                        status.update(label="❌ Download Failed", state="error")
                                        st.error(f"Failed to download image from S3 URL. Status: {img_response.status_code}")
                                        st.code(img_response.text[:500], language="xml")
                                else:
                                    # Backward compatibility: Check if response is raw bytes
                                    status.update(label="⚠️ Warning", state="complete")
                                    st.warning("API returned success but no URL. Checking for raw bytes...")
                            except ValueError:
                                # JSON decode failed, maybe it IS raw bytes (older API version)
                                image_bytes = response.content
                                status.update(label="✅ Complete (Raw Bytes)", state="complete")
                                st.image(image_bytes, caption="Processed", use_container_width=True)

                        else:
                            status.update(label="❌ API Error", state="error")
                            try:
                                err_msg = response.json().get("error", response.text)
                            except:
                                err_msg = response.text
                            st.error(f"API Failed: {err_msg}")
                            
                    except Exception as e:
                        status.update(label="❌ Connection Error", state="error")
                        st.error(f"Could not connect to API: {e}")
                else:
                    status.update(label="❌ Upload Failed", state="error")
    else:
        st.info("Result will appear here...")