import gradio as gr
import requests
from PIL import Image
from transformers import pipeline, Pipeline
import os
from dotenv import load_dotenv
import io  # <-- NEW: Import the 'io' library

# --- 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...")
image_detector: Pipeline = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32", local_files_only=True)
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:
        # --- MODIFIED IMAGE DOWNLOAD BLOCK ---
        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'}
        
        # 1. Let requests handle the full download and any redirects
        response = requests.get(image_url, timeout=15, headers=headers)
        response.raise_for_status()  # Raise an exception for bad status codes (like 404 or 500)

        # 2. Open the image from the downloaded content in memory
        image = Image.open(io.BytesIO(response.content))
        # --- END OF MODIFIED BLOCK ---
    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)