#!/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
import sys
import asyncio
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")

# Constants
MAX_RESPONSE_TOKENS = 20000  # Stay well under 25K limit
DEFAULT_LIMIT = 100
SMALL_CATEGORY_THRESHOLD = 50  # Categories with fewer items returned in full


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


async def check_for_updates():
    """Check PyPI for newer version and log if available."""
    try:
        import urllib.request
        from packaging import version
        from comtrade_mcp import __version__ as current_version

        # Quick check with 2 second timeout
        req = urllib.request.Request(
            "https://pypi.org/pypi/comtrade-mcp-server/json",
            headers={'User-Agent': 'comtrade-mcp-server'}
        )

        with urllib.request.urlopen(req, timeout=2) as response:
            data = json.loads(response.read())
            latest_version = data['info']['version']

            if version.parse(latest_version) > version.parse(current_version):
                print(f"⚠️  Update available: v{current_version} → v{latest_version}", file=sys.stderr)
                print(f"   Run: uv cache clean comtrade-mcp-server && uvx comtrade-mcp-server", file=sys.stderr)
    except Exception:
        # Silently fail - don't block server startup
        pass


def filter_dataframe(df, search: Optional[str] = None, limit: int = DEFAULT_LIMIT, offset: int = 0):
    """Filter and paginate DataFrame with search capability."""
    if df is None or df.empty:
        return df, {"total": 0, "showing": 0, "has_more": False}

    total_records = len(df)
    filtered_df = df

    # Apply search filter if provided
    if search and search.strip():
        search_term = search.strip().lower()
        # Search across all string columns
        mask = filtered_df.astype(str).apply(
            lambda row: row.str.lower().str.contains(search_term, regex=False).any(),
            axis=1
        )
        filtered_df = filtered_df[mask]

    filtered_total = len(filtered_df)

    # Apply pagination
    paginated_df = filtered_df.iloc[offset:offset + limit]

    metadata = {
        "total": total_records,
        "filtered": filtered_total if search else total_records,
        "showing": len(paginated_df),
        "offset": offset,
        "limit": limit,
        "has_more": (offset + limit) < filtered_total
    }

    return paginated_df, metadata


def format_response(df, metadata: dict, category: str = "data") -> str:
    """Format DataFrame with metadata for compact JSON response."""
    if df is None or df.empty:
        return json.dumps({"data": [], "metadata": metadata})

    # Compact JSON (no indentation)
    data_json = df.to_dict(orient="records")

    response = {
        "data": data_json,
        "metadata": metadata
    }

    if metadata.get("has_more"):
        response["hint"] = f"Use offset={metadata['offset'] + metadata['limit']} to get next page"

    return json.dumps(response)


@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 search_reference to find 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. Supports pagination.",
            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)"
                    },
                    "limit": {
                        "type": "integer",
                        "description": f"Max records to return (default: {DEFAULT_LIMIT})"
                    },
                    "offset": {
                        "type": "integer",
                        "description": "Number of records to skip (default: 0)"
                    }
                },
                "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. Supports pagination.",
            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)"
                    },
                    "limit": {
                        "type": "integer",
                        "description": f"Max records to return (default: {DEFAULT_LIMIT})"
                    },
                    "offset": {
                        "type": "integer",
                        "description": "Number of records to skip (default: 0)"
                    }
                },
                "required": ["typeCode", "freqCode", "clCode", "period"]
            }
        ),
        Tool(
            name="list_references",
            description="List all available reference tables (countries, commodities, etc.). "
            "Free to use. Supports pagination.",
            inputSchema={
                "type": "object",
                "properties": {
                    "category": {
                        "type": "string",
                        "description": "Filter by category (optional, e.g., 'cmd:HS')"
                    },
                    "limit": {
                        "type": "integer",
                        "description": f"Max records to return (default: {DEFAULT_LIMIT})"
                    },
                    "offset": {
                        "type": "integer",
                        "description": "Number of records to skip (default: 0)"
                    }
                },
                "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'. "
            "Returns paginated results with search support. Free to use.",
            inputSchema={
                "type": "object",
                "properties": {
                    "category": {
                        "type": "string",
                        "description": "Reference category (e.g., 'reporter', 'partner', 'flow')"
                    },
                    "search": {
                        "type": "string",
                        "description": "Search term to filter results (optional)"
                    },
                    "limit": {
                        "type": "integer",
                        "description": f"Max records to return (default: {DEFAULT_LIMIT}, max: 500)"
                    },
                    "offset": {
                        "type": "integer",
                        "description": "Number of records to skip for pagination (default: 0)"
                    }
                },
                "required": ["category"]
            }
        ),
        Tool(
            name="search_reference",
            description="Search for specific codes in a reference table by keyword. "
            "More efficient than get_reference for finding specific entries. "
            "Returns top 20 matches. Free to use.",
            inputSchema={
                "type": "object",
                "properties": {
                    "category": {
                        "type": "string",
                        "description": "Reference category to search (e.g., 'reporter', 'partner')"
                    },
                    "keyword": {
                        "type": "string",
                        "description": "Keyword to search for (case-insensitive)"
                    }
                },
                "required": ["category", "keyword"]
            }
        ),
        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:
                result = df.to_json(orient="records")
                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")
                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")
            )

            limit = arguments.get("limit", DEFAULT_LIMIT)
            offset = arguments.get("offset", 0)

            filtered_df, metadata = filter_dataframe(df, limit=limit, offset=offset)

            if filtered_df is not None and not filtered_df.empty:
                result = format_response(filtered_df, metadata)
                return [TextContent(type="text", text=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)
            )

            limit = arguments.get("limit", DEFAULT_LIMIT)
            offset = arguments.get("offset", 0)

            filtered_df, metadata = filter_dataframe(df, limit=limit, offset=offset)

            if filtered_df is not None and not filtered_df.empty:
                result = format_response(filtered_df, metadata)
                return [TextContent(type="text", text=result)]
            else:
                return [TextContent(type="text", text="No metadata found.")]

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

            limit = arguments.get("limit", DEFAULT_LIMIT)
            offset = arguments.get("offset", 0)

            filtered_df, metadata = filter_dataframe(df, limit=limit, offset=offset)

            if filtered_df is not None and not filtered_df.empty:
                result = format_response(filtered_df, metadata)
                return [TextContent(type="text", text=result)]
            else:
                return [TextContent(type="text", text="No references found.")]

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

            search = arguments.get("search")
            limit = min(arguments.get("limit", DEFAULT_LIMIT), 500)  # Cap at 500
            offset = arguments.get("offset", 0)

            filtered_df, metadata = filter_dataframe(df, search=search, limit=limit, offset=offset)

            if filtered_df is not None and not filtered_df.empty:
                result = format_response(filtered_df, metadata, category=arguments["category"])
                return [TextContent(type="text", text=result)]
            else:
                if search:
                    return [TextContent(type="text", text=f"No results found for search term: '{search}'")]
                else:
                    return [TextContent(type="text", text="No reference data found.")]

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

            keyword = arguments["keyword"]

            # Search with top 20 limit
            filtered_df, metadata = filter_dataframe(df, search=keyword, limit=20, offset=0)

            if filtered_df is not None and not filtered_df.empty:
                result = format_response(filtered_df, metadata, category=arguments["category"])
                return [TextContent(
                    type="text",
                    text=f"Search results for '{keyword}':\n\n{result}"
                )]
            else:
                return [TextContent(
                    type="text",
                    text=f"No results found for '{keyword}' in category '{arguments['category']}'"
                )]

        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")
                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 run_server():
    """Run the MCP server using stdio transport."""
    # Check for updates on startup
    await check_for_updates()

    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await app.run(
            read_stream,
            write_stream,
            app.create_initialization_options()
        )


def main():
    """Entry point for the MCP server."""
    asyncio.run(run_server())


if __name__ == "__main__":
    main()
