#!/usr/bin/env python3
"""
Qwen Image Generation MCP Server
Provides image generation via Qwen (DashScope) API
"""

import os
import requests
from urllib.parse import urlparse
from mcp.server.fastmcp import FastMCP


# Initialize MCP server
mcp = FastMCP("qwen-image-generator")


@mcp.tool()
def generate_image(
    prompt: str,
    negative_prompt: str | None = None,
    prompt_extend: bool = True,
    size: str = "1024*1024",
    n: int = 1,
    image_url: str | None = None,
    output_path: str | None = None
) -> str:
    """
    Generate images using Qwen (通义千问) image generation API.
    
    Args:
        prompt: The image generation prompt describing what to generate (max 800 chars)
        negative_prompt: Negative prompt to avoid certain elements (max 500 chars)
        prompt_extend: Enable prompt extend to enhance the prompt (default: true)
        watermark: Add watermark to generated image (default: false)
        size: Image resolution. Options: 1024*1024 (1:1), 1344*768 (16:9), 768*1344 (9:16), 1184*864 (4:3), 864*1184 (3:4)
        n: Number of images to generate, 1-6 (default: 1)
        image_url: Optional reference image URL for image-to-image generation
        output_path: Optional local path to save the generated image
    
    Returns:
        Image URLs and generation status
    """
    api_key = os.environ.get("DASHSCOPE_API_KEY", "")
    api_base = "https://dashscope.aliyuncs.com"
    
    if not api_key:
        return "Error: DASHSCOPE_API_KEY environment variable is not set"
    
    # Build content array
    content = [{"text": prompt}]
    
    # Add reference image if provided
    if image_url:
        content.append({"image_url": {"url": image_url}})
    
    # Build parameters dict
    parameters = {
        "prompt_extend": prompt_extend,
        "size": size,
        "n": n
    }
    
    # Only add negative_prompt if provided
    if negative_prompt:
        parameters["negative_prompt"] = negative_prompt
    
    # Build request payload
    payload = {
        "model": "qwen-image-2.0-pro",
        "input": {
            "messages": [
                {
                    "role": "user",
                    "content": content
                }
            ]
        },
        "parameters": parameters
    }
    
    # Make API request
    url = f"{api_base}/api/v1/services/aigc/multimodal-generation/generation"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    try:
        response = requests.post(url, headers=headers, json=payload, timeout=180)
        response.raise_for_status()
        result = response.json()
        
        # Check for API errors
        if "error" in result:
            error_msg = result.get("error", {}).get("message", "Unknown error")
            return f"API Error: {error_msg}"
        
        # Parse response
        choices = result.get("output", {}).get("choices", [])
        usage = result.get("usage", {})
        
        # Extract image URLs
        image_urls = []
        for choice in choices:
            message = choice.get("message", {})
            content_items = message.get("content", [])
            for item in content_items:
                if "image" in item:
                    image_urls.append(item["image"])
        
        width = usage.get("width", 1024)
        height = usage.get("height", 1024)
        request_id = result.get("request_id", "N/A")
        
        # Save image if output_path provided
        saved_path = None
        if output_path and image_urls:
            try:
                img_response = requests.get(image_urls[0], timeout=30)
                img_response.raise_for_status()
                
                # Determine file extension
                parsed = urlparse(image_urls[0])
                ext = os.path.splitext(parsed.path)[1] if "." in parsed.path else ".png"
                if not ext or len(ext) > 5:
                    ext = ".png"
                
                # Ensure directory exists
                os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else ".", exist_ok=True)
                
                # Add extension if not present
                if not output_path.endswith(ext):
                    output_path = output_path + ext
                
                with open(output_path, "wb") as f:
                    f.write(img_response.content)
                saved_path = os.path.abspath(output_path)
            except Exception as e:
                return f"Failed to save image: {str(e)}"
        
        # Build response
        response_text = f"Successfully generated {n} image(s) ({width}x{height})\n\n"
        response_text += f"Request ID: {request_id}\n\n"
        
        if saved_path:
            response_text += f"Saved to: {saved_path}\n\n"
        
        response_text += "Image URLs:\n"
        for i, img_url in enumerate(image_urls, 1):
            response_text += f"  {i}. {img_url}\n"
        
        return response_text
        
    except requests.exceptions.RequestException as e:
        return f"Request Error: {str(e)}"
    except Exception as e:
        return f"Unexpected Error: {str(e)}"


if __name__ == "__main__":
    mcp.run()