#!/usr/bin/env python3
"""
Flight Tracker MCP Server
Uses AviationStack API (has free tier)
"""
import os
import sys
import json
import asyncio
import logging
from typing import Any, Dict, List, Optional
from datetime import datetime
import aiohttp
from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationOptions
import mcp.server.stdio
import mcp.types as types

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class FlightTrackerServer:
    def __init__(self):
        # AviationStack API (requires registration for free API key)
        # API key read from environment variable
        self.api_key = os.getenv("AVIATIONSTACK_API_KEY", "")
        self.base_url = "http://api.aviationstack.com/v1"
        self.server = Server("flight-tracker")
        
        # Setup tool list handler
        @self.server.list_tools()
        async def handle_list_tools() -> List[types.Tool]:
            return [
                self._track_flight_tool(),
                self._search_flights_tool()
            ]
        
        # Setup tool call handler
        @self.server.call_tool()
        async def handle_call_tool(name: str, arguments: Optional[Dict[str, Any]]) -> List[types.TextContent]:
            return await self._handle_tool_call(name, arguments or {})
    
    def _track_flight_tool(self) -> types.Tool:
        return types.Tool(
            name="track_flight",
            description="Get real-time status of a specific flight",
            inputSchema={
                "type": "object",
                "properties": {
                    "flight_number": {
                        "type": "string",
                        "description": "Flight number (e.g., AA100, UA456)"
                    }
                },
                "required": ["flight_number"]
            }
        )
    
    def _search_flights_tool(self) -> types.Tool:
        return types.Tool(
            name="search_flights",
            description="Search for flights between airports",
            inputSchema={
                "type": "object",
                "properties": {
                    "origin": {
                        "type": "string",
                        "description": "Origin airport code (e.g., JFK, LAX)"
                    },
                    "destination": {
                        "type": "string",
                        "description": "Destination airport code (e.g., LHR, CDG)"
                    },
                    "date": {
                        "type": "string",
                        "description": "Flight date in YYYY-MM-DD format"
                    }
                },
                "required": ["origin", "destination", "date"]
            }
        )
    
    async def _handle_tool_call(self, name: str, arguments: Dict[str, Any]) -> List[types.TextContent]:
        """Handle tool calls"""
        try:
            if not self.api_key:
                return [types.TextContent(
                    type="text",
                    text="Error: AVIATIONSTACK_API_KEY not set. Please set it in environment variables."
                )]
            
            if name == "track_flight":
                result = await self._track_flight(arguments["flight_number"])
            elif name == "search_flights":
                result = await self._search_flights(
                    arguments["origin"],
                    arguments["destination"],
                    arguments["date"]
                )
            else:
                return [types.TextContent(
                    type="text",
                    text=f"Unknown tool: {name}"
                )]
            
            return [types.TextContent(
                type="text",
                text=json.dumps(result, indent=2)
            )]
            
        except Exception as e:
            logger.error(f"Error in {name}: {e}")
            return [types.TextContent(
                type="text",
                text=f"Error: {str(e)}"
            )]
    
    async def _track_flight(self, flight_number: str) -> Dict[str, Any]:
        """Track a specific flight"""
        try:
            async with aiohttp.ClientSession() as session:
                params = {
                    'access_key': self.api_key,
                    'flight_iata': flight_number
                }
                async with session.get(f"{self.base_url}/flights", params=params) as response:
                    if response.status == 200:
                        data = await response.json()
                        if data.get('data'):
                            return {
                                "flight_number": flight_number,
                                "status": "found",
                                "flights": data['data']
                            }
                        else:
                            return {
                                "flight_number": flight_number,
                                "status": "not_found",
                                "message": "No flight data found"
                            }
                    else:
                        return {
                            "flight_number": flight_number,
                            "status": "error",
                            "message": f"API error: {response.status}"
                        }
        except Exception as e:
            return {
                "flight_number": flight_number,
                "status": "error",
                "message": f"Error: {str(e)}"
            }
    
    async def _search_flights(self, origin: str, destination: str, date: str) -> Dict[str, Any]:
        """Search flights between airports"""
        try:
            async with aiohttp.ClientSession() as session:
                params = {
                    'access_key': self.api_key,
                    'dep_iata': origin,
                    'arr_iata': destination,
                    'flight_date': date
                }
                async with session.get(f"{self.base_url}/flights", params=params) as response:
                    if response.status == 200:
                        data = await response.json()
                        flights = data.get('data', [])
                        return {
                            "origin": origin,
                            "destination": destination,
                            "date": date,
                            "flights": flights,
                            "count": len(flights)
                        }
                    else:
                        error_text = await response.text()
                        return {
                            "origin": origin,
                            "destination": destination,
                            "date": date,
                            "flights": [],
                            "error": f"API error {response.status}: {error_text}"
                        }
        except Exception as e:
            return {
                "origin": origin,
                "destination": destination,
                "date": date,
                "flights": [],
                "error": f"Error: {str(e)}"
            }
    
    async def run(self):
        """Run the server"""
        async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
            await self.server.run(
                read_stream,
                write_stream,
                InitializationOptions(
                    server_name="flight-tracker",
                    server_version="0.1.0",
                    capabilities=self.server.get_capabilities(
                        notification_options=NotificationOptions(),
                        experimental_capabilities={},
                    ),
                ),
            )

async def main():
    """Main function"""
    server = FlightTrackerServer()
    await server.run()

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

