# Import
from mcp.server.fastmcp import FastMCP, Image
from jaxa.earth import je
import requests

# Initialize FastMCP server
mcp = FastMCP("JAXA_Earth_API_Assistant")

# Search collection ID/Bands from JAXA Earth API
@mcp.tool()
async def search_collections_id():
    """
    This guide introduces data available through the JAXA Earth API.
    Based on user requests, please select and respond with the appropriate dataset ID and bands.
    Returned text is a list of all datasets available via the JAXA Earth API.
    Each parameter is described as follows:

    id: The ID uniquely identifying the dataset
    title: The dataset title
    description: The dataset description
    bands: The IDs of the data included in the dataset. Multiple IDs are separated by commas (,).
    keywords: Keywords associated with the dataset. These include the name of the observing satellite or the space agency managing the data.
    startDate: A string representing the start date and time of the dataset period in ISO8601 format. This indicates that data from this date onward is available.
    endDate: A string representing the end date and time of the dataset period in ISO8601 format. A value of “present” indicates that data is still being updated daily.
    The end of each dataset's parameters is indicated by “---”.
    bbox: Geographic extent of the dataset (bounding box). For EPSG:4326, this indicates [minimum longitude, minimum latitude, maximum longitude, maximum latitude]. For EPSG:3995 or EPSG:3031, this indicates [minimum X, minimum Y, maximum X, maximum Y].
    epsg: EPSG code indicating the projection method of the dataset.

    Based on the above, please select the optimal dataset ID and bands for the user's request and provide your response.

    """

    # Load JAXA Earth API datasets information
    JE_TEXT_PATH = "https://data.earth.jaxa.jp/app/mcp/catalog.md"
    response = requests.get(JE_TEXT_PATH)
    je_text  = response.text

    # Return dataset information text
    return je_text

# Show satellite image using JAXA Earth API
@mcp.tool()
async def show_images(
        collection:str = "JAXA.EORC_ALOS.PRISM_AW3D30.v3.2_global",
        band: str      = "DSM",
        dlim: list[str,str                ] = ["2021-01-01T00:00:00" ,
                                               "2021-01-01T00:00:00"], 
        bbox: list[float,float,float,float] = [135.0,37.5,140.0,42.5],
    ) -> list:
    """show satellite image using JAXA Earth API based on user input.

    Args:
        collection: JAXA Earth API collection ID.
        band: band name in the collection.
        dlim: date range limit request from start to end. yyyy-mm-ddThh:mm:ss formatted date string.
        bbox: Geographic extent of the dataset (bounding box). For EPSG:4326, this indicates [minimum longitude, minimum latitude, maximum longitude, maximum latitude]. For EPSG:3995 or EPSG:3031, this indicates [minimum X, minimum Y, maximum X, maximum Y].
    """

    # Set ppu (pixel per unit) by target image size
    image_size = 300 # image size target
    ppu        = image_size/(bbox[2] - bbox[0])

    # Get an image
    data = je.ImageCollection(collection=collection,ssl_verify=True)\
             .filter_date(dlim=dlim)\
             .filter_resolution(ppu=ppu)\
             .filter_bounds(bbox=bbox)\
             .select(band=band)\
             .get_images()

    # Process and show an image
    img_data = je.ImageProcess(data)\
                 .show_images(output="buffer")
    
    # Return image data as png buffers
    output = []
    for png_buffer in img_data.png_buffers:
        output.append(Image(data=png_buffer, format="png"))
    return output

# Calclate satellite data's spatial statistics using JAXA Earth API
@mcp.tool()
async def calc_spatial_stats(
        collection:str = "JAXA.EORC_ALOS.PRISM_AW3D30.v3.2_global",
        band: str      = "DSM",
        dlim: list[str,str                ] = ["2021-01-01T00:00:00" ,
                                               "2021-01-01T00:00:00"], 
        bbox: list[float,float,float,float] = [135.0,37.5,140.0,42.5],
    ) -> dict:
    """Calcurate satellite data's spatial statistics values using JAXA Earth API based on user input.

    Args:
        collection: JAXA Earth API collection ID.
        band: band name in the collection.
        dlim: date range limit request from start to end. yyyy-mm-ddThh:mm:ss formatted date string.
        bbox: Geographic extent of the dataset (bounding box). For EPSG:4326, this indicates [minimum longitude, minimum latitude, maximum longitude, maximum latitude]. For EPSG:3995 or EPSG:3031, this indicates [minimum X, minimum Y, maximum X, maximum Y].
    """

    # Set ppu (pixel per unit) by target image size
    image_size = 300 # image size target
    ppu        = image_size/(bbox[2] - bbox[0])

    # Get an image
    data = je.ImageCollection(collection=collection,ssl_verify=True)\
             .filter_date(dlim=dlim)\
             .filter_resolution(ppu=ppu)\
             .filter_bounds(bbox=bbox)\
             .select(band=band)\
             .get_images()

    # Process and show an image
    img_data = je.ImageProcess(data)\
                 .calc_spatial_stats()\
    
    # Return
    return img_data.timeseries

# Show satellite data's spatial statistics using JAXA Earth API
@mcp.tool()
async def show_spatial_stats(
        collection:str = "JAXA.EORC_ALOS.PRISM_AW3D30.v3.2_global",
        band: str      = "DSM",
        dlim: list[str,str                ] = ["2021-01-01T00:00:00" ,
                                               "2021-01-01T00:00:00"], 
        bbox: list[float,float,float,float] = [135.0,37.5,140.0,42.5],
    ) -> dict:
    """Show satellite data's spatial statistics result image using JAXA Earth API based on user input.

    Args:
        collection: JAXA Earth API collection ID.
        band: band name in the collection.
        dlim: date range limit request from start to end. yyyy-mm-ddThh:mm:ss formatted date string.
        bbox: Geographic extent of the dataset (bounding box). For EPSG:4326, this indicates [minimum longitude, minimum latitude, maximum longitude, maximum latitude]. For EPSG:3995 or EPSG:3031, this indicates [minimum X, minimum Y, maximum X, maximum Y].
    """

    # Set ppu (pixel per unit) by target image size
    image_size = 300 # image size target
    ppu        = image_size/(bbox[2] - bbox[0])

    # Get an image
    data = je.ImageCollection(collection=collection,ssl_verify=True)\
             .filter_date(dlim=dlim)\
             .filter_resolution(ppu=ppu)\
             .filter_bounds(bbox=bbox)\
             .select(band=band)\
             .get_images()

    # Process and show an image
    img_data = je.ImageProcess(data)\
                 .calc_spatial_stats()\
                 .show_spatial_stats(output="buffer")
    
    # Return image data as png buffers
    output_images = []
    for png_buffer in img_data.png_buffers_stats:
        output_images.append(Image(data=png_buffer, format="png"))

    # Return
    return output_images

# Run the MCP server
def main():
    mcp.run(transport='stdio')

if __name__ == "__main__":
    main()