import uuid
from datetime import UTC
from datetime import datetime
from datetime import timedelta

import jwt

from amsdal_ml.mcp_server.mcp_oauth import oauth_router
from amsdal_ml.mcp_server.mcp_oauth.services.models import RegisterRequest
from amsdal_ml.mcp_server.mcp_oauth.services.models import RegisterResponse
from amsdal_ml.mcp_server.mcp_oauth.services.types import JWT_ALGORITHM
from amsdal_ml.mcp_server.mcp_oauth.services.types import ClientIdTokenClaims
from amsdal_ml.mcp_server.mcp_oauth.services.types import OAuthConstants
from amsdal_ml.mcp_server.mcp_oauth.services.utils import get_signing_key
from amsdal_ml.ml_config import ml_config


@oauth_router.post('/register', response_model=RegisterResponse)
async def register(req: RegisterRequest):
    signing_key = get_signing_key()

    now = datetime.now(tz=UTC)
    iat = int(now.timestamp())
    expiration_time = now + timedelta(days=ml_config.oauth_client_id_expiration_days)
    expiration_ts = int(expiration_time.timestamp())
    jti = str(uuid.uuid4())

    payload = {
        ClientIdTokenClaims.REDIRECT_URIS.value: req.redirect_uris,
        ClientIdTokenClaims.GRANT_TYPES.value: req.grant_types,
        ClientIdTokenClaims.RESPONSE_TYPES.value: req.response_types,
        ClientIdTokenClaims.TOKEN_ENDPOINT_AUTH_METHOD.value: OAuthConstants.TOKEN_ENDPOINT_AUTH_NONE,
        ClientIdTokenClaims.ISSUED_AT.value: iat,
        ClientIdTokenClaims.EXPIRATION.value: expiration_ts,
        ClientIdTokenClaims.JWT_ID.value: jti,
        ClientIdTokenClaims.ISSUER.value: ml_config.oauth_issuer,
    }

    token = jwt.encode(
        payload,
        signing_key,
        algorithm=JWT_ALGORITHM,
    )

    return RegisterResponse(
        client_id=token,
        expires_at=expiration_ts,
        client_id_issued_at=iat,
    )
