#!/usr/bin/env python3
"""
A1D MCP Server with Gradio - Official MCP Implementation
Following Gradio's official MCP documentation
"""

import gradio as gr
import os
from utils import A1DAPIClient, prepare_request_data, format_response_with_preview
from config import TOOLS_CONFIG

def get_api_client():
    """Get API client with current API key"""
    api_key = os.getenv("A1D_API_KEY")
    if not api_key:
        raise ValueError("A1D_API_KEY environment variable is required")
    return A1DAPIClient(api_key)

def remove_background(image_url: str) -> str:
    """
    Remove background from images using AI.

    Args:
        image_url (str): The URL of the image to remove background from. Must be a valid HTTP/HTTPS URL pointing to an image file.

    Returns:
        str: A message indicating the result and the URL of the processed image.
    """
    try:
        if not image_url or not image_url.strip():
            return "❌ Error: Please provide an image URL"
        
        client = get_api_client()
        data = prepare_request_data("remove_bg", image_url=image_url)
        response = client.make_request_with_result(
            TOOLS_CONFIG["remove_bg"]["api_endpoint"],
            data,
            timeout=120
        )
        message, media_url = format_response_with_preview(response, "remove_bg")
        return f"{message}\n\nResult URL: {media_url}" if media_url else message
    except Exception as e:
        return f"❌ Error: {str(e)}"

def upscale_image(image_url: str, scale: str = "2") -> str:
    """
    Upscale images using AI.

    Args:
        image_url (str): The URL of the image to upscale. Must be a valid HTTP/HTTPS URL pointing to an image file.
        scale (str): Scale factor for upscaling. Choose from "2", "4", "8", or "16". Higher values produce larger images but take longer to process.

    Returns:
        str: A message indicating the result and the URL of the processed image.
    """
    try:
        if not image_url or not image_url.strip():
            return "❌ Error: Please provide an image URL"
        
        client = get_api_client()
        data = prepare_request_data("image_upscaler", image_url=image_url, scale=int(scale))
        response = client.make_request_with_result(
            TOOLS_CONFIG["image_upscaler"]["api_endpoint"],
            data,
            timeout=120
        )
        message, media_url = format_response_with_preview(response, "image_upscaler")
        return f"{message}\n\nResult URL: {media_url}" if media_url else message
    except Exception as e:
        return f"❌ Error: {str(e)}"

def generate_image(prompt: str) -> str:
    """
    Generate images using AI from text prompts.

    Args:
        prompt (str): Text description of the image to generate. Be descriptive and specific for better results. Example: "A beautiful sunset over mountains with vibrant orange and purple colors".

    Returns:
        str: A message indicating the result and the URL of the generated image.
    """
    try:
        if not prompt or not prompt.strip():
            return "❌ Error: Please provide a text prompt"
        
        client = get_api_client()
        data = prepare_request_data("image_generator", prompt=prompt)
        response = client.make_request_with_result(
            TOOLS_CONFIG["image_generator"]["api_endpoint"],
            data,
            timeout=120
        )
        message, media_url = format_response_with_preview(response, "image_generator")
        return f"{message}\n\nResult URL: {media_url}" if media_url else message
    except Exception as e:
        return f"❌ Error: {str(e)}"

def vectorize_image(image_url: str) -> str:
    """
    Convert images to vector format using AI.

    Args:
        image_url (str): The URL of the image to convert to vector format. Must be a valid HTTP/HTTPS URL pointing to an image file.

    Returns:
        str: A message indicating the result and the URL of the vectorized image.
    """
    try:
        if not image_url or not image_url.strip():
            return "❌ Error: Please provide an image URL"
        
        client = get_api_client()
        data = prepare_request_data("image_vectorization", image_url=image_url)
        response = client.make_request_with_result(
            TOOLS_CONFIG["image_vectorization"]["api_endpoint"],
            data,
            timeout=120
        )
        message, media_url = format_response_with_preview(response, "image_vectorization")
        return f"{message}\n\nResult URL: {media_url}" if media_url else message
    except Exception as e:
        return f"❌ Error: {str(e)}"

def extend_image(image_url: str) -> str:
    """
    Extend images using AI.

    Args:
        image_url (str): The URL of the image to extend. Must be a valid HTTP/HTTPS URL pointing to an image file.

    Returns:
        str: A message indicating the result and the URL of the extended image.
    """
    try:
        if not image_url or not image_url.strip():
            return "❌ Error: Please provide an image URL"
        
        client = get_api_client()
        data = prepare_request_data("image_extends", image_url=image_url)
        response = client.make_request_with_result(
            TOOLS_CONFIG["image_extends"]["api_endpoint"],
            data,
            timeout=120
        )
        message, media_url = format_response_with_preview(response, "image_extends")
        return f"{message}\n\nResult URL: {media_url}" if media_url else message
    except Exception as e:
        return f"❌ Error: {str(e)}"

def upscale_video(video_url: str) -> str:
    """
    Upscale videos using AI.

    Args:
        video_url (str): The URL of the video to upscale. Must be a valid HTTP/HTTPS URL pointing to a video file (MP4, AVI, MOV, etc.).

    Returns:
        str: A message indicating the result and the URL of the upscaled video.
    """
    try:
        if not video_url or not video_url.strip():
            return "❌ Error: Please provide a video URL"
        
        client = get_api_client()
        data = prepare_request_data("video_upscaler", video_url=video_url)
        response = client.make_request_with_result(
            TOOLS_CONFIG["video_upscaler"]["api_endpoint"],
            data,
            timeout=300
        )
        message, media_url = format_response_with_preview(response, "video_upscaler")
        return f"{message}\n\nResult URL: {media_url}" if media_url else message
    except Exception as e:
        return f"❌ Error: {str(e)}"

# Create the Gradio interface using the official MCP approach
demo = gr.Interface(
    fn=remove_background,
    inputs=[gr.Textbox(label="Image URL", placeholder="https://example.com/image.jpg", value="https://asset2.qieman.com/user-avatar/20240721/1071394_e15eb706-485d-4604-ad28-ac9830de1fe4.jpg")],
    outputs=[gr.Textbox(label="Result")],
    title="🤖 A1D MCP Server - Universal AI Tools",
    description="Powerful AI image and video processing tools for any MCP-compatible client.",
    examples=[
        ["https://asset2.qieman.com/user-avatar/20240721/1071394_e15eb706-485d-4604-ad28-ac9830de1fe4.jpg"],
        ["https://images.unsplash.com/photo-1506905925346-21bda4d32df4"]
    ]
)

if __name__ == "__main__":
    # Check for API key
    if not os.getenv("A1D_API_KEY"):
        print("❌ Error: A1D_API_KEY environment variable is required")
        print("Please set your API key in the Space settings")
        exit(1)
    
    print("🚀 Starting A1D MCP Server (Official Implementation)...")
    print("✅ API key found")
    print("🌐 MCP Server will be available at: /gradio_api/mcp/sse")
    
    # Launch with MCP server enabled (official method)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        mcp_server=True  # This is the official way to enable MCP
    )