from pydantic import AnyHttpUrl

from mcp.server.auth.provider import AccessToken, TokenVerifier
from mcp.server.auth.settings import AuthSettings
from mcp.server.fastmcp import FastMCP
from fastapi import FastAPI, Request
from starlette.responses import JSONResponse, RedirectResponse
import contextlib
import httpx


class SimpleTokenVerifier(TokenVerifier):
    """Simple token verifier for demonstration."""

    async def verify_token(self, token: str) -> AccessToken | None:
        pass  # This is where you would implement actual token validation


# Create FastMCP instance as a Resource Server
mcp = FastMCP(
    "Weather Service",
    # Token verifier for authentication
    token_verifier=SimpleTokenVerifier(),
    # Auth settings for RFC 9728 Protected Resource Metadata
    auth=AuthSettings(
        issuer_url=AnyHttpUrl("https://huggingface.co/.well-known/oauth-authorization-server"),  # Authorization Server URL
        resource_server_url=AnyHttpUrl("https://freddyaboulton-fastmcp-oauth.hf.space/"),  # This server's URL
        required_scopes=["user"],
    ),
)


@mcp.tool()
async def get_weather(city: str = "London") -> dict[str, str]:
    """Get weather data for a city"""
    return {
        "city": city,
        "temperature": "22",
        "condition": "Partly cloudy",
        "humidity": "65%",
    }

@contextlib.asynccontextmanager
async def lifespan(app: FastAPI):
    async with contextlib.AsyncExitStack() as stack:
        await stack.enter_async_context(mcp.session_manager.run())
        yield


app = FastAPI(lifespan=lifespan)
app.mount("/weather_mcp", mcp.streamable_http_app())

@app.get("/")
async def _():
    return JSONResponse({"message": "Welcome to the Weather Service!"})

@app.get("/.well-known/oauth-protected-resource")
async def _():
    return RedirectResponse("/weather_mcp/.well-known/oauth-protected-resource")

@app.post("/register")
async def _():
    return RedirectResponse("https://huggingface.co/oauth/register")

@app.post("/register")
async def register_proxy(request: Request):
    """Forward requests to Hugging Face OAuth register endpoint"""
    target_url = "https://huggingface.co/oauth/register"
    
    async with httpx.AsyncClient() as client:
        # Prepare the request data
        headers = dict(request.headers)
        # Remove host header to avoid conflicts
        headers.pop("host", None)
        
        # Get query parameters
        query_params = str(request.url.query) if request.url.query else None
        target_url_with_params = f"{target_url}?{query_params}" if query_params else target_url
        
        # Get request body if present
        body = await request.body()

        print("body", body)
        
        # Forward the request
        response = await client.request(
            method=request.method,
            url=target_url_with_params,
            headers=headers,
            content=body,
            follow_redirects=False
        )
        
        # Return the response with the same status code, headers, and content
        return Response(
            content=response.content,
            status_code=response.status_code,
            headers=dict(response.headers)
        )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)