#!/usr/bin/env python3
"""
UN Comtrade MCP Server

This MCP server provides access to UN Comtrade API for international trade data.
"""

import os
import json
from typing import Any, Optional
import comtradeapicall
from mcp.server import Server
from mcp.types import Tool, TextContent
import mcp.server.stdio


# Initialize the MCP server
app = Server("comtrade-mcp-server")


def get_subscription_key() -> Optional[str]:
    """Get subscription key from environment variable."""
    return os.getenv("COMTRADE_SUBSCRIPTION_KEY")


@app.list_tools()
async def list_tools() -> list[Tool]:
    """List available Comtrade API tools."""
    return [
        Tool(
            name="preview_trade_data",
            description="Preview final trade data (limited to 500 records, no subscription key required). "
            "Returns a sample of trade data for exploration.",
            inputSchema={
                "type": "object",
                "properties": {
                    "typeCode": {
                        "type": "string",
                        "description": "Product type: 'C' for Goods or 'S' for Services",
                        "enum": ["C", "S"]
                    },
                    "freqCode": {
                        "type": "string",
                        "description": "Frequency: 'A' for Annual or 'M' for Monthly",
                        "enum": ["A", "M"]
                    },
                    "clCode": {
                        "type": "string",
                        "description": "Classification code (e.g., 'HS', 'SITC')"
                    },
                    "period": {
                        "type": "string",
                        "description": "Time period (YYYY for annual, YYYYMM for monthly, e.g., '2022' or '202205')"
                    },
                    "reporterCode": {
                        "type": "string",
                        "description": "Reporter country code (use get_reference with 'reporter' to get codes)"
                    },
                    "cmdCode": {
                        "type": "string",
                        "description": "Commodity code(s), comma-separated (e.g., '91' or '91,90')"
                    },
                    "flowCode": {
                        "type": "string",
                        "description": "Trade flow code: 'M' for imports, 'X' for exports, etc."
                    },
                    "partnerCode": {
                        "type": "string",
                        "description": "Partner country code (optional)"
                    },
                    "breakdownMode": {
                        "type": "string",
                        "description": "Breakdown mode: 'classic' or 'plus' (optional)",
                        "enum": ["classic", "plus"]
                    },
                    "includeDesc": {
                        "type": "boolean",
                        "description": "Include descriptions (default: true)"
                    }
                },
                "required": ["typeCode", "freqCode", "clCode", "period", "reporterCode", "flowCode"]
            }
        ),
        Tool(
            name="get_trade_data",
            description="Get final trade data with subscription key (limited to 250K records). "
            "Requires COMTRADE_SUBSCRIPTION_KEY environment variable.",
            inputSchema={
                "type": "object",
                "properties": {
                    "typeCode": {
                        "type": "string",
                        "description": "Product type: 'C' for Goods or 'S' for Services",
                        "enum": ["C", "S"]
                    },
                    "freqCode": {
                        "type": "string",
                        "description": "Frequency: 'A' for Annual or 'M' for Monthly",
                        "enum": ["A", "M"]
                    },
                    "clCode": {
                        "type": "string",
                        "description": "Classification code (e.g., 'HS', 'SITC')"
                    },
                    "period": {
                        "type": "string",
                        "description": "Time period (YYYY for annual, YYYYMM for monthly)"
                    },
                    "reporterCode": {
                        "type": "string",
                        "description": "Reporter country code"
                    },
                    "cmdCode": {
                        "type": "string",
                        "description": "Commodity code(s), comma-separated"
                    },
                    "flowCode": {
                        "type": "string",
                        "description": "Trade flow code"
                    },
                    "partnerCode": {
                        "type": "string",
                        "description": "Partner country code (optional)"
                    },
                    "maxRecords": {
                        "type": "integer",
                        "description": "Maximum records to return (default: 2500)"
                    },
                    "breakdownMode": {
                        "type": "string",
                        "description": "Breakdown mode: 'classic' or 'plus'",
                        "enum": ["classic", "plus"]
                    },
                    "includeDesc": {
                        "type": "boolean",
                        "description": "Include descriptions (default: true)"
                    }
                },
                "required": ["typeCode", "freqCode", "clCode", "period", "reporterCode", "flowCode"]
            }
        ),
        Tool(
            name="get_data_availability",
            description="Check what trade data is available for specific criteria. "
            "Free to use without subscription key.",
            inputSchema={
                "type": "object",
                "properties": {
                    "typeCode": {
                        "type": "string",
                        "description": "Product type: 'C' for Goods or 'S' for Services",
                        "enum": ["C", "S"]
                    },
                    "freqCode": {
                        "type": "string",
                        "description": "Frequency: 'A' for Annual or 'M' for Monthly",
                        "enum": ["A", "M"]
                    },
                    "clCode": {
                        "type": "string",
                        "description": "Classification code (e.g., 'HS', 'SITC')"
                    },
                    "period": {
                        "type": "string",
                        "description": "Time period (YYYY for annual, YYYYMM for monthly)"
                    },
                    "reporterCode": {
                        "type": "string",
                        "description": "Reporter country code (optional, use None for all countries)"
                    }
                },
                "required": ["typeCode", "freqCode", "clCode", "period"]
            }
        ),
        Tool(
            name="get_metadata",
            description="Get metadata and publication notes for specific trade data. "
            "Free to use without subscription key.",
            inputSchema={
                "type": "object",
                "properties": {
                    "typeCode": {
                        "type": "string",
                        "description": "Product type: 'C' for Goods or 'S' for Services",
                        "enum": ["C", "S"]
                    },
                    "freqCode": {
                        "type": "string",
                        "description": "Frequency: 'A' for Annual or 'M' for Monthly",
                        "enum": ["A", "M"]
                    },
                    "clCode": {
                        "type": "string",
                        "description": "Classification code (e.g., 'HS', 'SITC')"
                    },
                    "period": {
                        "type": "string",
                        "description": "Time period (YYYY for annual, YYYYMM for monthly)"
                    },
                    "reporterCode": {
                        "type": "string",
                        "description": "Reporter country code (optional)"
                    },
                    "showHistory": {
                        "type": "boolean",
                        "description": "Show historical publication notes (default: false)"
                    }
                },
                "required": ["typeCode", "freqCode", "clCode", "period"]
            }
        ),
        Tool(
            name="list_references",
            description="List all available reference tables (countries, commodities, etc.). "
            "Free to use.",
            inputSchema={
                "type": "object",
                "properties": {
                    "category": {
                        "type": "string",
                        "description": "Filter by category (optional, e.g., 'cmd:HS')"
                    }
                },
                "required": []
            }
        ),
        Tool(
            name="get_reference",
            description="Get contents of a specific reference table (e.g., country codes, commodity codes). "
            "Common categories: 'reporter', 'partner', 'flow', 'customs', 'mot'. Free to use.",
            inputSchema={
                "type": "object",
                "properties": {
                    "category": {
                        "type": "string",
                        "description": "Reference category (e.g., 'reporter', 'partner', 'flow')"
                    }
                },
                "required": ["category"]
            }
        ),
        Tool(
            name="convert_country_iso3_to_code",
            description="Convert ISO3 country codes to Comtrade country codes. "
            "Example: 'USA,FRA,CHE' -> Comtrade codes. Free to use.",
            inputSchema={
                "type": "object",
                "properties": {
                    "iso3_codes": {
                        "type": "string",
                        "description": "Comma-separated ISO3 country codes (e.g., 'USA,FRA,CHE,ITA')"
                    }
                },
                "required": ["iso3_codes"]
            }
        ),
        Tool(
            name="get_live_updates",
            description="Get recent data releases from Comtrade. Requires subscription key.",
            inputSchema={
                "type": "object",
                "properties": {},
                "required": []
            }
        ),
    ]


@app.call_tool()
async def call_tool(name: str, arguments: Any) -> list[TextContent]:
    """Handle tool calls for Comtrade API operations."""

    try:
        if name == "preview_trade_data":
            df = comtradeapicall.previewFinalData(
                typeCode=arguments["typeCode"],
                freqCode=arguments["freqCode"],
                clCode=arguments["clCode"],
                period=arguments["period"],
                reporterCode=arguments["reporterCode"],
                cmdCode=arguments.get("cmdCode"),
                flowCode=arguments["flowCode"],
                partnerCode=arguments.get("partnerCode"),
                partner2Code=None,
                customsCode=None,
                motCode=None,
                maxRecords=arguments.get("maxRecords", 500),
                format_output="JSON",
                aggregateBy=None,
                breakdownMode=arguments.get("breakdownMode"),
                countOnly=None,
                includeDesc=arguments.get("includeDesc", True)
            )

            if df is not None and not df.empty:
                # Convert DataFrame to JSON for better formatting
                result = df.to_json(orient="records", indent=2)
                return [TextContent(
                    type="text",
                    text=f"Found {len(df)} records:\n\n{result}"
                )]
            else:
                return [TextContent(type="text", text="No data found for the specified criteria.")]

        elif name == "get_trade_data":
            subscription_key = get_subscription_key()
            if not subscription_key:
                return [TextContent(
                    type="text",
                    text="Error: COMTRADE_SUBSCRIPTION_KEY environment variable not set. "
                    "This tool requires a valid subscription key."
                )]

            df = comtradeapicall.getFinalData(
                subscription_key=subscription_key,
                typeCode=arguments["typeCode"],
                freqCode=arguments["freqCode"],
                clCode=arguments["clCode"],
                period=arguments["period"],
                reporterCode=arguments["reporterCode"],
                cmdCode=arguments.get("cmdCode"),
                flowCode=arguments["flowCode"],
                partnerCode=arguments.get("partnerCode"),
                partner2Code=None,
                customsCode=None,
                motCode=None,
                maxRecords=arguments.get("maxRecords", 2500),
                format_output="JSON",
                aggregateBy=None,
                breakdownMode=arguments.get("breakdownMode"),
                countOnly=None,
                includeDesc=arguments.get("includeDesc", True)
            )

            if df is not None and not df.empty:
                result = df.to_json(orient="records", indent=2)
                return [TextContent(
                    type="text",
                    text=f"Found {len(df)} records:\n\n{result}"
                )]
            else:
                return [TextContent(type="text", text="No data found for the specified criteria.")]

        elif name == "get_data_availability":
            df = comtradeapicall._getFinalDataAvailability(
                typeCode=arguments["typeCode"],
                freqCode=arguments["freqCode"],
                clCode=arguments["clCode"],
                period=arguments["period"],
                reporterCode=arguments.get("reporterCode")
            )

            if df is not None and not df.empty:
                result = df.to_json(orient="records", indent=2)
                return [TextContent(
                    type="text",
                    text=f"Data availability ({len(df)} entries):\n\n{result}"
                )]
            else:
                return [TextContent(type="text", text="No availability data found.")]

        elif name == "get_metadata":
            df = comtradeapicall._getMetadata(
                typeCode=arguments["typeCode"],
                freqCode=arguments["freqCode"],
                clCode=arguments["clCode"],
                period=arguments["period"],
                reporterCode=arguments.get("reporterCode"),
                showHistory=arguments.get("showHistory", False)
            )

            if df is not None and not df.empty:
                result = df.to_json(orient="records", indent=2)
                return [TextContent(
                    type="text",
                    text=f"Metadata ({len(df)} entries):\n\n{result}"
                )]
            else:
                return [TextContent(type="text", text="No metadata found.")]

        elif name == "list_references":
            df = comtradeapicall.listReference(
                category=arguments.get("category")
            )

            if df is not None and not df.empty:
                result = df.to_json(orient="records", indent=2)
                return [TextContent(
                    type="text",
                    text=f"Available references ({len(df)} entries):\n\n{result}"
                )]
            else:
                return [TextContent(type="text", text="No references found.")]

        elif name == "get_reference":
            df = comtradeapicall.getReference(
                category=arguments["category"]
            )

            if df is not None and not df.empty:
                result = df.to_json(orient="records", indent=2)
                return [TextContent(
                    type="text",
                    text=f"Reference data for '{arguments['category']}' ({len(df)} entries):\n\n{result}"
                )]
            else:
                return [TextContent(type="text", text="No reference data found.")]

        elif name == "convert_country_iso3_to_code":
            code_string = comtradeapicall.convertCountryIso3ToCode(
                countryIsoCode=arguments["iso3_codes"]
            )

            return [TextContent(
                type="text",
                text=f"ISO3 codes '{arguments['iso3_codes']}' -> Comtrade codes: {code_string}"
            )]

        elif name == "get_live_updates":
            subscription_key = get_subscription_key()
            if not subscription_key:
                return [TextContent(
                    type="text",
                    text="Error: COMTRADE_SUBSCRIPTION_KEY environment variable not set."
                )]

            df = comtradeapicall.getLiveUpdate(subscription_key=subscription_key)

            if df is not None and not df.empty:
                result = df.to_json(orient="records", indent=2)
                return [TextContent(
                    type="text",
                    text=f"Recent data releases ({len(df)} entries):\n\n{result}"
                )]
            else:
                return [TextContent(type="text", text="No recent releases found.")]

        else:
            return [TextContent(type="text", text=f"Unknown tool: {name}")]

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


async def main():
    """Run the MCP server using stdio transport."""
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await app.run(
            read_stream,
            write_stream,
            app.create_initialization_options()
        )


if __name__ == "__main__":
    import asyncio
    asyncio.run(main())
