#!/usr/bin/env python3
"""
MCP Server for Azure Maps.

This server provides tools to interact with the Azure Maps REST Services,
specifically for retrieving static maps.
"""

import os
import httpx
import base64
from typing import Optional, List, Literal
from pydantic import BaseModel, Field, ConfigDict
from mcp.server.fastmcp import FastMCP, Context

# Initialize the MCP server
mcp = FastMCP("azure_maps_mcp")

# Constants
API_BASE_URL = "https://atlas.microsoft.com/map/static"
API_VERSION = "2024-04-01"
DEFAULT_MAP_SIZE = "500,500"
DEFAULT_FORMAT = "png"

# Imagery Sets Mapping
# Mapping Bing styles to Azure (layer, style) tuples or just layer
# Azure layers: basic, hybrid, labels, terra
# Azure styles: main, dark
IMAGERY_MAP = {
    "Aerial": ("hybrid", "main"), # detailed satellite/hybrid
    "AerialWithLabels": ("hybrid", "main"),
    "AerialWithLabelsOnDemand": ("hybrid", "main"),
    "Birdseye": ("hybrid", "main"), # Best approximation
    "BirdseyeWithLabels": ("hybrid", "main"),
    "BirdseyeV2": ("hybrid", "main"),
    "BirdseyeV2WithLabels": ("hybrid", "main"),
    "CanvasDark": ("basic", "dark"),
    "CanvasLight": ("basic", "main"), # lighter main
    "CanvasGray": ("basic", "main"), # Best approximation
    "Road": ("basic", "main"),
    "RoadOnDemand": ("basic", "main")
}

ImagerySet = Literal[
    "Aerial",
    "AerialWithLabels",
    "AerialWithLabelsOnDemand",
    "Birdseye",
    "BirdseyeWithLabels",
    "BirdseyeV2",
    "BirdseyeV2WithLabels",
    "CanvasDark",
    "CanvasLight",
    "CanvasGray",
    "Road",
    "RoadOnDemand"
]

class GetStaticMapInput(BaseModel):
    """Input parameters for getting a static map."""
    model_config = ConfigDict(extra="forbid")

    centerPoint: str = Field(
        ...,
        description="Center point of the map as 'latitude,longitude' (e.g., '47.610,-122.107')."
    )
    zoomLevel: int = Field(
        ...,
        description="Zoom level of the map (0-20).",
        ge=0,
        le=20
    )
    imagerySet: ImagerySet = Field(
        default="Aerial",
        description="The type of imagery to display (mapped to Azure Maps layers)."
    )
    mapSize: str = Field(
        default=DEFAULT_MAP_SIZE,
        description="Width and height of the map in pixels, formatted as 'width,height' (e.g., '500,500')."
    )
    format: Literal["png", "jpg", "jpeg"] = Field(
        default=DEFAULT_FORMAT,
        description="Image format of the returned map."
    )
    pushpins: Optional[List[str]] = Field(
        default=None,
        description="List of pushpins. Format from Bing 'style;location;title' will be converted. Location must be 'lat,lon'."
    )
    mapLayer: Optional[Literal["Traffic", "Transit"]] = Field(
        default=None,
        description="Optional layer. 'Traffic' adds traffic flow. 'Transit' is not directly supported in static image."
    )

async def _get_api_key(ctx: Context) -> str:
    """Retrieves the API key from environment or context."""
    api_key = os.environ.get("AZURE_MAPS_SUBSCRIPTION_KEY")
    if not api_key:
        # Try to elicit from user if supported by client, otherwise fail
        try:
            api_key = await ctx.elicit("Please provide your Azure Maps Subscription Key:", input_type="password")
        except Exception:
            pass
            
    if not api_key:
         raise ValueError("AZURE_MAPS_SUBSCRIPTION_KEY environment variable is not set.")
    return api_key

@mcp.tool()
async def get_static_map(params: GetStaticMapInput, ctx: Context) -> dict:
    """
    Retrieves a static map image from Azure Maps.
    
    Returns a base64 encoded string of the image, wrapped in a JSON object 
    describing the image.
    """
    api_key = await _get_api_key(ctx)
    
    # Parse center point (Lat,Lon -> Lon,Lat)
    try:
        lat, lon = params.centerPoint.split(",")
        center = f"{lon.strip()},{lat.strip()}"
    except ValueError:
        return {"error": "centerPoint must be 'latitude,longitude'"}

    # Parse map size
    try:
        width, height = params.mapSize.split(",")
    except ValueError:
        width, height = "500", "500"

    # Map imagery set
    layer, style = IMAGERY_MAP.get(params.imagerySet, ("basic", "main"))

    query_params = {
        "api-version": API_VERSION,
        "subscription-key": api_key,
        "center": center,
        "zoom": params.zoomLevel,
        "width": width,
        "height": height,
        "format": params.format,
        "layer": layer,
        "style": style
    }
    
    # Handle Traffic (Azure Maps Static Image doesn't have simple 'Traffic' layer param usually, 
    # but let's check if we can append it to layer or use traffic flow tiles.
    # Standard static API: layer=basic, hybrid, labels. 
    # To keep it simple and safe, we will ignore Transit. 
    # If Traffic is requested, we might not be able to fulfill it easily in one request without 'traffic' layer support 
    # which varies by API version. We'll omit for now to ensure stability or use 'hybrid' which might show some?)
    # Actually, some docs mention 'traffic' param boolean? No.
    # We will log or ignore for now as it's a migration.

    # Handle pushpins
    # Bing: style;location;title -> Azure: pins=default||'title'|lon lat
    if params.pushpins:
        # Azure allows multiple 'pins' parameters
        # Format: style|label|lon lat (Simplified)
        # Reference: pins=default||'Label'|12.34 56.78
        
        # We need to construct the 'pins' param string. 
        # Since httpx can take a list of tuples for multiple params with same name.
        
        # Let's prepare the list of query params
        api_params = list(query_params.items())
        
        for pin in params.pushpins:
            try:
                parts = pin.split(";")
                # Expected: style;lat,lon;title or style;lat,lon
                p_style = "default" # We map everything to default for now
                p_loc = ""
                p_title = ""
                
                if len(parts) >= 2:
                    p_loc_raw = parts[1]
                    if len(parts) >= 3:
                        p_title = parts[2]
                else:
                    continue # Invalid format
                
                # Parse lat,lon -> lon lat
                p_lat, p_lon = p_loc_raw.split(",")
                azure_loc = f"{p_lon.strip()} {p_lat.strip()}"
                
                # Construct pin param
                # default|co<color>|la<label_anchor>|lc<label_color>||'Label'|lon lat
                # Simple: default||'Label'|lon lat
                
                pin_str = f"default||'{p_title}'|{azure_loc}" if p_title else f"default|||{azure_loc}"
                api_params.append(("pins", pin_str))
                
            except Exception:
                continue # Skip malformed pins
    else:
        api_params = list(query_params.items())

    async with httpx.AsyncClient() as client:
        response = await client.get(API_BASE_URL, params=api_params)
        
        if response.status_code != 200:
            try:
                error_data = response.json()
                error_msg = error_data.get("error", {}).get("message", response.text)
                return {"error": f"{response.status_code} - {error_msg}"}
            except:
                return {"error": f"{response.status_code} - {response.text}"}
                
        # Success - return base64 image
        image_data = response.content
        base64_img = base64.b64encode(image_data).decode('utf-8')
        
        return {
            "format": params.format,
            "width": width,
            "height": height,
            "description": f"Static map of {params.centerPoint} at zoom {params.zoomLevel}",
            "image_data_base64": base64_img
        }

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