"""
AI Image Detector
Minimal. Yes/No + Confidence.
Uses Smogy detector (balanced precision/recall, fewer false positives)
"""

import os
os.environ['HF_HOME'] = '/tmp/hf_cache'
os.environ['TRANSFORMERS_CACHE'] = '/tmp/hf_cache'

import streamlit as st
from PIL import Image
import io

st.set_page_config(
    page_title="AI?",
    page_icon="🔍",
    layout="centered"
)

st.markdown("""
<style>
    #MainMenu, footer, header {visibility: hidden;}
    .block-container {padding-top: 2rem; padding-bottom: 0; max-width: 500px;}
</style>
""", unsafe_allow_html=True)


@st.cache_resource
def load_detector():
    """Load Smogy AI detector - better balanced, fewer false positives"""
    from transformers import pipeline
    return pipeline(
        "image-classification",
        model="Smogy/SMOGY-Ai-images-detector",
        device=-1
    )


def predict(detector, image: Image.Image) -> tuple:
    """Returns (is_ai, ai_confidence, human_confidence)"""
    img = image.convert('RGB')
    
    # Resize if too large
    max_dim = 1024
    if max(img.size) > max_dim:
        ratio = max_dim / max(img.size)
        img = img.resize((int(img.size[0] * ratio), int(img.size[1] * ratio)), Image.LANCZOS)
    
    results = detector(img)
    
    ai_score = 0.0
    human_score = 0.0
    
    for r in results:
        label = r['label'].lower()
        if 'ai' in label or 'artificial' in label or 'fake' in label:
            ai_score = max(ai_score, r['score'])
        elif 'human' in label or 'real' in label:
            human_score = max(human_score, r['score'])
    
    # Conservative: only flag AI if confidence > 70%
    threshold = 0.70
    is_ai = ai_score > threshold and ai_score > human_score
    
    return is_ai, ai_score, human_score


def main():
    st.markdown("<h2 style='text-align:center; margin:0;'>🔍</h2>", unsafe_allow_html=True)
    
    uploaded = st.file_uploader("", type=['png', 'jpg', 'jpeg', 'webp'], label_visibility="collapsed")
    
    if uploaded:
        image = Image.open(io.BytesIO(uploaded.read()))
        st.image(image, use_container_width=True)
        
        with st.spinner(""):
            try:
                detector = load_detector()
                is_ai, ai_conf, human_conf = predict(detector, image)
                
                if is_ai:
                    pct = int(ai_conf * 100)
                    st.markdown(f"""
                    <div style='text-align:center; padding:1.5rem 0;'>
                        <div style='font-size:4rem; font-weight:bold; color:#e74c3c;'>AI</div>
                        <div style='font-size:2rem; color:#666;'>{pct}%</div>
                    </div>
                    """, unsafe_allow_html=True)
                else:
                    pct = int(human_conf * 100)
                    st.markdown(f"""
                    <div style='text-align:center; padding:1.5rem 0;'>
                        <div style='font-size:4rem; font-weight:bold; color:#27ae60;'>REAL</div>
                        <div style='font-size:2rem; color:#666;'>{pct}%</div>
                    </div>
                    """, unsafe_allow_html=True)
                    
            except Exception as e:
                st.error(str(e))


if __name__ == "__main__":
    main()