#!/usr/bin/env python3
"""
MiniMax Image Generation MCP Server
Provides image generation via MiniMax API
"""

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


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


def read_local_image(file_path: str) -> str | None:
    """
    Read local image file and return as base64 encoded string.
    Returns the base64 string or None if failed.
    """
    try:
        with open(file_path, "rb") as f:
            image_data = f.read()
            return base64.b64encode(image_data).decode("utf-8")
    except Exception as e:
        return None


def build_subject_reference(subject_ref: str, subject_type: str) -> dict | None:
    """
    Build subject_reference object from URL or local file path.
    
    Args:
        subject_ref: URL or local file path
        subject_type: Type of subject (character, product, logo, etc.)
    
    Returns:
        dict with subject_reference structure or None if failed
    """
    # Check if it's a URL or local file
    if subject_ref.startswith(("http://", "https://")):
        # It's a URL
        return {
            "type": subject_type,
            "image_file": subject_ref
        }
    else:
        # It's a local file - convert to base64
        base64_data = read_local_image(subject_ref)
        if base64_data:
            return {
                "type": subject_type,
                "image_file": f"data:image/jpeg;base64,{base64_data}"
            }
        return None


@mcp.tool()
def generate_image(
    prompt: str,
    model: str = "image-01",
    aspect_ratio: str = "1:1",
    n: int = 1,
    prompt_optimizer: bool = False,
    subject_reference: str | None = None,
    subject_type: str = "character",
    seed: int | None = None,
    output_path: str | None = None
) -> str:
    """
    Generate images using MiniMax image generation API.
    
    Args:
        prompt: The image generation prompt describing what to generate
        model: Image generation model (default: image-01)
        aspect_ratio: Image aspect ratio (default: 1:1, options: 1:1, 16:9, 9:16, 4:3, 3:4)
        n: Number of images to generate (default: 1, max: 3)
        prompt_optimizer: Enable prompt optimizer (default: false)
        subject_reference: Reference image URL or local file path for image-to-image generation
        subject_type: Subject reference type (default: character, options: character, product, logo, video_subject, other)
        seed: Random seed for reproducible generation
        output_path: Optional local path to save the generated image
    
    Returns:
        Image URLs and generation status
    """
    api_key = os.environ.get("MINIMAX_API_KEY", "")
    api_base = "https://api.minimaxi.com"
    
    if not api_key:
        return "Error: MINIMAX_API_KEY environment variable is not set"
    
    # Build request payload
    payload = {
        "model": model,
        "prompt": prompt,
        "aspect_ratio": aspect_ratio,
        "response_format": "url",
        "n": n,
        "prompt_optimizer": prompt_optimizer
    }
    
    # Add seed if provided
    if seed is not None:
        payload["seed"] = seed
    
    # Add subject_reference for image-to-image generation
    if subject_reference:
        subject_ref = build_subject_reference(subject_reference, subject_type)
        if subject_ref:
            payload["subject_reference"] = [subject_ref]
    
    # Make API request
    url = f"{api_base}/v1/image_generation"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    try:
        response = requests.post(url, headers=headers, json=payload, timeout=60)
        response.raise_for_status()
        result = response.json()
        
        # Check for API errors
        if result.get("base_resp", {}).get("status_code") != 0:
            error_msg = result.get("base_resp", {}).get("status_msg", "Unknown error")
            return f"API Error: {error_msg}"
        
        # Extract image URLs
        image_urls = result.get("data", {}).get("image_urls", [])
        metadata = result.get("metadata", {})
        request_id = result.get("id", "N/A")
        success_count = metadata.get("success_count", len(image_urls))
        
        # Save image if output_path provided
        saved_paths = []
        if output_path and image_urls:
            try:
                import time
                timestamp = int(time.time())
                
                for i, img_url in enumerate(image_urls, 1):
                    img_response = requests.get(img_url, timeout=30)
                    img_response.raise_for_status()
                    
                    # Determine file extension
                    parsed = urlparse(img_url)
                    ext = os.path.splitext(parsed.path)[1] if "." in parsed.path else ".jpeg"
                    if not ext or len(ext) > 5:
                        ext = ".jpeg"
                    
                    # Handle multiple images
                    if len(image_urls) > 1:
                        base_path = output_path.rsplit('.', 1)[0] if '.' in output_path else output_path
                        file_ext = output_path.rsplit('.', 1)[1] if '.' in output_path else ext
                        img_path = f"{base_path}_{i}_{timestamp}.{file_ext}"
                    else:
                        if not output_path.endswith(ext):
                            img_path = output_path + ext
                        else:
                            img_path = output_path
                    
                    # Ensure directory exists
                    os.makedirs(os.path.dirname(img_path) if os.path.dirname(img_path) else ".", exist_ok=True)
                    
                    with open(img_path, "wb") as f:
                        f.write(img_response.content)
                    saved_paths.append(os.path.abspath(img_path))
            except Exception as e:
                return f"Failed to save image: {str(e)}"
        
        # Build response
        response_text = f"Successfully generated {success_count} image(s)\n\n"
        response_text += f"Request ID: {request_id}\n\n"
        
        if saved_paths:
            response_text += "Saved to:\n"
            for path in saved_paths:
                response_text += f"  - {path}\n"
            response_text += "\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()