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
    user_agent = ""
    request_path = ""

    # 1. Try from request headers (for MCP clients)
    try:
        request = gr.request()
        if request and hasattr(request, 'headers'):
            headers = dict(request.headers)
            api_key = get_api_key_from_headers(headers)
            user_agent = headers.get('user-agent', '')
            request_path = getattr(request, 'url', {}).path if hasattr(
                request, 'url') else ""
            print(f"🔍 Request headers found - User-Agent: {user_agent}")
            print(f"🔍 Request path: {request_path}")
            print(
                f"🔍 API key from headers: {'Found' if api_key else 'Not found'}")
    except Exception as e:
        print(f"🔍 No request context available: {e}")

    # 2. Check if running on Hugging Face Space
    is_space = os.getenv("SPACE_ID") is not None
    space_api_key = os.getenv("A1D_API_KEY")
    print(
        f"🔍 Environment check - Is Space: {is_space}, Space API key: {'Found' if space_api_key else 'Not found'}")

    # 3. Determine if this is a web browser request or MCP client request
    is_web_request = False
    is_mcp_request = False

    # Check if this is an MCP request
    if request_path and ('/mcp/' in request_path or '/gradio_api/mcp' in request_path):
        is_mcp_request = True
        print("🔍 Detected MCP API request")

    if user_agent:
        user_agent_lower = user_agent.lower()
        # Web browsers typically have 'mozilla' in user agent
        is_web_request = ('mozilla' in user_agent_lower or
                          'chrome' in user_agent_lower or
                          'safari' in user_agent_lower or
                          'edge' in user_agent_lower)
        print(f"🔍 Request type detection - Is web request: {is_web_request}")
    else:
        # If no user agent, assume it's NOT a web request (likely MCP client)
        is_web_request = False
        print("🔍 No User-Agent found - assuming MCP client request")

    # 4. STRICT RULE: MCP requests MUST have API key
    if is_mcp_request and not api_key:
        error_msg = (
            "🔑 API key is REQUIRED for MCP requests!\n\n"
            "This is an MCP API endpoint. You must provide your API key.\n"
            "Get your API key at https://a1d.ai\n\n"
            "Configuration example:\n"
            '{\n'
            '  "mcpServers": {\n'
            '    "a1d": {\n'
            '      "command": "npx",\n'
            '      "args": [\n'
            '        "mcp-remote@latest",\n'
            '        "https://aigchacker-a1d-mcp-server.hf.space/gradio_api/mcp/sse",\n'
            '        "--header",\n'
            '        "API_KEY:${MCP_API_KEY}"\n'
            '      ],\n'
            '      "env": {\n'
            '        "MCP_API_KEY": "your_a1d_api_key_here"\n'
            '      }\n'
            '    }\n'
            '  }\n'
            '}'
        )
        print(f"❌ MCP API key validation failed: {error_msg}")
        raise ValueError(error_msg)

    # 5. Use Space API key ONLY for web browser requests on Hugging Face Space
    if not api_key and is_space and space_api_key and is_web_request and not is_mcp_request:
        print("📡 Using API key from Space environment variable (web demo)")
        return A1DAPIClient(space_api_key)

    # 6. For all other cases, user API key is mandatory
    if not api_key:
        error_msg = (
            "🔑 API key is required!\n\n"
            "Please provide API_KEY in request headers.\n"
            "Get your API key at https://a1d.ai\n\n"
            "Configuration example:\n"
            '{\n'
            '  "mcpServers": {\n'
            '    "a1d": {\n'
            '      "command": "npx",\n'
            '      "args": [\n'
            '        "mcp-remote@latest",\n'
            '        "https://aigchacker-a1d-mcp-server.hf.space/gradio_api/mcp/sse",\n'
            '        "--header",\n'
            '        "API_KEY:${MCP_API_KEY}"\n'
            '      ],\n'
            '      "env": {\n'
            '        "MCP_API_KEY": "your_a1d_api_key_here"\n'
            '      }\n'
            '    }\n'
            '  }\n'
            '}'
        )
        print(f"❌ API key validation failed: {error_msg}")
        raise ValueError(error_msg)

    print("🔑 Using API key from MCP client 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


# MCP Documentation Component
def create_mcp_docs():
    """Create MCP documentation component"""
    return gr.Markdown("""
## 🔧 MCP Client Configuration

### ⚠️ API Key Required for Client Usage

When using with your own MCP client (Claude Desktop, Cursor, etc.), you **must** provide your API key:

```json
{
  "mcpServers": {
    "a1d": {
      "command": "npx",
      "args": [
        "mcp-remote@latest",
        "https://aigchacker-a1d-mcp-server.hf.space/gradio_api/mcp/sse",
        "--header",
        "API_KEY:${MCP_API_KEY}"
      ],
      "env": {
        "MCP_API_KEY": "your_a1d_api_key_here"
      }
    }
  }
}
```

**🔑 API key is mandatory for client usage.** Get your API key at [A1D.ai](https://a1d.ai/home/api).

---

### 🌐 Using the Hosted Demo on Hugging Face Space

The hosted demo at [https://huggingface.co/spaces/aigchacker/a1d-mcp-server](https://huggingface.co/spaces/aigchacker/a1d-mcp-server) uses our provided API key for demonstration purposes only, with limited usage. For production use, please obtain your own API key.

---

### 🔑 How to Get Your A1D API Key:

1. **Visit A1D Website**: Go to [https://a1d.ai](https://a1d.ai)
2. **Sign Up/Login**: Create an account or login to your existing account
3. **Access Dashboard**: Navigate to your user dashboard
4. **Generate API Key**: Look for "API Keys" or "Developer" section
5. **Copy Your Key**: Copy the generated API key
6. **Replace in Config**: Replace `your_a1d_api_key_here` with your actual API key

---

### 📋 Available MCP Tools:
- `remove_bg_wrapper` - Remove background from images
- `image_upscaler_wrapper` - Upscale images (2x/4x/8x/16x)
- `video_upscaler_wrapper` - Upscale videos
- `image_vectorization_wrapper` - Convert images to vector format
- `image_extends_wrapper` - Extend images using AI
- `image_generator_wrapper` - Generate images from text prompts

### 🌐 MCP Endpoints:
- **SSE Endpoint**: `https://aigchacker-a1d-mcp-server.hf.space/gradio_api/mcp/sse`
- **Schema**: `https://aigchacker-a1d-mcp-server.hf.space/gradio_api/mcp/schema`

### 💡 Usage Summary:
- **Hosted Demo**: Works directly in browser with provided API key
- **MCP Client**: Requires your own API key for production use
""")

# Create Gradio interfaces for each tool


def create_gradio_app():
    """Create the main Gradio application with all tools"""

    # Background Removal Interface with MCP docs
    with gr.Blocks(title="🎭 Background Removal") as remove_bg_interface:
        gr.Markdown("# 🎭 Background Removal")
        gr.Markdown("Remove background from images using AI")

        with gr.Row():
            with gr.Column():
                bg_input = gr.Textbox(
                    label="Image URL",
                    placeholder="https://example.com/image.jpg",
                    info="Enter the URL of the image to remove background from"
                )
                bg_button = gr.Button("Remove Background", variant="primary")
            with gr.Column():
                bg_result = gr.Textbox(label="Result")
                bg_preview = gr.Image(label="Preview")

        bg_button.click(
            fn=remove_bg_wrapper,
            inputs=[bg_input],
            outputs=[bg_result, bg_preview]
        )

        # Add MCP documentation
        create_mcp_docs()

    # Image Upscaler Interface with MCP docs
    with gr.Blocks(title="🔍 Image Upscaler") as image_upscaler_interface:
        gr.Markdown("# 🔍 Image Upscaler")
        gr.Markdown("Upscale images using AI with specified scale factor")

        with gr.Row():
            with gr.Column():
                up_input = gr.Textbox(
                    label="Image URL",
                    placeholder="https://example.com/image.jpg",
                    info="Enter the URL of the image to upscale"
                )
                up_scale = gr.Dropdown(
                    choices=[2, 4, 8, 16],
                    value=2,
                    label="Scale Factor",
                    info="Choose the upscaling factor"
                )
                up_button = gr.Button("Upscale Image", variant="primary")
            with gr.Column():
                up_result = gr.Textbox(label="Result")
                up_preview = gr.Image(label="Preview")

        up_button.click(
            fn=image_upscaler_wrapper,
            inputs=[up_input, up_scale],
            outputs=[up_result, up_preview]
        )

        # Add MCP documentation
        create_mcp_docs()

    # Video Upscaler Interface with MCP docs
    with gr.Blocks(title="🎬 Video Upscaler") as video_upscaler_interface:
        gr.Markdown("# 🎬 Video Upscaler")
        gr.Markdown("Upscale videos using AI")

        with gr.Row():
            with gr.Column():
                vid_input = gr.Textbox(
                    label="Video URL",
                    placeholder="https://example.com/video.mp4",
                    info="Enter the URL of the video to upscale"
                )
                vid_button = gr.Button("Upscale Video", variant="primary")
            with gr.Column():
                vid_result = gr.Textbox(label="Result")
                vid_preview = gr.Video(label="Preview")

        vid_button.click(
            fn=video_upscaler_wrapper,
            inputs=[vid_input],
            outputs=[vid_result, vid_preview]
        )

        # Add MCP documentation
        create_mcp_docs()

    # Image Vectorization Interface with MCP docs
    with gr.Blocks(title="📐 Image Vectorization") as image_vectorization_interface:
        gr.Markdown("# 📐 Image Vectorization")
        gr.Markdown("Convert images to vector format using AI")

        with gr.Row():
            with gr.Column():
                vec_input = gr.Textbox(
                    label="Image URL",
                    placeholder="https://example.com/image.jpg",
                    info="Enter the URL of the image to convert to vector format"
                )
                vec_button = gr.Button("Vectorize Image", variant="primary")
            with gr.Column():
                vec_result = gr.Textbox(label="Result")
                vec_preview = gr.Image(label="Preview")

        vec_button.click(
            fn=image_vectorization_wrapper,
            inputs=[vec_input],
            outputs=[vec_result, vec_preview]
        )

        # Add MCP documentation
        create_mcp_docs()

    # Image Extension Interface with MCP docs
    with gr.Blocks(title="🖼️ Image Extension") as image_extends_interface:
        gr.Markdown("# 🖼️ Image Extension")
        gr.Markdown("Extend images using AI")

        with gr.Row():
            with gr.Column():
                ext_input = gr.Textbox(
                    label="Image URL",
                    placeholder="https://example.com/image.jpg",
                    info="Enter the URL of the image to extend"
                )
                ext_button = gr.Button("Extend Image", variant="primary")
            with gr.Column():
                ext_result = gr.Textbox(label="Result")
                ext_preview = gr.Image(label="Preview")

        ext_button.click(
            fn=image_extends_wrapper,
            inputs=[ext_input],
            outputs=[ext_result, ext_preview]
        )

        # Add MCP documentation
        create_mcp_docs()

    # Image Generator Interface with MCP docs
    with gr.Blocks(title="🎨 Image Generator") as image_generator_interface:
        gr.Markdown("# 🎨 Image Generator")
        gr.Markdown("Generate images using AI from text prompts")

        with gr.Row():
            with gr.Column():
                gen_input = gr.Textbox(
                    label="Text Prompt",
                    placeholder="A beautiful sunset over mountains",
                    info="Enter a text description to generate an image",
                    lines=3
                )
                gen_button = gr.Button("Generate Image", variant="primary")
            with gr.Column():
                gen_result = gr.Textbox(label="Result")
                gen_preview = gr.Image(label="Preview")

        gen_button.click(
            fn=image_generator_wrapper,
            inputs=[gen_input],
            outputs=[gen_result, gen_preview]
        )

        # Add MCP documentation
        create_mcp_docs()

    # 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__":
    print("🚀 Starting A1D MCP Server...")
    print("=" * 70)

    # Check environment and show configuration
    is_space = os.getenv("SPACE_ID") is not None
    space_api_key = os.getenv("A1D_API_KEY")

    if is_space and space_api_key:
        print("✅ Running on Hugging Face Space with demo API key")
        print("🌐 Web demo: Works directly with provided API key")
        print("🔑 MCP clients: Must provide their own API key")
    elif is_space:
        print("⚠️  Running on Hugging Face Space without API key")
        print("🔑 All users must provide their own API key")
    else:
        print("🖥️  Running locally")
        if space_api_key:
            print("✅ Local API key found")
        else:
            print("🔑 Users must provide their own API key")

    print(f"\n🎯 Title: {GRADIO_CONFIG['title']}")
    print(
        f"🌐 Server: http://{GRADIO_CONFIG['server_name']}:{GRADIO_CONFIG['server_port']}")
    print(
        f"🔧 MCP Endpoint: http://{GRADIO_CONFIG['server_name']}:{GRADIO_CONFIG['server_port']}/gradio_api/mcp/sse")
    print(f"📋 Available Tools: {len(TOOLS_CONFIG)} AI tools")

    print(f"\n📖 Authentication Strategy:")
    print("   - Space demo: Uses provided API key (limited usage)")
    print("   - MCP clients: Must provide own API key (production usage)")
    print("   - Get API key at: https://a1d.ai")
    print("=" * 70)

    # Create and launch the app
    demo = create_gradio_app()

    # Launch the Gradio app with MCP server enabled
    demo.launch(
        server_name=GRADIO_CONFIG["server_name"],
        server_port=GRADIO_CONFIG["server_port"],
        share=GRADIO_CONFIG["share"],
        mcp_server=True  # Enable MCP server functionality
    )