import gradio as gr
import requests
from PIL import Image
from transformers import pipeline, Pipeline
import os
from dotenv import load_dotenv

# --- Configuration ---
load_dotenv()
VALID_BEARER_TOKEN = os.getenv("VALID_BEARER_TOKEN")
OWNER_PHONE_NUMBER = os.getenv("OWNER_PHONE_NUMBER")

# --- AI Model Setup ---
print("Loading AI Image Detection model...")
# FIX 1: Use the correct pipeline for this task.
# 'zero-shot-image-classification' is designed to work with 'candidate_labels'.
image_detector: Pipeline = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32")
print("✅ Model loaded successfully.")

# --- Tool 1: The Main Analysis Function ---
def analyze_image_authenticity(image_url: str) -> dict:
    """
    Analyzes an image from a URL to determine if it is real or AI-generated.
    
    Args:
        image_url: The URL of the image to analyze.
        
    Returns:
        A dictionary with the analysis results and probability scores.
    """
    if not image_url:
        raise gr.Error("Image URL parameter is missing.")
    
    print(f"Analyzing image from URL: {image_url}")
    try:
        # FIX 2: Add a User-Agent header to mimic a real browser request.
        # This prevents websites (including Puch's image hosting) from blocking the download.
        headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.36'}
        image = Image.open(requests.get(image_url, stream=True, timeout=10, headers=headers).raw)
    except Exception as e:
        raise gr.Error(f"Could not load image from URL. Error: {str(e)}")

    labels = ["a real photograph", "a computer-generated image", "an illustration or drawing"]
    results = image_detector(image, candidate_labels=labels)
    
    print(f"Analysis successful. Results: {results}")
    return {"analysis_results": results}

# --- Tool 2: The Validation Function for Sharing ---
def validate() -> str:
    """
    Validates the server ownership for sharing. Returns the owner's phone number.
    
    Returns:
        The owner's phone number as a string.
    """
    return OWNER_PHONE_NUMBER

# --- Gradio Interface ---
demo = gr.TabbedInterface(
    [
        gr.Interface(
            fn=analyze_image_authenticity,
            inputs=[gr.Textbox(label="Image URL")],
            outputs=[gr.JSON(label="Analysis Results")],
            title="AI Image Authenticity Detector"
        ),
        gr.Interface(
            fn=validate,
            inputs=[],
            outputs="text",
            title="Validation Tool"
        )
    ],
    ["Image Analyzer", "Validator"]
)

# --- Launch the App and MCP Server ---
demo.launch(mcp_server=True)