"""
Example MCP Server Implementation for Document Translation API
This demonstrates how to wrap the translation API as an MCP server
"""
import asyncio
import httpx
from typing import Any
from mcp.server.models import InitializationOptions
from mcp.server import NotificationOptions, Server
from mcp.server.stdio import stdio_server
from mcp import types

# API Configuration
API_BASE_URL = "http://localhost:8000"


class DocumentTranslatorMCP:
    """MCP Server for Document Translation API"""
    
    def __init__(self):
        self.server = Server("document-translator")
        self.http_client = None
        self._setup_handlers()
    
    def _setup_handlers(self):
        """Set up MCP tool handlers"""
        
        @self.server.list_tools()
        async def handle_list_tools() -> list[types.Tool]:
            """List available tools"""
            return [
                types.Tool(
                    name="translate_document",
                    description="Translate a document (Excel, Word, or PowerPoint) while preserving all formatting",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "file_path": {
                                "type": "string",
                                "description": "Path to the document file to translate"
                            },
                            "target_language": {
                                "type": "string",
                                "description": "Target language code (e.g., 'es', 'fr', 'de')"
                            },
                            "source_language": {
                                "type": "string",
                                "description": "Source language code (default: 'auto' for auto-detection)",
                                "default": "auto"
                            },
                            "output_path": {
                                "type": "string",
                                "description": "Path where the translated document should be saved"
                            }
                        },
                        "required": ["file_path", "target_language", "output_path"]
                    }
                ),
                types.Tool(
                    name="get_supported_languages",
                    description="Get list of supported language codes for translation",
                    inputSchema={
                        "type": "object",
                        "properties": {}
                    }
                ),
                types.Tool(
                    name="check_api_health",
                    description="Check if the translation API is healthy and operational",
                    inputSchema={
                        "type": "object",
                        "properties": {}
                    }
                )
            ]
        
        @self.server.call_tool()
        async def handle_call_tool(
            name: str,
            arguments: dict[str, Any] | None
        ) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
            """Handle tool calls"""
            
            if name == "translate_document":
                return await self._translate_document(arguments)
            elif name == "get_supported_languages":
                return await self._get_supported_languages()
            elif name == "check_api_health":
                return await self._check_health()
            else:
                raise ValueError(f"Unknown tool: {name}")
    
    async def _translate_document(self, args: dict[str, Any]) -> list[types.TextContent]:
        """Translate a document via the API"""
        file_path = args["file_path"]
        target_language = args["target_language"]
        source_language = args.get("source_language", "auto")
        output_path = args["output_path"]
        
        try:
            async with httpx.AsyncClient(timeout=300.0) as client:
                # Upload and translate the document
                with open(file_path, "rb") as f:
                    files = {"file": (file_path, f)}
                    data = {
                        "target_language": target_language,
                        "source_language": source_language
                    }
                    
                    response = await client.post(
                        f"{API_BASE_URL}/translate",
                        files=files,
                        data=data
                    )
                    
                    if response.status_code == 200:
                        # Save the translated document
                        with open(output_path, "wb") as output:
                            output.write(response.content)
                        
                        return [
                            types.TextContent(
                                type="text",
                                text=f"✅ Document translated successfully!\n\n"
                                     f"Original: {file_path}\n"
                                     f"Translated: {output_path}\n"
                                     f"Language: {source_language} → {target_language}\n"
                                     f"Size: {len(response.content)} bytes"
                            )
                        ]
                    else:
                        error_detail = response.json().get("detail", "Unknown error")
                        return [
                            types.TextContent(
                                type="text",
                                text=f"❌ Translation failed: {error_detail}"
                            )
                        ]
        
        except Exception as e:
            return [
                types.TextContent(
                    type="text",
                    text=f"❌ Error during translation: {str(e)}"
                )
            ]
    
    async def _get_supported_languages(self) -> list[types.TextContent]:
        """Get supported languages from the API"""
        try:
            async with httpx.AsyncClient() as client:
                response = await client.get(f"{API_BASE_URL}/languages")
                
                if response.status_code == 200:
                    data = response.json()
                    languages = data.get("supported_languages", {})
                    
                    lang_list = "\n".join([f"- {code}: {name}" for code, name in languages.items()])
                    
                    return [
                        types.TextContent(
                            type="text",
                            text=f"📚 Supported Languages:\n\n{lang_list}\n\n"
                                 f"Note: {data.get('note', '')}"
                        )
                    ]
                else:
                    return [
                        types.TextContent(
                            type="text",
                            text="❌ Failed to retrieve supported languages"
                        )
                    ]
        
        except Exception as e:
            return [
                types.TextContent(
                    type="text",
                    text=f"❌ Error: {str(e)}"
                )
            ]
    
    async def _check_health(self) -> list[types.TextContent]:
        """Check API health"""
        try:
            async with httpx.AsyncClient() as client:
                response = await client.get(f"{API_BASE_URL}/health")
                
                if response.status_code == 200:
                    data = response.json()
                    return [
                        types.TextContent(
                            type="text",
                            text=f"✅ API is healthy!\n\n"
                                 f"Status: {data.get('status')}\n"
                                 f"Translation Service: {data.get('translation_service')}"
                        )
                    ]
                else:
                    return [
                        types.TextContent(
                            type="text",
                            text="❌ API is not responding correctly"
                        )
                    ]
        
        except Exception as e:
            return [
                types.TextContent(
                    type="text",
                    text=f"❌ Cannot connect to API: {str(e)}"
                )
            ]
    
    async def run(self):
        """Run the MCP server"""
        async with stdio_server() as (read_stream, write_stream):
            await self.server.run(
                read_stream,
                write_stream,
                InitializationOptions(
                    server_name="document-translator",
                    server_version="1.0.0",
                    capabilities=self.server.get_capabilities(
                        notification_options=NotificationOptions(),
                        experimental_capabilities={}
                    )
                )
            )


async def main():
    """Main entry point"""
    mcp_server = DocumentTranslatorMCP()
    await mcp_server.run()


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