#!/usr/bin/env python3
"""
QR Generator MCP Server - Use/Usinggoqr.meFreeAPI生成二维码
"""
import os
import sys
import json
import asyncio
import logging
from typing import Any, Dict, List, Optional
import base64
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 QRGeneratorServer:
    def __init__(self):
        # Use/UsingFree的 goqr.me API
        self.base_url = "https://api.qrserver.com/v1"
        self.server = Server("qr-generator")
        
        # 设置工具ListHandle/Process器
        @self.server.list_tools()
        async def handle_list_tools() -> List[types.Tool]:
            return [
                self._generate_qr_code_tool(),
                self._generate_qr_with_logo_tool(),
                self._generate_vcard_qr_tool(),
                self._generate_wifi_qr_tool(),
                self._generate_sms_qr_tool()
            ]
        
        # 设置工具调用Handle/Process器
        @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 _generate_qr_code_tool(self) -> types.Tool:
        return types.Tool(
            name="generate_qr_code",
            description="Generate a QR code for any text or URL",
            inputSchema={
                "type": "object",
                "properties": {
                    "data": {
                        "type": "string",
                        "description": "Text or URL to encode in QR code"
                    },
                    "size": {
                        "type": "integer",
                        "description": "Size of QR code in pixels (default: 200)",
                        "default": 200,
                        "minimum": 100,
                        "maximum": 1000
                    },
                    "color": {
                        "type": "string",
                        "description": "Color of QR code in hex format (default: 000000)",
                        "default": "000000"
                    },
                    "bgcolor": {
                        "type": "string",
                        "description": "Background color in hex format (default: FFFFFF)",
                        "default": "FFFFFF"
                    },
                    "format": {
                        "type": "string",
                        "enum": ["png", "jpg", "svg"],
                        "description": "Output format (default: png)",
                        "default": "png"
                    }
                },
                "required": ["data"]
            }
        )
    
    def _generate_qr_with_logo_tool(self) -> types.Tool:
        return types.Tool(
            name="generate_qr_with_logo",
            description="Generate a QR code with a logo in the center",
            inputSchema={
                "type": "object",
                "properties": {
                    "data": {
                        "type": "string",
                        "description": "Text or URL to encode"
                    },
                    "logo_url": {
                        "type": "string",
                        "description": "URL of the logo image"
                    },
                    "size": {
                        "type": "integer",
                        "default": 300
                    }
                },
                "required": ["data", "logo_url"]
            }
        )
    
    def _generate_vcard_qr_tool(self) -> types.Tool:
        return types.Tool(
            name="generate_vcard_qr",
            description="Generate a QR code for a contact vCard",
            inputSchema={
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string",
                        "description": "Full name"
                    },
                    "phone": {
                        "type": "string",
                        "description": "Phone number"
                    },
                    "email": {
                        "type": "string",
                        "description": "Email address"
                    },
                    "organization": {
                        "type": "string",
                        "description": "Company/Organization name"
                    },
                    "url": {
                        "type": "string",
                        "description": "Website URL"
                    }
                },
                "required": ["name"]
            }
        )
    
    def _generate_wifi_qr_tool(self) -> types.Tool:
        return types.Tool(
            name="generate_wifi_qr",
            description="Generate a QR code for WiFi network connection",
            inputSchema={
                "type": "object",
                "properties": {
                    "ssid": {
                        "type": "string",
                        "description": "WiFi network name (SSID)"
                    },
                    "password": {
                        "type": "string",
                        "description": "WiFi password"
                    },
                    "security": {
                        "type": "string",
                        "enum": ["WPA", "WEP", "nopass"],
                        "description": "Security type (default: WPA)",
                        "default": "WPA"
                    },
                    "hidden": {
                        "type": "boolean",
                        "description": "Is network hidden? (default: false)",
                        "default": False
                    }
                },
                "required": ["ssid"]
            }
        )
    
    def _generate_sms_qr_tool(self) -> types.Tool:
        return types.Tool(
            name="generate_sms_qr",
            description="Generate a QR code for sending SMS",
            inputSchema={
                "type": "object",
                "properties": {
                    "phone": {
                        "type": "string",
                        "description": "Phone number to send SMS to"
                    },
                    "message": {
                        "type": "string",
                        "description": "SMS message content"
                    }
                },
                "required": ["phone"]
            }
        )
    
    async def _generate_qr(self, data: str, **params) -> Dict[str, Any]:
        """Generate QR code using goqr.me API"""
        url = f"{self.base_url}/create-qr-code/"
        
        # Default parameters
        default_params = {
            "data": data,
            "size": "200x200",
            "format": "png"
        }
        
        # Update with provided parameters
        if "size" in params:
            size = params["size"]
            default_params["size"] = f"{size}x{size}"
            del params["size"]
        
        default_params.update(params)
        
        async with aiohttp.ClientSession() as session:
            try:
                async with session.get(url, params=default_params, timeout=30) as response:
                    if response.status == 200:
                        # Get image data
                        image_data = await response.read()
                        # Convert to base64
                        base64_data = base64.b64encode(image_data).decode('utf-8')
                        
                        return {
                            "success": True,
                            "format": default_params.get("format", "png"),
                            "size": default_params.get("size", "200x200"),
                            "data_url": f"data:image/{default_params.get('format', 'png')};base64,{base64_data}",
                            "base64": base64_data
                        }
                    else:
                        return {"error": f"API returned status {response.status}"}
                        
            except asyncio.TimeoutError:
                return {"error": "Request timeout"}
            except Exception as e:
                return {"error": f"Request failed: {str(e)}"}
    
    async def _handle_tool_call(self, name: str, arguments: Dict[str, Any]) -> List[types.TextContent]:
        """Handle tool calls"""
        try:
            if name == "generate_qr_code":
                data = arguments["data"]
                params = {}
                
                if "color" in arguments:
                    params["color"] = arguments["color"]
                if "bgcolor" in arguments:
                    params["bgcolor"] = arguments["bgcolor"]
                if "format" in arguments:
                    params["format"] = arguments["format"]
                if "size" in arguments:
                    params["size"] = arguments["size"]
                
                result = await self._generate_qr(data, **params)
                
            elif name == "generate_qr_with_logo":
                # Note: goqr.me doesn't support logos directly
                # We'll generate a regular QR code and note this limitation
                data = arguments["data"]
                size = arguments.get("size", 300)
                
                result = await self._generate_qr(data, size=size)
                if result.get("success"):
                    result["note"] = "Logo embedding not supported by free API. Generated regular QR code."
                
            elif name == "generate_vcard_qr":
                # Build vCard format
                vcard_lines = ["BEGIN:VCARD", "VERSION:3.0"]
                
                if "name" in arguments:
                    vcard_lines.append(f"FN:{arguments['name']}")
                if "phone" in arguments:
                    vcard_lines.append(f"TEL:{arguments['phone']}")
                if "email" in arguments:
                    vcard_lines.append(f"EMAIL:{arguments['email']}")
                if "organization" in arguments:
                    vcard_lines.append(f"ORG:{arguments['organization']}")
                if "url" in arguments:
                    vcard_lines.append(f"URL:{arguments['url']}")
                
                vcard_lines.append("END:VCARD")
                vcard_data = "\n".join(vcard_lines)
                
                result = await self._generate_qr(vcard_data, size=250)
                
            elif name == "generate_wifi_qr":
                # Build WiFi format
                ssid = arguments["ssid"]
                password = arguments.get("password", "")
                security = arguments.get("security", "WPA")
                hidden = arguments.get("hidden", False)
                
                # WiFi QR format: WIFI:T:WPA;S:mynetwork;P:mypass;H:true;;
                wifi_data = f"WIFI:T:{security};S:{ssid};P:{password};H:{'true' if hidden else 'false'};;"
                
                result = await self._generate_qr(wifi_data, size=250)
                
            elif name == "generate_sms_qr":
                # Build SMS format
                phone = arguments["phone"]
                message = arguments.get("message", "")
                
                # SMS QR format: smsto:+1234567890:Hello
                sms_data = f"smsto:{phone}:{message}" if message else f"smsto:{phone}"
                
                result = await self._generate_qr(sms_data)
                
            else:
                result = {"error": f"Unknown tool: {name}"}
            
            return [types.TextContent(
                type="text",
                text=json.dumps(result, indent=2)
            )]
            
        except Exception as e:
            logger.error(f"Error in tool {name}: {e}")
            return [types.TextContent(
                type="text",
                text=json.dumps({"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="qr-generator",
                    server_version="0.1.0",
                    capabilities=self.server.get_capabilities(
                        notification_options=NotificationOptions(),
                        experimental_capabilities={},
                    ),
                ),
            )

async def main():
    server = QRGeneratorServer()
    await server.run()

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