app.py

"""
A1D MCP Server - Gradio Application
Universal AI Tools for image and video processing
"""

import gradio as gr
import os
from typing import Optional, Tuple, Union
from utils import A1DAPIClient, validate_url, validate_scale, prepare_request_data, format_response_with_preview
from config import GRADIO_CONFIG, TOOLS_CONFIG
from mcp_handler import get_api_key_from_headers


# Initialize API client
def get_api_client():
    """Get API client with current API key"""
    # Try to get API key from multiple sources
    api_key = None

    # 1. Try from request headers (for MCP clients)
    try:
        request = gr.request()
        if request and hasattr(request, 'headers'):
            api_key = get_api_key_from_headers(dict(request.headers))
    except:
        pass

    # 2. Fallback to environment variable
    if not api_key:
        api_key = os.getenv("A1D_API_KEY")

    if not api_key:
        raise ValueError(
            "API key is required. Set A1D_API_KEY environment variable or provide API_KEY in request headers.")

    return A1DAPIClient(api_key)


def remove_bg(image_url: str) -> Tuple[str, Optional[str]]:
    """Remove background from images using AI.

    Args:
        image_url: The URL of the image to remove background from

    Returns:
        Tuple of (result_message, media_url_for_preview)
    """
    try:
        if not validate_url(image_url):
            return "❌ Error: Invalid image URL format", None

        client = get_api_client()
        data = prepare_request_data("remove_bg", image_url=image_url)

        # Use the new method that waits for result
        response = client.make_request_with_result(
            TOOLS_CONFIG["remove_bg"]["api_endpoint"],
            data,
            timeout=120  # 2 minutes timeout
        )

        return format_response_with_preview(response, "remove_bg")

    except Exception as e:
        return f"❌ Error: {str(e)}", None


def image_upscaler(image_url: str, scale: int = 2) -> Tuple[str, Optional[str]]:
    """Upscale images using AI with specified scale factor.

    Args:
        image_url: The URL of the image to upscale
        scale: Scale factor for upscaling (2, 4, 8, or 16). Default: 2

    Returns:
        Tuple of (result_message, media_url_for_preview)
    """
    try:
        if not validate_url(image_url):
            return "❌ Error: Invalid image URL format", None

        if not validate_scale(scale):
            return "❌ Error: Scale must be 2, 4, 8, or 16", None

        client = get_api_client()
        data = prepare_request_data(
            "image_upscaler", image_url=image_url, scale=scale)

        response = client.make_request_with_result(
            TOOLS_CONFIG["image_upscaler"]["api_endpoint"],
            data,
            timeout=120
        )

        return format_response_with_preview(response, "image_upscaler")

    except Exception as e:
        return f"❌ Error: {str(e)}", None


def video_upscaler(video_url: str) -> Tuple[str, Optional[str]]:
    """Upscale videos using AI.

    Args:
        video_url: The URL of the video to upscale

    Returns:
        Tuple of (result_message, media_url_for_preview)
    """
    try:
        if not validate_url(video_url):
            return "❌ Error: Invalid video URL format", None

        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  # 5 minutes for video processing
        )

        return format_response_with_preview(response, "video_upscaler")

    except Exception as e:
        return f"❌ Error: {str(e)}", None


def image_vectorization(image_url: str) -> Tuple[str, Optional[str]]:
    """Convert images to vector format using AI.

    Args:
        image_url: The URL of the image to vectorize

    Returns:
        Tuple of (result_message, media_url_for_preview)
    """
    try:
        if not validate_url(image_url):
            return "❌ Error: Invalid image URL format", None

        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
        )

        return format_response_with_preview(response, "image_vectorization")

    except Exception as e:
        return f"❌ Error: {str(e)}", None


def image_extends(image_url: str) -> Tuple[str, Optional[str]]:
    """Extend images using AI.

    Args:
        image_url: The URL of the image to extend

    Returns:
        Tuple of (result_message, media_url_for_preview)
    """
    try:
        if not validate_url(image_url):
            return "❌ Error: Invalid image URL format", None

        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
        )

        return format_response_with_preview(response, "image_extends")

    except Exception as e:
        return f"❌ Error: {str(e)}", None


def image_generator(prompt: str) -> Tuple[str, Optional[str]]:
    """Generate images using AI from text prompts.

    Args:
        prompt: Text prompt to generate image from

    Returns:
        Tuple of (result_message, media_url_for_preview)
    """
    try:
        if not prompt or not prompt.strip():
            return "❌ Error: Prompt is required and cannot be empty", None

        client = get_api_client()
        data = prepare_request_data("image_generator", prompt=prompt.strip())

        response = client.make_request_with_result(
            TOOLS_CONFIG["image_generator"]["api_endpoint"],
            data,
            timeout=120
        )

        return format_response_with_preview(response, "image_generator")

    except Exception as e:
        return f"❌ Error: {str(e)}", None


# Wrapper functions for Gradio interface
def remove_bg_wrapper(image_url: str):
    """Wrapper for remove_bg that returns message and media for Gradio

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

    Returns:
        Tuple of (result_message, media_url_for_preview)
    """
    message, media_url = remove_bg(image_url)
    return message, media_url if media_url else None


def image_upscaler_wrapper(image_url: str, scale: int):
    """Wrapper for image_upscaler that returns message and media for Gradio

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

    Returns:
        Tuple of (result_message, media_url_for_preview)
    """
    message, media_url = image_upscaler(image_url, scale)
    return message, media_url if media_url else None


def video_upscaler_wrapper(video_url: str):
    """Wrapper for video_upscaler that returns message and media for Gradio

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

    Returns:
        Tuple of (result_message, media_url_for_preview)
    """
    message, media_url = video_upscaler(video_url)
    return message, media_url if media_url else None


def image_vectorization_wrapper(image_url: str):
    """Wrapper for image_vectorization that returns message and media for Gradio

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

    Returns:
        Tuple of (result_message, media_url_for_preview)
    """
    message, media_url = image_vectorization(image_url)
    return message, media_url if media_url else None


def image_extends_wrapper(image_url: str):
    """Wrapper for image_extends that returns message and media for Gradio

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

    Returns:
        Tuple of (result_message, media_url_for_preview)
    """
    message, media_url = image_extends(image_url)
    return message, media_url if media_url else None


def image_generator_wrapper(prompt: str):
    """Wrapper for image_generator that returns message and media for Gradio

    Args:
        prompt: 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:
        Tuple of (result_message, media_url_for_preview)
    """
    message, media_url = image_generator(prompt)
    return message, media_url if media_url else None


# Create Gradio interfaces for each tool
def create_gradio_app():
    """Create the main Gradio application with all tools"""

    # Create individual interfaces for each tool
    remove_bg_interface = gr.Interface(
        fn=remove_bg_wrapper,
        inputs=[
            gr.Textbox(
                label="Image URL",
                placeholder="https://example.com/image.jpg",
                info="Enter the URL of the image to remove background from"
            )
        ],
        outputs=[
            gr.Textbox(label="Result"),
            gr.Image(label="Preview")
        ],
        title="🎭 Background Removal",
        description="Remove background from images using AI"
    )

    image_upscaler_interface = gr.Interface(
        fn=image_upscaler_wrapper,
        inputs=[
            gr.Textbox(
                label="Image URL",
                placeholder="https://example.com/image.jpg",
                info="Enter the URL of the image to upscale"
            ),
            gr.Dropdown(
                choices=[2, 4, 8, 16],
                value=2,
                label="Scale Factor",
                info="Choose the upscaling factor"
            )
        ],
        outputs=[
            gr.Textbox(label="Result"),
            gr.Image(label="Preview")
        ],
        title="🔍 Image Upscaler",
        description="Upscale images using AI with specified scale factor"
    )

    video_upscaler_interface = gr.Interface(
        fn=video_upscaler_wrapper,
        inputs=[
            gr.Textbox(
                label="Video URL",
                placeholder="https://example.com/video.mp4",
                info="Enter the URL of the video to upscale"
            )
        ],
        outputs=[
            gr.Textbox(label="Result"),
            gr.Video(label="Preview")
        ],
        title="🎬 Video Upscaler",
        description="Upscale videos using AI"
    )

    image_vectorization_interface = gr.Interface(
        fn=image_vectorization_wrapper,
        inputs=[
            gr.Textbox(
                label="Image URL",
                placeholder="https://example.com/image.jpg",
                info="Enter the URL of the image to convert to vector format"
            )
        ],
        outputs=[
            gr.Textbox(label="Result"),
            gr.Image(label="Preview")
        ],
        title="📐 Image Vectorization",
        description="Convert images to vector format using AI"
    )

    image_extends_interface = gr.Interface(
        fn=image_extends_wrapper,
        inputs=[
            gr.Textbox(
                label="Image URL",
                placeholder="https://example.com/image.jpg",
                info="Enter the URL of the image to extend"
            )
        ],
        outputs=[
            gr.Textbox(label="Result"),
            gr.Image(label="Preview")
        ],
        title="🖼️ Image Extension",
        description="Extend images using AI"
    )

    image_generator_interface = gr.Interface(
        fn=image_generator_wrapper,
        inputs=[
            gr.Textbox(
                label="Text Prompt",
                placeholder="A beautiful sunset over mountains",
                info="Enter a text description to generate an image",
                lines=3
            )
        ],
        outputs=[
            gr.Textbox(label="Result"),
            gr.Image(label="Preview")
        ],
        title="🎨 Image Generator",
        description="Generate images using AI from text prompts"
    )

    # Create tabbed interface
    demo = gr.TabbedInterface(
        [
            remove_bg_interface,
            image_upscaler_interface,
            video_upscaler_interface,
            image_vectorization_interface,
            image_extends_interface,
            image_generator_interface
        ],
        [
            "Background Removal",
            "Image Upscaler",
            "Video Upscaler",
            "Image Vectorization",
            "Image Extension",
            "Image Generator"
        ],
        title=GRADIO_CONFIG["title"],
        theme=GRADIO_CONFIG["theme"]
    )

    return demo


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: export A1D_API_KEY=your_api_key_here")
        exit(1)

    # Create and launch the app
    demo = create_gradio_app()

    # Launch the Gradio app
    demo.launch(
        server_name=GRADIO_CONFIG["server_name"],
        server_port=GRADIO_CONFIG["server_port"],
        share=GRADIO_CONFIG["share"]
    )