#!/usr/bin/env python3
"""
MCP Deployer Server - Deploy MCP Servers to Modal.com on the fly

This MCP server provides tools to:
1. Deploy Python MCP code to Modal.com
2. Manage multiple deployments
3. Get deployment URLs and status
4. Run tests against deployed servers

All deployments use minimal resources with cold starts allowed (no billing when idle).
"""

import json
import os
import re
import subprocess
import tempfile
import hashlib
from datetime import datetime
from pathlib import Path
from typing import Optional

from fastmcp import FastMCP

# Initialize the MCP server
mcp = FastMCP("mcp-deployer")

# Deployment registry path
DEPLOYMENTS_DIR = Path(__file__).parent / "deployments"
DEPLOYMENTS_DIR.mkdir(exist_ok=True)
REGISTRY_FILE = DEPLOYMENTS_DIR / "registry.json"

# Modal wrapper template
MODAL_WRAPPER_TEMPLATE = '''#!/usr/bin/env python3
"""
Auto-generated Modal deployment for MCP Server: {app_name}
Generated at: {timestamp}
"""

import modal

# App configuration with minimal resources and cold starts allowed
app = modal.App("{app_name}")

# Image with required dependencies
image = modal.Image.debian_slim(python_version="3.12").pip_install(
    "fastapi==0.115.14",
    "fastmcp>=2.10.0",
    "pydantic>=2.0.0",
    "requests>=2.28.0",
    "uvicorn>=0.20.0",
    {extra_deps}
)

def make_mcp_server():
    """Create the MCP server with user-defined tools"""
    from fastmcp import FastMCP
    
    mcp = FastMCP("{server_name}")
    
    # === USER-DEFINED TOOLS START ===
{user_code_indented}
    # === USER-DEFINED TOOLS END ===
    
    return mcp


@app.function(
    image=image,
    # Cost optimization: minimal resources, allow cold starts
    cpu=0.25,           # 1/4 CPU core (cheapest)
    memory=256,         # 256 MB memory (minimal)
    timeout=300,        # 5 min timeout
    # Scale to zero when not in use (no billing when idle)
    scaledown_window=2, # Scale down after 2 seconds of inactivity
)
@modal.asgi_app()
def web():
    """ASGI web endpoint for the MCP server"""
    from fastapi import FastAPI
    
    mcp = make_mcp_server()
    mcp_app = mcp.http_app(transport="streamable-http", stateless_http=True)
    
    fastapi_app = FastAPI(
        title="{server_name}",
        description="Auto-deployed MCP Server on Modal.com",
        lifespan=mcp_app.router.lifespan_context
    )
    fastapi_app.mount("/", mcp_app, "mcp")
    
    return fastapi_app


# Test function to verify deployment
@app.function(image=image)
async def test_server():
    """Test the deployed MCP server"""
    from fastmcp import Client
    from fastmcp.client.transports import StreamableHttpTransport
    
    transport = StreamableHttpTransport(url=f"{{web.get_web_url()}}/mcp/")
    client = Client(transport)
    
    async with client:
        tools = await client.list_tools()
        return {{
            "status": "ok",
            "tools_count": len(tools),
            "tools": [t.name for t in tools]
        }}
'''


def load_registry() -> dict:
    """Load the deployment registry"""
    if REGISTRY_FILE.exists():
        return json.loads(REGISTRY_FILE.read_text())
    return {"deployments": {}}


def save_registry(registry: dict):
    """Save the deployment registry"""
    REGISTRY_FILE.write_text(json.dumps(registry, indent=2))


def generate_app_name(server_name: str) -> str:
    """Generate a unique Modal app name from server name"""
    # Sanitize name: lowercase, alphanumeric and hyphens only
    sanitized = re.sub(r'[^a-z0-9-]', '-', server_name.lower())
    sanitized = re.sub(r'-+', '-', sanitized).strip('-')
    
    # Add short hash for uniqueness
    hash_suffix = hashlib.md5(f"{server_name}{datetime.now().isoformat()}".encode()).hexdigest()[:6]
    
    return f"mcp-{sanitized[:40]}-{hash_suffix}"


def extract_imports_and_code(user_code: str) -> tuple[list[str], str]:
    """Extract import statements and separate from function code"""
    lines = user_code.strip().split('\n')
    imports = []
    code_lines = []
    
    for line in lines:
        stripped = line.strip()
        if stripped.startswith('import ') or stripped.startswith('from '):
            # Extract package name for pip install
            if stripped.startswith('from '):
                match = re.match(r'from\s+(\w+)', stripped)
                if match:
                    imports.append(match.group(1))
            else:
                match = re.match(r'import\s+(\w+)', stripped)
                if match:
                    imports.append(match.group(1))
        code_lines.append(line)
    
    return imports, '\n'.join(code_lines)


def indent_code(code: str, spaces: int = 4) -> str:
    """Indent code by specified number of spaces"""
    indent = ' ' * spaces
    return '\n'.join(indent + line if line.strip() else line for line in code.split('\n'))


@mcp.tool()
def deploy_mcp_server(
    server_name: str,
    mcp_tools_code: str,
    extra_pip_packages: list[str] = None,
    description: str = None
) -> dict:
    """
    Deploy an MCP server with custom tools to Modal.com.
    
    The deployed server will:
    - Use minimal CPU (0.25 cores) and memory (256MB)
    - Scale to zero when not in use (no billing when idle)
    - Allow cold starts (2-5 second startup time)
    - Be accessible via a public URL
    
    Args:
        server_name: Name for your MCP server (e.g., "fuel-prices", "weather-api")
        mcp_tools_code: Python code defining your MCP tools using @mcp.tool() decorator.
                       Do NOT include FastMCP initialization - just the tool functions.
                       Example:
                       ```
                       @mcp.tool()
                       def get_time(timezone: str = "UTC") -> str:
                           '''Get current time in specified timezone'''
                           from datetime import datetime
                           from zoneinfo import ZoneInfo
                           return datetime.now(ZoneInfo(timezone)).isoformat()
                       ```
        extra_pip_packages: Additional pip packages to install (e.g., ["pandas", "numpy"])
        description: Optional description for the server
    
    Returns:
        dict with:
        - success: bool
        - app_name: Modal app name
        - url: Deployed server URL (after deployment completes)
        - mcp_endpoint: Full MCP endpoint URL (url + /mcp/)
        - deployment_id: Unique deployment identifier
        - message: Status message
    """
    try:
        # Generate unique app name
        app_name = generate_app_name(server_name)
        
        # Extract imports and prepare extra dependencies
        detected_imports, cleaned_code = extract_imports_and_code(mcp_tools_code)
        
        # Combine detected imports with explicitly specified packages
        all_packages = list(set(detected_imports + (extra_pip_packages or [])))
        
        # Filter out standard library packages
        stdlib = {'os', 'sys', 'json', 're', 'datetime', 'time', 'typing', 'pathlib', 
                  'collections', 'functools', 'itertools', 'math', 'random', 'string',
                  'hashlib', 'base64', 'urllib', 'zoneinfo', 'asyncio'}
        extra_deps = [pkg for pkg in all_packages if pkg.lower() not in stdlib]
        
        # Format extra dependencies for template
        extra_deps_str = ',\n    '.join(f'"{pkg}"' for pkg in extra_deps) if extra_deps else ''
        
        # Indent user code for template insertion
        user_code_indented = indent_code(cleaned_code, spaces=4)
        
        # Generate Modal wrapper code
        modal_code = MODAL_WRAPPER_TEMPLATE.format(
            app_name=app_name,
            server_name=server_name,
            timestamp=datetime.now().isoformat(),
            extra_deps=extra_deps_str,
            user_code_indented=user_code_indented
        )
        
        # Create deployment file
        deploy_dir = DEPLOYMENTS_DIR / app_name
        deploy_dir.mkdir(exist_ok=True)
        deploy_file = deploy_dir / "app.py"
        deploy_file.write_text(modal_code)
        
        # Also save original user code for reference
        (deploy_dir / "original_tools.py").write_text(mcp_tools_code)
        
        # Deploy to Modal
        result = subprocess.run(
            ["modal", "deploy", str(deploy_file)],
            capture_output=True,
            text=True,
            timeout=300  # 5 minute timeout for deployment
        )
        
        if result.returncode != 0:
            return {
                "success": False,
                "error": "Deployment failed",
                "stdout": result.stdout,
                "stderr": result.stderr,
                "app_name": app_name,
                "deploy_file": str(deploy_file)
            }
        
        # Extract URL from deployment output
        # Modal outputs URLs in format: https://xxx--yyy.modal.run
        url_match = re.search(r'https://[a-zA-Z0-9-]+--[a-zA-Z0-9-]+\.modal\.run', result.stdout)
        deployed_url = url_match.group(0) if url_match else None
        
        # If URL not found in stdout, try to get it programmatically
        if not deployed_url:
            try:
                import modal
                remote_func = modal.Function.from_name(app_name, "web")
                deployed_url = remote_func.get_web_url()
            except Exception:
                deployed_url = f"https://<workspace>--{app_name}-web.modal.run"
        
        # Update registry
        registry = load_registry()
        deployment_id = f"deploy-{app_name}"
        registry["deployments"][deployment_id] = {
            "app_name": app_name,
            "server_name": server_name,
            "url": deployed_url,
            "mcp_endpoint": f"{deployed_url}/mcp/" if deployed_url else None,
            "description": description,
            "created_at": datetime.now().isoformat(),
            "status": "deployed",
            "extra_packages": extra_deps,
            "deploy_file": str(deploy_file)
        }
        save_registry(registry)
        
        return {
            "success": True,
            "app_name": app_name,
            "url": deployed_url,
            "mcp_endpoint": f"{deployed_url}/mcp/" if deployed_url else None,
            "deployment_id": deployment_id,
            "deploy_file": str(deploy_file),
            "detected_packages": extra_deps,
            "message": f"✅ Successfully deployed '{server_name}' to Modal.com!\n"
                      f"🔗 URL: {deployed_url}\n"
                      f"📡 MCP Endpoint: {deployed_url}/mcp/\n"
                      f"💡 Use this URL in your MCP client configuration."
        }
        
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "error": "Deployment timed out after 5 minutes",
            "app_name": app_name if 'app_name' in dir() else None
        }
    except Exception as e:
        return {
            "success": False,
            "error": str(e),
            "app_name": app_name if 'app_name' in dir() else None
        }


@mcp.tool()
def list_deployments() -> dict:
    """
    List all deployed MCP servers.
    
    Returns:
        dict with:
        - success: bool
        - total: Number of deployments
        - deployments: List of deployment details
    """
    try:
        registry = load_registry()
        deployments = registry.get("deployments", {})
        
        deployment_list = []
        for dep_id, info in deployments.items():
            deployment_list.append({
                "deployment_id": dep_id,
                "app_name": info.get("app_name"),
                "server_name": info.get("server_name"),
                "url": info.get("url"),
                "mcp_endpoint": info.get("mcp_endpoint"),
                "status": info.get("status"),
                "created_at": info.get("created_at"),
                "description": info.get("description")
            })
        
        # Sort by creation date (newest first)
        deployment_list.sort(key=lambda x: x.get("created_at", ""), reverse=True)
        
        return {
            "success": True,
            "total": len(deployment_list),
            "deployments": deployment_list
        }
        
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }


@mcp.tool()
def get_deployment_status(deployment_id: str = None, app_name: str = None) -> dict:
    """
    Get detailed status of a deployed MCP server.
    
    Args:
        deployment_id: The deployment ID (e.g., "deploy-mcp-fuel-prices-abc123")
        app_name: Or the Modal app name (e.g., "mcp-fuel-prices-abc123")
    
    Returns:
        dict with deployment details and status
    """
    try:
        registry = load_registry()
        deployments = registry.get("deployments", {})
        
        # Find deployment by ID or app_name
        info = None
        found_id = None
        
        if deployment_id and deployment_id in deployments:
            info = deployments[deployment_id]
            found_id = deployment_id
        elif app_name:
            for dep_id, dep_info in deployments.items():
                if dep_info.get("app_name") == app_name:
                    info = dep_info
                    found_id = dep_id
                    break
        
        if not info:
            return {
                "success": False,
                "error": f"Deployment not found: {deployment_id or app_name}"
            }
        
        # Try to get current URL from Modal
        try:
            import modal
            remote_func = modal.Function.from_name(info["app_name"], "web")
            current_url = remote_func.get_web_url()
            info["url"] = current_url
            info["mcp_endpoint"] = f"{current_url}/mcp/"
            info["live"] = True
        except Exception:
            info["live"] = False
        
        return {
            "success": True,
            "deployment_id": found_id,
            **info
        }
        
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }


@mcp.tool()
def delete_deployment(deployment_id: str = None, app_name: str = None, confirm: bool = False) -> dict:
    """
    Delete a deployed MCP server from Modal.
    
    Args:
        deployment_id: The deployment ID to delete
        app_name: Or the Modal app name to delete
        confirm: Must be True to confirm deletion
    
    Returns:
        dict with deletion status
    """
    if not confirm:
        return {
            "success": False,
            "error": "Must set confirm=True to delete deployment"
        }
    
    try:
        registry = load_registry()
        deployments = registry.get("deployments", {})
        
        # Find deployment
        target_app_name = None
        found_id = None
        
        if deployment_id and deployment_id in deployments:
            target_app_name = deployments[deployment_id].get("app_name")
            found_id = deployment_id
        elif app_name:
            for dep_id, dep_info in deployments.items():
                if dep_info.get("app_name") == app_name:
                    target_app_name = app_name
                    found_id = dep_id
                    break
        
        if not target_app_name:
            return {
                "success": False,
                "error": f"Deployment not found: {deployment_id or app_name}"
            }
        
        # Stop the Modal app
        result = subprocess.run(
            ["modal", "app", "stop", target_app_name],
            capture_output=True,
            text=True,
            timeout=60
        )
        
        # Remove from registry
        if found_id in deployments:
            del deployments[found_id]
            registry["deployments"] = deployments
            save_registry(registry)
        
        # Clean up local files
        deploy_dir = DEPLOYMENTS_DIR / target_app_name
        if deploy_dir.exists():
            import shutil
            shutil.rmtree(deploy_dir)
        
        return {
            "success": True,
            "app_name": target_app_name,
            "deployment_id": found_id,
            "message": f"✅ Deleted deployment '{target_app_name}'"
        }
        
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "error": "Deletion timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }


@mcp.tool()
def test_mcp_server(url: str = None, deployment_id: str = None) -> dict:
    """
    Test a deployed MCP server by listing its tools.
    
    Args:
        url: The MCP server URL (e.g., "https://xxx.modal.run")
        deployment_id: Or the deployment ID to test
    
    Returns:
        dict with test results including available tools
    """
    try:
        # Get URL from deployment_id if not provided
        if not url and deployment_id:
            registry = load_registry()
            if deployment_id in registry.get("deployments", {}):
                url = registry["deployments"][deployment_id].get("url")
        
        if not url:
            return {
                "success": False,
                "error": "Must provide either url or deployment_id"
            }
        
        # Ensure URL ends with /mcp/
        mcp_url = url.rstrip('/') + '/mcp/'
        
        # Test using requests (simpler than async client)
        import requests
        
        # First, check if the server is reachable
        response = requests.get(url, timeout=30)
        
        if response.status_code == 200:
            return {
                "success": True,
                "url": url,
                "mcp_endpoint": mcp_url,
                "status_code": response.status_code,
                "message": f"✅ Server is reachable at {url}\n"
                          f"📡 MCP endpoint: {mcp_url}\n"
                          f"💡 Connect using: npx @modelcontextprotocol/inspector {mcp_url}"
            }
        else:
            return {
                "success": False,
                "url": url,
                "status_code": response.status_code,
                "error": f"Server returned status {response.status_code}"
            }
            
    except requests.exceptions.Timeout:
        return {
            "success": False,
            "error": "Request timed out - server may be cold starting (try again in 5-10 seconds)"
        }
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }


def _generate_modal_code_internal(
    server_name: str,
    mcp_tools_code: str,
    extra_pip_packages: list[str] = None
) -> dict:
    """Internal function to generate Modal deployment code"""
    try:
        app_name = generate_app_name(server_name)
        
        # Extract imports and prepare extra dependencies
        detected_imports, cleaned_code = extract_imports_and_code(mcp_tools_code)
        all_packages = list(set(detected_imports + (extra_pip_packages or [])))
        
        stdlib = {'os', 'sys', 'json', 're', 'datetime', 'time', 'typing', 'pathlib', 
                  'collections', 'functools', 'itertools', 'math', 'random', 'string',
                  'hashlib', 'base64', 'urllib', 'zoneinfo', 'asyncio'}
        extra_deps = [pkg for pkg in all_packages if pkg.lower() not in stdlib]
        
        extra_deps_str = ',\n    '.join(f'"{pkg}"' for pkg in extra_deps) if extra_deps else ''
        user_code_indented = indent_code(cleaned_code, spaces=4)
        
        modal_code = MODAL_WRAPPER_TEMPLATE.format(
            app_name=app_name,
            server_name=server_name,
            timestamp=datetime.now().isoformat(),
            extra_deps=extra_deps_str,
            user_code_indented=user_code_indented
        )
        
        return {
            "success": True,
            "app_name": app_name,
            "code": modal_code,
            "detected_packages": extra_deps,
            "message": "Generated Modal deployment code. Review and use deploy_mcp_server() to deploy."
        }
        
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }


# Expose internal function for testing
generate_modal_code = _generate_modal_code_internal


@mcp.tool()
def preview_modal_code(
    server_name: str,
    mcp_tools_code: str,
    extra_pip_packages: list[str] = None
) -> dict:
    """
    Generate Modal deployment code WITHOUT deploying.
    Useful for reviewing the generated code before deployment.
    
    Args:
        server_name: Name for your MCP server
        mcp_tools_code: Python code defining your MCP tools
        extra_pip_packages: Additional pip packages to install
    
    Returns:
        dict with generated Modal code
    """
    return _generate_modal_code_internal(server_name, mcp_tools_code, extra_pip_packages)


if __name__ == "__main__":
    import sys
    import os

    host = "0.0.0.0"
    # Use PORT env var (set by Hugging Face Spaces to 7860), fallback to 8002 for local dev
    port = int(os.environ.get("PORT", 8002))

    # Command line argument can still override
    if len(sys.argv) > 1:
        port = int(sys.argv[1])
    
    print(f"🚀 MCP Deployer Server")
    print(f"📡 Available at: http://{host}:{port}")
    print(f"🔧 SSE endpoint: http://{host}:{port}/sse")
    print(f"\n📋 Available Tools:")
    print(f"   • deploy_mcp_server - Deploy MCP code to Modal.com")
    print(f"   • list_deployments - List all deployments")
    print(f"   • get_deployment_status - Check deployment status")
    print(f"   • delete_deployment - Remove a deployment")
    print(f"   • test_mcp_server - Test a deployed server")
    print(f"   • preview_modal_code - Preview generated code")
    print(f"\n⏳ Starting server...\n")
    
    mcp.run(transport="sse", host=host, port=port)