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

st.set_page_config(page_title="AI Image Detector", page_icon="🔍")

st.title("AI Image Detector")
st.write("Upload an image to check if it's AI-generated")


api_key = "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"
st.session_state.api_key = api_key

def process_image(image_bytes, api_key):
    header_auth = f"Bearer {api_key}"
    invoke_url = "https://ai.api.nvidia.com/v1/cv/hive/ai-generated-image-detection"
    
    # Convert image bytes to base64
    image_b64 = base64.b64encode(image_bytes).decode()
    
    payload = {
        "input": [f"data:image/png;base64,{image_b64}"]
    }
    headers = {
        "Content-Type": "application/json",
        "Authorization": header_auth,
        "Accept": "application/json",
    }
    
    try:
        response = requests.post(invoke_url, headers=headers, json=payload)
        response.raise_for_status()
        result = response.json()
        
        # Check if response contains the expected structure
        if 'data' in result and len(result['data']) > 0:
            first_result = result['data'][0]
            if 'is_ai_generated' in first_result:
                return {
                    'confidence': first_result['is_ai_generated'],
                    'sources': first_result.get('possible_sources', {}),
                    'status': first_result.get('status', 'UNKNOWN')
                }
        
        st.error("Unexpected response format from API")
        return None
            
    except requests.exceptions.RequestException as e:
        st.error(f"Error processing image: {str(e)}")
        return None

# File uploader
uploaded_file = st.file_uploader("Choose an image...", type=['png', 'jpg', 'jpeg'])

if uploaded_file is not None and api_key:
    # Display the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_container_width=True)
    
    # Convert image to bytes
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format=image.format)
    img_byte_arr = img_byte_arr.getvalue()
    
    # Process the image
    with st.spinner("Analyzing image..."):
        result = process_image(img_byte_arr, api_key)
        
        if result and result['status'] == 'SUCCESS':
            confidence = result['confidence']
            sources = result['sources']
            
            st.write("---")
            st.write("### Result")
            
            # Determine if image is AI-generated (using 50% threshold)
            is_ai_generated = "Yes" if confidence >= 0.5 else "No"
            
            # Display result with appropriate styling
            if is_ai_generated == "Yes":
                st.error(f"Is this image AI-generated? **{is_ai_generated}**")
                
                # Show top 3 possible sources if AI-generated
                if sources:
                    st.write("Top possible AI models used:")
                    sorted_sources = sorted(sources.items(), key=lambda x: x[1], reverse=True)[:3]
                    for source, prob in sorted_sources:
                        if prob > 0.01:  # Only show sources with >1% probability
                            st.write(f"- {source}: {prob:.1%}")
            else:
                st.success(f"Is this image AI-generated? **{is_ai_generated}**")
            
            # Show confidence score in smaller text
            st.caption(f"Confidence score: {confidence:.2%}")

elif not api_key and uploaded_file is not None:
    st.warning("Please enter your NVIDIA API key first")

# Add footer with instructions
st.markdown("---")
st.markdown("""
---
### How to use:

1. Upload an image (PNG, JPG, or JPEG)
2. Wait for the analysis result
3. Get a ** Yes/No ** answer based on whether the image is AI-generated

""")