#!/usr/bin/env python3

# Generated by Claude

"""

Mock Google OAuth2 server for testing.

This FastAPI server mocks Google's OAuth2 endpoints and always returns
successful authentication responses.

Run with: python tests/oauth_mock.py
Test with: curl -X POST http://localhost:8080/oauth2/token -d "grant_type=authorization_code&code=test"
"""

import base64
import json
import time
from datetime import datetime, timedelta
from typing import Optional
from uuid import uuid4

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding, rsa
from fastapi import FastAPI, Form, HTTPException, Request
from fastapi.responses import JSONResponse

# Configurable email for mock responses
_mock_email = "test@example.com"

app = FastAPI(title="Mock Google OAuth2 Server", version="1.0.0")

# Generate RSA key pair for signing tokens
_private_key = rsa.generate_private_key(
    public_exponent=65537, key_size=2048, backend=default_backend()
)
_public_key = _private_key.public_key()

# Export public key components for JWKS
_public_numbers = _public_key.public_numbers()
_modulus = (
    base64.urlsafe_b64encode(
        _public_numbers.n.to_bytes((_public_numbers.n.bit_length() + 7) // 8, "big")
    )
    .decode()
    .rstrip("=")
)
_exponent = (
    base64.urlsafe_b64encode(
        _public_numbers.e.to_bytes((_public_numbers.e.bit_length() + 7) // 8, "big")
    )
    .decode()
    .rstrip("=")
)


def create_mock_id_token(email: str, sub: str, nonce: str, base_url: str, client_id: str) -> str:
    """Create a JWT ID token with a real RS256 signature."""
    header = {"alg": "RS256", "typ": "JWT", "kid": "mock_key_id"}
    payload = {
        "iss": base_url,
        "aud": client_id,
        "azp": client_id,
        "sub": sub,
        "email": email,
        "email_verified": True,
        "name": f"Test User {sub}",
        "given_name": "Test",
        "family_name": "User",
        "exp": int((datetime.now() + timedelta(hours=1)).timestamp()),
        "iat": int(datetime.now().timestamp()),
        "hd": "example.com",  # Hosted domain
    }

    # Add nonce if provided
    if nonce:
        payload["nonce"] = nonce

    # Encode header and payload
    header_b64 = (
        base64.urlsafe_b64encode(json.dumps(header, separators=(",", ":")).encode())
        .decode()
        .rstrip("=")
    )
    payload_b64 = (
        base64.urlsafe_b64encode(json.dumps(payload, separators=(",", ":")).encode())
        .decode()
        .rstrip("=")
    )

    # Create signing input
    signing_input = f"{header_b64}.{payload_b64}".encode()

    # Sign with private key using RS256 (RSA with SHA-256)
    signature = _private_key.sign(signing_input, padding.PKCS1v15(), hashes.SHA256())

    # Encode signature
    signature_b64 = base64.urlsafe_b64encode(signature).decode().rstrip("=")

    return f"{header_b64}.{payload_b64}.{signature_b64}"


# Store mock tokens for refresh and authorization codes
mock_token_store = {}
mock_auth_code_store = {}  # Store nonce and other data for auth codes


@app.get("/.well-known/openid-configuration")
async def openid_configuration(request: Request):
    """Mock OpenID Connect configuration endpoint."""
    base_url = f"{request.url.scheme}://{request.url.netloc}"
    return JSONResponse(
        {
            "issuer": base_url,
            "authorization_endpoint": f"{base_url}/oauth2/auth",
            "token_endpoint": f"{base_url}/oauth2/token",
            "userinfo_endpoint": f"{base_url}/oauth2/userinfo",
            "jwks_uri": f"{base_url}/oauth2/certs",
            "response_types_supported": ["code", "token", "id_token"],
            "subject_types_supported": ["public"],
            "id_token_signing_alg_values_supported": ["RS256"],
            "scopes_supported": ["openid", "email", "profile"],
            "claims_supported": [
                "iss",
                "sub",
                "aud",
                "exp",
                "iat",
                "email",
                "email_verified",
            ],
        }
    )


@app.post("/oauth2/token")
async def token_endpoint(
    request: Request,
    grant_type: str = Form(...),
    code: Optional[str] = Form(None),
    refresh_token: Optional[str] = Form(None),
    client_id: Optional[str] = Form(None),
    client_secret: Optional[str] = Form(None),
    redirect_uri: Optional[str] = Form(None),
):
    """Mock OAuth2 token endpoint - always returns success."""

    if grant_type == "authorization_code":
        # Initial token request with authorization code
        access_token = f"AT{uuid4()}"
        new_refresh_token = f"RT{uuid4()}"

        # Retrieve nonce and redirect_uri from auth code store if it exists
        auth_code_data = mock_auth_code_store.get(code, {})
        nonce = auth_code_data.get("nonce")
        stored_redirect_uri = auth_code_data.get("redirect_uri")

        # Validate redirect_uri if it was provided in both the auth and token requests
        if redirect_uri and stored_redirect_uri and redirect_uri != stored_redirect_uri:
            raise HTTPException(
                status_code=400,
                detail={
                    "error": "invalid_grant",
                    "error_description": "redirect_uri mismatch",
                },
            )

        base_url = f"{request.url.scheme}://{request.url.netloc}"
        id_token = create_mock_id_token(
            email=_mock_email, sub="123456789", nonce=nonce, base_url=base_url, client_id=client_id
        )

        # Store token data for potential refresh
        mock_token_store[new_refresh_token] = {
            "email": _mock_email,
            "sub": "123456789",
            "access_token": access_token,
            "created_at": datetime.now(),
        }

        # Clean up used auth code
        if code in mock_auth_code_store:
            del mock_auth_code_store[code]

        return JSONResponse(
            {
                "access_token": access_token,
                "refresh_token": new_refresh_token,
                "token_type": "Bearer",
                "expires_in": 3600,
                "id_token": id_token,
                "scope": "openid email profile",
            }
        )

    elif grant_type == "refresh_token":
        # Refresh token request
        if not refresh_token:
            raise HTTPException(status_code=400, detail="refresh_token required")

        # Get stored user data or use defaults
        stored_data = mock_token_store.get(
            refresh_token, {"email": _mock_email, "sub": "123456789"}
        )

        new_access_token = f"mock_access_token_refreshed_{int(time.time())}"
        base_url = f"{request.url.scheme}://{request.url.netloc}"
        new_id_token = create_mock_id_token(
            email=stored_data.get("email", _mock_email),
            sub=stored_data.get("sub", "123456789"),
            nonce="",
            base_url=base_url,
            client_id=client_id,
        )

        # Update stored data
        mock_token_store[refresh_token]["access_token"] = new_access_token

        return JSONResponse(
            {
                "access_token": new_access_token,
                "token_type": "Bearer",
                "expires_in": 3600,
                "id_token": new_id_token,
            }
        )

    else:
        raise HTTPException(
            status_code=400,
            detail={
                "error": "unsupported_grant_type",
                "error_description": f"Grant type '{grant_type}' not supported",
            },
        )


@app.get("/oauth2/userinfo")
async def userinfo_endpoint():
    """Mock OAuth2 userinfo endpoint."""
    return JSONResponse(
        {
            "sub": "123456789",
            "email": _mock_email,
            "email_verified": True,
            "name": "Test User",
            "given_name": "Test",
            "family_name": "User",
            "picture": "https://example.com/avatar.jpg",
            "locale": "en",
            "hd": "example.com",
        }
    )


@app.get("/oauth2/auth")
async def authorize_endpoint(
    request: Request,
    client_id: str,
    redirect_uri: str,
    response_type: str = "code",
    scope: str = "openid email",
    state: Optional[str] = None,
    nonce: Optional[str] = None,
):
    """Mock OAuth2 authorization endpoint."""
    from fastapi.responses import RedirectResponse

    # Always approve and redirect with mock authorization code
    auth_code = f"mock_auth_code_{int(time.time())}"

    # Store nonce and redirect_uri with the auth code for later retrieval
    mock_auth_code_store[auth_code] = {
        "nonce": nonce,
        "redirect_uri": redirect_uri,
    }

    params = f"code={auth_code}"
    if state:
        params += f"&state={state}"

    # Add id_token and access_token to the response
    base_url = f"{request.url.scheme}://{request.url.netloc}"
    id_token = create_mock_id_token(
        email=_mock_email, sub="123456789", nonce=nonce, base_url=base_url, client_id=client_id
    )
    access_token = f"mock_access_token_{int(time.time())}"

    params += f"&id_token={id_token}"
    params += f"&access_token={access_token}"
    params += "&token_type=Bearer"

    return RedirectResponse(url=f"{redirect_uri}?{params}")


@app.get("/oauth2/certs")
async def certs_endpoint():
    """JWKS endpoint with real public key."""
    return JSONResponse(
        {
            "keys": [
                {
                    "kty": "RSA",
                    "use": "sig",
                    "kid": "mock_key_id",
                    "n": _modulus,
                    "e": _exponent,
                    "alg": "RS256",
                }
            ]
        }
    )


@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {"status": "ok", "message": "Mock OAuth2 server is running"}


@app.post("/set_email")
async def set_email(email: str = Form(...)):
    """Set the email address to be used in mock OAuth2 responses."""
    global _mock_email
    _mock_email = email
    return JSONResponse(
        {
            "status": "ok",
            "message": f"Email set to {email}",
            "email": _mock_email,
        }
    )


@app.get("/")
async def root(request: Request):
    """Root endpoint with server info."""
    base_url = f"{request.url.scheme}://{request.url.netloc}"
    return {
        "name": "Mock Google OAuth2 Server",
        "version": "1.0.0",
        "endpoints": {
            "authorization": "/oauth2/auth",
            "token": "/oauth2/token",
            "userinfo": "/oauth2/userinfo",
            "jwks": "/oauth2/certs",
            "openid_configuration": "/.well-known/openid-configuration",
        },
        "test_commands": {
            "get_token": f"curl -X POST {base_url}/oauth2/token -d 'grant_type=authorization_code&code=test'",
            "refresh_token": f"curl -X POST {base_url}/oauth2/token -d 'grant_type=refresh_token&refresh_token=YOUR_REFRESH_TOKEN'",
            "get_userinfo": f"curl {base_url}/oauth2/userinfo",
            "set_email": f"curl -X POST {base_url}/set_email -d 'email=custom@example.com'",
        },
    }
