from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request

class SmartSchemeMiddleware(BaseHTTPMiddleware):
    """
    Middleware to set the correct URL scheme (http or https) based on environment.
    - When running locally (host is 'localhost' or '127.0.0.1'), the scheme remains 'http'.
    - When deployed behind a proxy (e.g., on Hugging Face Spaces), the middleware checks for
      the 'x-forwarded-proto' header and uses its value to set the scheme (typically 'https').
    This helps avoid mixed content issues (e.g., when generating static URLs in templates),
    especially when deploying to environments where HTTPS is terminated at the proxy level.
    """
    async def dispatch(self, request: Request, call_next):
        # Detect local development by checking host
        host = request.headers.get("host", "")
        is_local = host.startswith("127.0.0.1") or host.startswith("localhost")

        # Only override scheme if not running locally
        if not is_local and "x-forwarded-proto" in request.headers:
            request.scope["scheme"] = request.headers["x-forwarded-proto"]

        return await call_next(request)