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...")
image_detector: Pipeline = pipeline("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:
        image = Image.open(requests.get(image_url, stream=True, timeout=10).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 (NEW) ---
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 ---
# We now create a list of Interfaces to expose both tools
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",
            description="Tool to analyze an image's authenticity."
        ),
        gr.Interface(
            fn=validate,
            inputs=[],
            outputs="text",
            title="Validation Tool",
            description="Used by the Puch AI platform to validate shared servers."
        )
    ],
    ["Image Analyzer", "Validator"]
)

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