"""
MCP Server implementation for text-to-image generation.

This server enables AI agents to generate images from text prompts using 
Doubao Seedream API or compatible image generation services.
AI agents should call this server when users request image generation,
visual content creation, or when asked to create pictures from descriptions.

IMPORTANT: Generated images are saved to the user's home directory under 'images' folder by default.
Users may need to move these files to their project directory as needed.
"""

import asyncio
import os
import sys
import logging
import tempfile
import requests
import json
import time
import random
from pathlib import Path

from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationOptions
import mcp.server.stdio
import mcp.types as types

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def setup_encoding():
    """Configure UTF-8 encoding for cross-platform compatibility."""
    if sys.platform == "win32":
        try:
            # Reconfigure all standard streams for UTF-8 on Windows
            sys.stdin.reconfigure(encoding='utf-8', errors='replace')
            sys.stdout.reconfigure(encoding='utf-8', errors='replace')
            sys.stderr.reconfigure(encoding='utf-8', errors='replace')
            os.environ['PYTHONIOENCODING'] = 'utf-8'
        except Exception as e:
            logger.warning(f"Failed to reconfigure stdio streams: {e}")


def image_to_base64(image_path: str) -> str:
    """Convert an image file to Base64 encoded string.
    
    Args:
        image_path: Path to the image file
        
    Returns:
        Base64 encoded string of the image with proper prefix
        
    Raises:
        FileNotFoundError: If the image file does not exist
        ValueError: If the file is not a valid image
    """
    import base64
    from pathlib import Path
    
    path = Path(image_path)
    if not path.exists():
        raise FileNotFoundError(f"Image file does not exist: {image_path}")
    
    # Determine image type from extension
    suffix = path.suffix.lower()
    if suffix not in ['.jpg', '.jpeg', '.png', '.webp', '.gif']:
        raise ValueError(f"Unsupported image format: {suffix}. Supported formats: .jpg, .jpeg, .png, .webp, .gif")
    
    # Read the image file
    with open(path, 'rb') as image_file:
        image_data = image_file.read()
    
    # Encode to base64
    base64_encoded = base64.b64encode(image_data).decode('utf-8')
    
    # Determine the appropriate MIME type
    mime_type = 'image/jpeg' if suffix in ['.jpg', '.jpeg'] else f'image/{suffix[1:]}'
    
    # Return as data URL format
    return f"data:{mime_type};base64,{base64_encoded}"


# Initialize server
server = Server("McpSeedream")

# Global configuration that can be set via MCP
config = {
    "api_url": os.getenv("ARK_API_URL", "https://ark.cn-beijing.volces.com/api/v3/images/generations"),
    "default_model": os.getenv("ARK_DEFAULT_MODEL", "doubao-seedream-4-5-251128"),
    "api_key": os.getenv("ARK_API_KEY")  # No default for API key for security reasons
}




@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    """List available text-to-image generation tools and configuration tools."""
    return [
        types.Tool(
            name="set_config",
            description="Configure API settings for image generation. Use this to set up the image generation service before generating images.",
            inputSchema={
                "type": "object",
                "properties": {
                    "api_url": {
                        "type": "string",
                        "description": "API endpoint URL for the image generation service."
                    },
                    "default_model": {
                        "type": "string",
                        "description": "Default model to use for image generation."
                    },
                    "api_key": {
                        "type": "string",
                        "description": "API key for authentication."
                    }
                },
                "required": []
            }
        ),
        types.Tool(
            name="generate_image",
            description="Generate an image from text prompt, single image, or multiple images using configured API. Use this when users ask to create images, generate pictures, visualize concepts, edit existing images, or turn text descriptions into visual content. Supports text-to-image, image-to-image, and multiple image fusion. IMPORTANT: You must specify the output directory and can optionally specify a target filename. If no target filename is provided, a filename will be generated using timestamp and random numbers.",
            inputSchema={
                "type": "object",
                "properties": {
                    "prompt": {
                        "type": "string",
                        "description": "Text description of the image to generate."
                    },
                    "model": {
                        "type": "string",
                        "description": "Model to use for generation (overrides default)."
                    },
                    "n": {
                        "type": "integer",
                        "description": "Number of images to generate (default: 1, max: 10).",
                        "default": 1
                    },
                    "size": {
                        "type": "string",
                        "description": "Size of the generated image (default: '2K'). For Doubao API, use '2K', '4K' or specific dimensions like '1024x1024'.",
                        "enum": ["256x256", "512x512", "1024x1024", "1792x1024", "1024x1792", "1K", "2K", "4K"]
                    },
                    "style": {
                        "type": "string",
                        "description": "Style of the generated image (default: 'vivid').",
                        "enum": ["vivid", "natural"]
                    },
                    "quality": {
                        "type": "string",
                        "description": "Quality of the generated image (default: 'standard').",
                        "enum": ["standard", "hd"]
                    },
                    "image": {
                        "oneOf": [
                            {"type": "string", "description": "URL or Base64 encoded image to use as reference for image generation"},
                            {"type": "array", "items": {"type": "string"}, "description": "Array of URLs or Base64 encoded images to use as references for image generation"}
                        ],
                        "description": "Single image URL/Base64 or array of image URLs/Base64 to use as reference for image generation. Optional parameter."
                    },
                    "sequential_image_generation": {
                        "type": "string",
                        "enum": ["auto", "disabled"],
                        "description": "Enable ('auto') or disable ('disabled') sequential image generation for creating related images. Default: 'disabled'."
                    },
                    "sequential_image_generation_options": {
                        "type": "object",
                        "properties": {
                            "max_images": {
                                "type": "integer",
                                "description": "Maximum number of images to generate in sequence when sequential_image_generation is 'auto'."
                            }
                        },
                        "description": "Options for sequential image generation."
                    },
                    "output_dir": {
                        "type": "string",
                        "description": "Directory to save generated images. This parameter is required."
                    },
                    "target_filename": {
                        "type": "string",
                        "description": "Optional target filename for the generated image. If not provided, a filename will be generated using timestamp and random numbers."
                    }
                },
                "required": ["prompt", "output_dir"]
            }
        ),
        types.Tool(
            name="convert_image_to_base64",
            description="Convert a local image file to Base64 encoded string for use with image generation APIs. Use this when you have a local image file that you want to use as reference for image generation.",
            inputSchema={
                "type": "object",
                "properties": {
                    "image_path": {
                        "type": "string",
                        "description": "Path to the local image file to convert to Base64."
                    }
                },
                "required": ["image_path"]
            }
        )
    ]


@server.call_tool()
async def handle_call_tool(
    name: str, arguments: dict | None
) -> list[types.TextContent | types.EmbeddedResource]:
    """Handle tool calls for configuration and image generation."""
    logger.info(f"Tool called: {name} with args: {arguments}")
    
    try:
        if not arguments:
            raise ValueError("Missing arguments")
        
        if name == "set_config":
            # Update configuration with provided values
            if "api_url" in arguments:
                config["api_url"] = arguments["api_url"]
            if "default_model" in arguments:
                config["default_model"] = arguments["default_model"]
            if "api_key" in arguments:
                config["api_key"] = arguments["api_key"]
            
            result_text = f"Configuration updated successfully.\nCurrent config: {json.dumps({k: v for k, v in config.items() if k != 'api_key'}, indent=2)}"
            return [types.TextContent(type="text", text=result_text)]
        
        elif name == "convert_image_to_base64":
            # Extract the image path
            image_path = arguments.get("image_path")
            if not image_path:
                raise ValueError("Missing image_path argument")
            
            try:
                # Convert the image to Base64
                base64_result = image_to_base64(image_path)
                
                # Return the Base64 string as text content
                return [types.TextContent(
                    type="text",
                    text=f"Base64 conversion successful. Use this string as the 'image' parameter value in generate_image: {base64_result}"
                )]
            except Exception as e:
                logger.error(f"Error converting image to base64: {e}")
                return [types.TextContent(type="text", text=f"Error converting image to base64: {str(e)}")]
        
        elif name == "generate_image":
            # Extract parameters
            prompt = arguments.get("prompt")
            if not prompt:
                raise ValueError("Missing prompt argument")
            
            n = arguments.get("n", 1)
            size = arguments.get("size", "2K")  # Default to 2K for Doubao API
            style = arguments.get("style", "vivid")
            quality = arguments.get("quality", "standard")
            model = arguments.get("model", config["default_model"])
            output_dir = arguments.get("output_dir")
            target_filename = arguments.get("target_filename")
            image = arguments.get("image")
            sequential_image_generation = arguments.get("sequential_image_generation", "disabled")
            sequential_image_generation_options = arguments.get("sequential_image_generation_options")
            
            # Validate parameters
            if n < 1 or n > 10:
                raise ValueError("n must be between 1 and 10")
            
            if not output_dir:
                raise ValueError("Missing output_dir argument")
            
            # Get API configuration
            api_url = config["api_url"]
            api_key = config["api_key"]
            
            if not api_url:
                raise ValueError("API URL not configured. Required: ARK_API_URL environment variable. Use set_config tool first if environment variables are not set.")
            # Ensure the API URL is the supported API since that's the only supported API
            if "ark.cn-beijing.volces.com" not in api_url and "doubao" not in api_url.lower() and "api.openai.com" not in api_url:
                raise ValueError("Only Doubao API or OpenAI API is supported in this version")
            if not api_key:
                raise ValueError("API key not configured. Required: ARK_API_KEY environment variable. Use set_config tool first if environment variables are not set.")
            
            # Prepare the request for a generic image generation API
            headers = {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            }
            
            # Prepare request payload based on the API provider
            if "ark.cn-beijing.volces.com" in api_url or "doubao" in api_url.lower():
                # For Doubao API - using the exact format from the official example
                # Convert size format for Doubao API (e.g., 1024x1024 -> 2K)
                doubao_size = size
                # Only convert standard dimensions to 2K if they're not already in the right format
                if size == "1024x1024":
                    doubao_size = "2K"
                elif size == "256x256" or size == "512x512":
                    # Keep these as they are since they're supported as pixel dimensions
                    doubao_size = size
                # For '1K', '2K', '4K' formats, use them as is
                # For other dimensions like '1792x1024', '1024x1792', they should be sent as-is to the API
                
                payload = {
                    "model": model or "doubao-seedream-4-5-251128",
                    "prompt": prompt,
                    "sequential_image_generation": sequential_image_generation,
                    "response_format": "url",
                    "size": doubao_size,
                    "stream": False,
                    "watermark": False
                }
                
                # Add optional parameters if provided
                if style and style != "vivid":
                    payload["style"] = style
                if quality and quality != "standard":
                    payload["quality"] = quality
                if image:
                    payload["image"] = image
                if sequential_image_generation == "auto" and sequential_image_generation_options:
                    payload["sequential_image_generation_options"] = sequential_image_generation_options
                
                # Use the exact endpoint format from the official example
                api_endpoint = api_url  # Use the full URL as provided in the example
            elif "api.openai.com" in api_url:
                # For OpenAI API
                payload = {
                    "model": model or config["default_model"],
                    "prompt": prompt,
                    "n": n,
                    "size": size,
                }
                
                # Add optional parameters if provided
                if style and style != "vivid":
                    payload["style"] = style
                if quality and quality != "standard":
                    payload["quality"] = quality
                if image:
                    payload["image"] = image
                if sequential_image_generation and sequential_image_generation != "disabled":
                    payload["sequential_image_generation"] = sequential_image_generation
                if sequential_image_generation_options:
                    payload["sequential_image_generation_options"] = sequential_image_generation_options
                
                # For OpenAI, we need to use the images/generations endpoint
                api_endpoint = f"{api_url.rstrip('/')}/images/generations"
            else:
                raise ValueError("Only Doubao API or OpenAI API is supported in this version")
            
            # Make the API request
            response = requests.post(api_endpoint, headers=headers, json=payload)
            response.raise_for_status()
            
            # Debug logging to help troubleshoot API issues
            logger.debug(f"API request sent to {api_endpoint} with payload: {json.dumps(payload, ensure_ascii=False)}")
            
            # Process the response
            response_data = response.json()
            results = []
            
            # Create output directory if it doesn't exist
            output_path = Path(output_dir)
            logger.info(f"Attempting to create/save images in directory: {output_path}")
            try:
                output_path.mkdir(parents=True, exist_ok=True)
            except Exception as e:
                logger.error(f"Failed to create output directory {output_path}: {e}")
                return [types.TextContent(type="text", text=f"Error: Failed to create output directory {output_path}: {str(e)}")]
            
            # Check if the API returns image data directly or as URLs
            if "ark.cn-beijing.volces.com" in api_url or "doubao" in api_url.lower():
                # Doubao API response handling
                if "data" in response_data:
                    for i, item in enumerate(response_data["data"]):
                        if "url" in item:
                            # Handle URL-based response
                            image_url = item["url"]
                            image_response = requests.get(image_url)
                            image_response.raise_for_status()
                            
                            # Save the image to the configured output directory
                            # Use target_filename if provided, otherwise generate using timestamp and random number
                            if target_filename:
                                image_filename = target_filename
                                # If there are multiple images, add a suffix to distinguish them
                                if n > 1:
                                    name_part, ext_part = os.path.splitext(target_filename)
                                    image_filename = f"{name_part}_{i+1}{ext_part}"
                            else:
                                timestamp = int(time.time())
                                random_num = random.randint(1000, 9999)
                                image_filename = f"generated_image_{i+1}_{timestamp}_{random_num}.png"
                            image_path = output_path / image_filename
                            try:
                                with open(image_path, "wb") as f:
                                    f.write(image_response.content)
                                logger.info(f"Image saved to {image_path}")
                            except Exception as e:
                                logger.error(f"Failed to save image to {image_path}: {e}")
                                return [types.TextContent(type="text", text=f"Error: Failed to save image to {image_path}: {str(e)}")]
                            
                            # Include text with the result showing the file path
                            results.append(types.TextContent(
                                type="text",
                                text=f"Image {i+1} generated successfully. File saved: {image_path}"
                            ))
                        elif "b64_json" in item:
                            # Handle base64 encoded image data
                            import base64
                            image_data = base64.b64decode(item["b64_json"])
                            
                            # Save the image to the configured output directory
                            # Use target_filename if provided, otherwise generate using timestamp and random number
                            if target_filename:
                                image_filename = target_filename
                                # If there are multiple images, add a suffix to distinguish them
                                if n > 1:
                                    name_part, ext_part = os.path.splitext(target_filename)
                                    image_filename = f"{name_part}_{i+1}{ext_part}"
                            else:
                                timestamp = int(time.time())
                                random_num = random.randint(1000, 9999)
                                image_filename = f"generated_image_{i+1}_{timestamp}_{random_num}.png"
                            image_path = output_path / image_filename
                            try:
                                with open(image_path, "wb") as f:
                                    f.write(image_data)
                                logger.info(f"Image saved to {image_path}")
                            except Exception as e:
                                logger.error(f"Failed to save image to {image_path}: {e}")
                                return [types.TextContent(type="text", text=f"Error: Failed to save image to {image_path}: {str(e)}")]
                            
                            # Include text with the result showing the file path
                            results.append(types.TextContent(
                                type="text",
                                text=f"Image {i+1} generated successfully. File saved: {image_path}"
                            ))
            elif "api.openai.com" in api_url:
                # Standard OpenAI-style API response
                for i, image_data in enumerate(response_data.get("data", [])):
                    image_url = image_data.get("url")
                    if not image_url:
                        continue
                    
                    # Download the image
                    image_response = requests.get(image_url)
                    image_response.raise_for_status()
                    
                    # Save the image to the configured output directory
                    # Use target_filename if provided, otherwise generate using timestamp and random number
                    if target_filename:
                        image_filename = target_filename
                        # If there are multiple images, add a suffix to distinguish them
                        if n > 1:
                            name_part, ext_part = os.path.splitext(target_filename)
                            image_filename = f"{name_part}_{i+1}{ext_part}"
                    else:
                        timestamp = int(time.time())
                        random_num = random.randint(1000, 9999)
                        image_filename = f"generated_image_{i+1}_{timestamp}_{random_num}.png"
                    image_path = output_path / image_filename
                    try:
                        with open(image_path, "wb") as f:
                            f.write(image_response.content)
                        logger.info(f"Image saved to {image_path}")
                    except Exception as e:
                        logger.error(f"Failed to save image to {image_path}: {e}")
                        return [types.TextContent(type="text", text=f"Error: Failed to save image to {image_path}: {str(e)}")]
                    
                    # Include text with the result showing the file path
                    results.append(types.TextContent(
                        type="text",
                        text=f"Image {i+1} generated successfully. File saved: {image_path}"
                    ))
            
            logger.info(f"Tool {name} execution successful. Generated {len(results)//2} images.")
            return results
        else:
            raise ValueError(f"Unknown tool: {name}")

    except Exception as e:
        logger.error(f"Error executing tool {name}: {e}")
        return [types.TextContent(type="text", text=f"Error: {str(e)}")]


async def run_server():
    """Run the MCP server."""
    logger.info("Starting McpSeedream - Text-to-Image Generation MCP Server...")
    
    # Setup encoding for cross-platform compatibility
    setup_encoding()
    
    # Run the server using stdin/stdout
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="McpSeedream",
                server_version="1.0.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )
