import os
from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import pipeline

app = Flask(__name__)
CORS(app)

print("Loading Hugging Face AI Image Detector model...")
# Use a highly-rated community model for AI image detection
try:
    detector = pipeline("image-classification", model="prithivMLmods/Deepfake-Detect-Siglip2")
    print("Model loaded successfully.")
except Exception as e:
    print(f"Error loading model: {e}")
    detector = None

@app.route('/predict', methods=['POST'])
def predict():
    if 'image' not in request.files:
        return jsonify({'error': 'No image provided'}), 400
    
    file = request.files['image']
    if file.filename == '':
        return jsonify({'error': 'No empty file allowed'}), 400
    
    try:
        # Save temp file
        temp_path = "temp_upload.jpg"
        file.save(temp_path)
        
        if not detector:
            return jsonify({'error': 'Model not initialized'}), 500

        # Predict using the pipeline
        results = detector(temp_path)
        
        # Clean up
        if os.path.exists(temp_path):
            os.remove(temp_path)
            
        # The pipeline returns a list of dicts: [{'label': 'artificial', 'score': 0.9}, {'label': 'human', 'score': 0.1}]
        # We need to find the AI/artificial/fake label probability
        prob_fake = 0.0
        predicted_label = "real"
        
        # Sort results to find highest score
        top_result = max(results, key=lambda x: x['score'])
        predicted_label = top_result['label']
        
        for res in results:
            label = res['label'].lower()
            if 'ai' in label or 'fake' in label or 'artificial' in label or 'generated' in label:
                prob_fake = res['score'] * 100
                break
        
        return jsonify({
            'label': predicted_label,
            'ai_probability': round(prob_fake, 2)
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)