import os
import base64
from aiohttp import web

class BasicAuthMiddleware:
    def __init__(self):
        self.username = os.getenv('COMFYUI_USERNAME', '')
        self.password = os.getenv('COMFYUI_PASSWORD', '')
        self.enabled = bool(self.username and self.password)

    @web.middleware
    async def handle(self, request, handler):
        if not self.enabled:
            return await handler(request)

        # Skip auth for WebSocket connections
        if request.path == '/ws':
            return await handler(request)

        # Get Authorization header
        auth_header = request.headers.get('Authorization')

        if not auth_header:
            return web.Response(
                status=401,
                headers={'WWW-Authenticate': 'Basic realm="ComfyUI Server"'}
            )

        try:
            auth_type, auth_string = auth_header.split(' ', 1)
            if auth_type.lower() != 'basic':
                raise ValueError('Invalid auth type')

            decoded = base64.b64decode(auth_string).decode('utf-8')
            provided_username, provided_password = decoded.split(':', 1)

            if provided_username == self.username and provided_password == self.password:
                return await handler(request)

        except Exception:
            pass

        return web.Response(
            status=401,
            headers={'WWW-Authenticate': 'Basic realm="ComfyUI Server"'}
        )

class BasicAuthSetup:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "enabled": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ("BASIC_AUTH",)
    FUNCTION = "setup_auth"
    CATEGORY = "utils"
    OUTPUT_NODE = True

    def setup_auth(self, enabled):
        return ({"enabled": enabled},)

NODE_CLASS_MAPPINGS = {
    "BasicAuthSetup": BasicAuthSetup
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "BasicAuthSetup": "Basic Auth Setup"
}

# Register the middleware
try:
    import server
    app = server.PromptServer.instance.app
    middleware = BasicAuthMiddleware()
    app.middlewares.insert(0, middleware.handle)
except Exception as e:
    print(f"Failed to register basic auth middleware: {e}")