import torch
from transformers import CLIPProcessor, CLIPModel
from PIL import Image

class ImageDeepfakeDetector:
    def __init__(self):
        print("⚡ Loading Image AI Model...")
        try:
            # Both model and processor now pull directly from Hugging Face!
            self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
            self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
            print("✅ Image Model Ready.")
        except Exception as e:
            print(f"❌ Error: {e}")
            exit()

    def predict(self, image_path):
        try:
            image = Image.open(image_path)
            
            # --- THE "PHYSICS" FIX ---
            # We are now detecting "Camera Physics" vs "Generative Patterns"
            labels = [
                # === REAL (Camera Physics) ===
                # Real cameras leave specific traces: noise, focus fall-off, organic skin.
                "photo taken with a canon or nikon dslr camera",    # Specific camera brand hints
                "real human skin with natural pores and imperfections", # Texture check
                "authentic photograph with natural lighting shadows",   # Lighting physics
                "candid photo from a smartphone camera",                # Lower quality real
                "high resolution raw photography",                      # HD Real

                # === FAKE (Generative Patterns) ===
                # AI generators leave specific traces: overly smooth, distorted backgrounds.
                "ai generated image by midjourney",          # Specific Generator
                "stable diffusion synthetic image",          # Specific Generator
                "deepfake face swap with artifacts",         # Manipulation
                "computer generated 3d character render",    # CGI look
                "perfectly smooth plastic skin texture",     # The "AI Glow"
                "unnatural distorted background details"     # Background check
            ]
            
            inputs = self.processor(
                text=labels, 
                images=image, 
                return_tensors="pt", 
                padding=True
            )

            with torch.no_grad():
                outputs = self.model(**inputs)
            
            probs = outputs.logits_per_image.softmax(dim=1)
            scores = probs.tolist()[0]
            
            # --- SCORING ---
            # Real = Indices 0 to 4 (5 labels)
            real_score = sum(scores[:5])
            
            # Fake = Indices 5 to 10 (6 labels)
            fake_score = sum(scores[5:])
            
            # --- "UNCERTAINTY" LOGIC ---
            # If the scores are very close (e.g., 49% Real vs 51% Fake), 
            # it means the AI is guessing. In court, "Innocent until proven guilty".
            # So, if the gap is small (< 5%), we assume it's REAL.
            
            score_diff = abs(real_score - fake_score)
            
            if fake_score > real_score:
                # It thinks it's fake, but is it sure?
                if score_diff < 0.05:
                    # Not sure enough -> Call it Real (but low confidence)
                    return "REAL", real_score
                else:
                    return "DEEPFAKE DETECTED", fake_score
            else:
                return "REAL", real_score

        except Exception as e:
            print(f"Error predicting {image_path}: {e}")
            return "ERROR", 0.0