import os
import gradio as gr
from transformers import AutoImageProcessor, SiglipForImageClassification
from PIL import Image
import torch

# Step 2: Load the AI Classifier (Siglip2)
# This model is specifically fine-tuned to distinguish between Real and AI-generated images
model_name = "prithivMLmods/Deepfake-Detect-Siglip2"
model = SiglipForImageClassification.from_pretrained(model_name)
processor = AutoImageProcessor.from_pretrained(model_name)

# Step 3: Perception Logic
def analyze_kyc_image(image):
    if image is None:
        return "ERROR: No image provided."
        
    try:
        # Convert to RGB if needed (handles RGBA or grayscale)
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image).convert("RGB")
        
        # Pre-process the image for the Siglip2 architecture
        inputs = processor(images=image, return_tensors="pt")
        
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            # Convert logits to probabilities
            probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
        
        # Get labels ('Fake' vs 'Real')
        labels = model.config.id2label
        
        # Identify the top prediction
        top_idx = 0 if probs[0] > probs[1] else 1
        verdict = labels[top_idx]
        confidence = probs[top_idx]
        
        # Step 4: Human-Readable Reporting
        if verdict.lower() == "fake":
            return f"🚨 ALERT: AI MANIPULATION DETECTED!\n\n📊 Fake Confidence: {confidence:.1%}\n(Model suggests this image is AI-generated or digitally altered)"
        else:
            return f"✅ SUCCESS: GENUINE PHOTOGRAPH VERIFIED.\n\n📊 Real Confidence: {confidence:.1%}\n(No signs of deepfake synthesis detected)"
            
    except Exception as e:
        return f"ERROR analyzing image: {str(e)}"

# Step 5: Build the Sentinel Interface (Gradio)
with gr.Blocks(title="IOB Sentinel: Image KYC Verifier") as demo:
    gr.Markdown("# IOB Sentinel: Image KYC Deepfake Detector")
    gr.Markdown("Uses Siglip2-Base to identify AI-generated portraits and digital face-swaps.")
    
    with gr.Tabs():
        with gr.TabItem("Upload ID Photo"):
            with gr.Row():
                with gr.Column():
                    img_input = gr.Image(type="numpy", label="Upload Identity Document / Photo")
                    submit_btn = gr.Button("Analyze Image", variant="primary")
                with gr.Column():
                    text_output = gr.Textbox(label="Verification Result", interactive=False)
            
            submit_btn.click(fn=analyze_kyc_image, inputs=img_input, outputs=text_output)
            
        with gr.TabItem("Live Capture"):
            with gr.Row():
                with gr.Column():
                    webcam_input = gr.Image(sources=["webcam"], type="numpy", label="Capture Live Photo")
                    webcam_btn = gr.Button("Analyze Capture", variant="primary")
                with gr.Column():
                    webcam_output = gr.Textbox(label="Verification Result", interactive=False)
            
            webcam_btn.click(fn=analyze_kyc_image, inputs=webcam_input, outputs=webcam_output)

if __name__ == "__main__":
    demo.launch()