import argparse
import sys
import os
from .server import mcp

class AuthMiddleware:
    """Simple ASGI middleware for Bearer Token authentication."""
    def __init__(self, app, token):
        self.app = app
        self.token = token

    async def __call__(self, scope, receive, send):
        if scope["type"] == "http":
            # Allow health checks (root path /) to pass without auth
            if scope.get("path") == "/":
                await self.app(scope, receive, send)
                return

            # Allow health checks or OPTIONS if needed
            headers = dict(scope.get("headers", []))
            auth_header = headers.get(b"authorization", b"").decode("utf-8")
            
            # Check for "Bearer <token>"
            # Be a bit more flexible with whitespace and casing of "Bearer"
            authorized = False
            if auth_header:
                parts = auth_header.split(None, 1)
                if len(parts) == 2 and parts[0].lower() == "bearer":
                    if parts[1].strip() == self.token.strip():
                        authorized = True
            
            if not authorized:
                await send({
                    'type': 'http.response.start',
                    'status': 401,
                    'headers': [(b'content-type', b'text/plain')],
                })
                await send({
                    'type': 'http.response.body',
                    'body': b'Unauthorized: Invalid or missing Bearer token',
                })
                return
        
        # Pass through to the original app
        await self.app(scope, receive, send)

def main():
    parser = argparse.ArgumentParser(description="Image Search MCP Server")
    parser.add_argument("--sse", action="store_true", help="Run in SSE mode (HTTP)")
    parser.add_argument("--port", type=int, default=8000, help="Port for SSE mode")
    parser.add_argument("--host", type=str, default="127.0.0.1", help="Host for SSE mode")
    
    args = parser.parse_args()

    if args.sse:
        try:
            import uvicorn
            # Use the official SSE app creator from FastMCP
            try:
                if hasattr(mcp, "http_app"):
                    # Modern FastMCP (2.x) uses http_app with transport="sse"
                    app = mcp.http_app(transport="sse")
                else:
                    from mcp.server.fastmcp import create_sse_app
                    app = create_sse_app(mcp, sse_path="/sse", message_path="/messages")
            except Exception as e:
                print(f"Warning: Failed to use standard SSE app creation: {e}")
                # Fallback for different versions
                if hasattr(mcp, "http_app"):
                    app = mcp.http_app(path="/")
                else:
                    app = getattr(mcp, "_app", mcp)

            # Add a basic root route for connectivity testing
            try:
                from starlette.responses import JSONResponse
                from starlette.routing import Route
                async def homepage(request):
                    return JSONResponse({
                        "status": "running",
                        "service": "Image Search MCP Server",
                        "endpoints": {
                            "sse": "/sse",
                            "messages": "/messages"
                        }
                    })
                # Check if we can add a route
                if hasattr(app, "routes"):
                    app.routes.insert(0, Route("/", homepage))
            except Exception:
                pass

            # Check for Auth Token
            auth_token = os.environ.get("MCP_AUTH_TOKEN")
            if auth_token:
                print(f"Authentication enabled. Require Bearer token.")
                app = AuthMiddleware(app, auth_token)
            
            print(f"Starting SSE server on {args.host}:{args.port}...")
            uvicorn.run(app, host=args.host, port=args.port, log_level="info")
        except Exception as e:
            print(f"Error starting SSE server: {e}")
            import traceback
            traceback.print_exc()
            sys.exit(1)
    else:
        # Stdio mode (Default)
        mcp.run()

if __name__ == "__main__":
    main()
