#!/usr/bin/env python3
"""
Alpha Vantage MCP Server
Provide股票市场DataQuery功能
"""

import os
import json
import asyncio
import aiohttp
from typing import Dict, List, Any, Optional
from datetime import datetime
import logging

from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationOptions
import mcp.server.stdio
import mcp.types as types

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

# 创建服务器实例
server = Server("alpha-vantage")

# Alpha Vantage API 配置
API_KEY = os.getenv("ALPHA_VANTAGE_API_KEY", "")
BASE_URL = "https://www.alphavantage.co/query"

@server.list_tools()
async def handle_list_tools() -> List[types.Tool]:
    """List all available tools"""
    return [
        types.Tool(
            name="get_quote",
            description="Get real-time stock quote",
            inputSchema={
                "type": "object",
                "properties": {
                    "symbol": {
                        "type": "string",
                        "description": "Stock symbol (e.g., AAPL)"
                    }
                },
                "required": ["symbol"]
            }
        ),
        types.Tool(
            name="get_daily",
            description="Get daily time series stock prices",
            inputSchema={
                "type": "object",
                "properties": {
                    "symbol": {
                        "type": "string",
                        "description": "Stock symbol"
                    },
                    "outputsize": {
                        "type": "string",
                        "description": "Number of data points: 'compact' (100) or 'full' (20+ years)",
                        "enum": ["compact", "full"],
                        "default": "compact"
                    }
                },
                "required": ["symbol"]
            }
        ),
        types.Tool(
            name="get_company_overview",
            description="Get company fundamental data and key metrics",
            inputSchema={
                "type": "object",
                "properties": {
                    "symbol": {
                        "type": "string",
                        "description": "Stock symbol"
                    }
                },
                "required": ["symbol"]
            }
        ),
        types.Tool(
            name="get_technical_indicator",
            description="Get technical indicators (SMA, EMA, RSI, MACD)",
            inputSchema={
                "type": "object",
                "properties": {
                    "symbol": {
                        "type": "string",
                        "description": "Stock symbol"
                    },
                    "indicator": {
                        "type": "string",
                        "description": "Indicator type",
                        "enum": ["SMA", "EMA", "RSI", "MACD"]
                    },
                    "interval": {
                        "type": "string",
                        "description": "Time interval",
                        "enum": ["1min", "5min", "15min", "30min", "60min", "daily", "weekly", "monthly"],
                        "default": "daily"
                    },
                    "time_period": {
                        "type": "integer",
                        "description": "Number of data points for calculation",
                        "default": 20
                    }
                },
                "required": ["symbol", "indicator"]
            }
        ),
        types.Tool(
            name="search_symbol",
            description="Search for stock symbols by company name or keyword",
            inputSchema={
                "type": "object",
                "properties": {
                    "keywords": {
                        "type": "string",
                        "description": "Company name or keywords"
                    }
                },
                "required": ["keywords"]
            }
        )
    ]
    
async def make_request(params: Dict[str, Any]) -> Dict[str, Any]:
    """Send request到 Alpha Vantage API"""
    if not API_KEY:
        raise ValueError("ALPHA_VANTAGE_API_KEY environment variable not set")
    
    # 添加 API key
    params["apikey"] = API_KEY
    
    async with aiohttp.ClientSession() as session:
        try:
            async with session.get(BASE_URL, params=params, timeout=30) as response:
                data = await response.json()
                
                # 检查Error
                if "Error Message" in data:
                    raise ValueError(f"API Error: {data['Error Message']}")
                elif "Note" in data:
                    raise ValueError(f"API Limit: {data['Note']}")
                elif not data:
                    raise ValueError("Empty response from API")
                    
                return data
                    
        except asyncio.TimeoutError:
            raise ValueError("Request timeout")
        except Exception as e:
            logger.error(f"Request error: {e}")
            raise

@server.call_tool()
async def handle_call_tool(
    name: str, arguments: Dict[str, Any]
) -> List[types.TextContent]:
    """Handle/Process工具调用"""
    try:
        if name == "get_quote":
            params = {
                "function": "GLOBAL_QUOTE",
                "symbol": arguments["symbol"]
            }
            data = await make_request(params)
            
            if "Global Quote" in data:
                quote = data["Global Quote"]
                result = f"""Stock Quote for {arguments['symbol']}:
- Price: ${quote.get('05. price', 'N/A')}
- Open: ${quote.get('02. open', 'N/A')}
- High: ${quote.get('03. high', 'N/A')}
- Low: ${quote.get('04. low', 'N/A')}
- Volume: {quote.get('06. volume', 'N/A')}
- Previous Close: ${quote.get('08. previous close', 'N/A')}
- Change: {quote.get('09. change', 'N/A')}
- Change %: {quote.get('10. change percent', 'N/A')}"""
            else:
                result = f"No quote data found for {arguments['symbol']}"
                
        elif name == "get_daily":
            params = {
                "function": "TIME_SERIES_DAILY",
                "symbol": arguments["symbol"],
                "outputsize": arguments.get("outputsize", "compact")
            }
            data = await make_request(params)
            
            if "Time Series (Daily)" in data:
                time_series = data["Time Series (Daily)"]
                dates = sorted(list(time_series.keys()), reverse=True)[:5]
                
                result = f"Daily Prices for {arguments['symbol']} (Last 5 days):\n"
                for date in dates:
                    day_data = time_series[date]
                    result += f"\n{date}:"
                    result += f"\n  Open: ${day_data['1. open']}"
                    result += f"\n  High: ${day_data['2. high']}"
                    result += f"\n  Low: ${day_data['3. low']}"
                    result += f"\n  Close: ${day_data['4. close']}"
                    result += f"\n  Volume: {day_data['5. volume']}"
            else:
                result = f"No daily price data found for {arguments['symbol']}"
                
        elif name == "get_company_overview":
            params = {
                "function": "OVERVIEW",
                "symbol": arguments["symbol"]
            }
            data = await make_request(params)
            
            if data and "Symbol" in data:
                result = f"""Company Overview for {data['Symbol']}:
- Name: {data.get('Name', 'N/A')}
- Description: {data.get('Description', 'N/A')[:200]}...
- Sector: {data.get('Sector', 'N/A')}
- Industry: {data.get('Industry', 'N/A')}
- Market Cap: ${data.get('MarketCapitalization', 'N/A')}
- P/E Ratio: {data.get('PERatio', 'N/A')}
- Dividend Yield: {data.get('DividendYield', 'N/A')}%
- 52 Week High: ${data.get('52WeekHigh', 'N/A')}
- 52 Week Low: ${data.get('52WeekLow', 'N/A')}"""
            else:
                result = f"No company overview found for {arguments['symbol']}"
                
        elif name == "get_technical_indicator":
            function_map = {
                "SMA": "SMA",
                "EMA": "EMA", 
                "RSI": "RSI",
                "MACD": "MACD"
            }
            
            params = {
                "function": function_map[arguments["indicator"]],
                "symbol": arguments["symbol"],
                "interval": arguments.get("interval", "daily"),
                "time_period": arguments.get("time_period", 20),
                "series_type": "close"
            }
            data = await make_request(params)
            
            # Get/Fetch技术指标Data的键名
            tech_key = f"Technical Analysis: {arguments['indicator']}"
            if tech_key in data:
                tech_data = data[tech_key]
                dates = sorted(list(tech_data.keys()), reverse=True)[:5]
                
                result = f"{arguments['indicator']} for {arguments['symbol']} (Last 5 periods):\n"
                for date in dates:
                    result += f"\n{date}: {tech_data[date][arguments['indicator']]}"
            else:
                result = f"No {arguments['indicator']} data found for {arguments['symbol']}"
                
        elif name == "search_symbol":
            params = {
                "function": "SYMBOL_SEARCH",
                "keywords": arguments["keywords"]
            }
            data = await make_request(params)
            
            if "bestMatches" in data:
                matches = data["bestMatches"][:5]
                result = f"Symbol Search Results for '{arguments['keywords']}':\n"
                for match in matches:
                    result += f"\n- {match['1. symbol']}: {match['2. name']}"
                    result += f"\n  Type: {match['3. type']}, Region: {match['4. region']}"
            else:
                result = f"No symbols found for '{arguments['keywords']}'"
                
        else:
            raise ValueError(f"Unknown tool: {name}")
            
        return [types.TextContent(type="text", text=result)]
            
    except Exception as e:
        error_msg = f"Error calling {name}: {str(e)}"
        logger.error(error_msg)
        return [types.TextContent(type="text", text=error_msg)]
    
async def main():
    # Run server
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="alpha-vantage",
                server_version="0.1.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                )
            )
        )

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