"""Token Exchange."""

import json
import logging
from datetime import datetime

import httpx
from httpx import HTTPStatusError
from msgspec import Struct, ValidationError
from msgspec.json import Decoder as JsonDecoder
from pytz import UTC

from slices_clientlib_core.base.const import BASE_HTTP_TIMEOUT_S
from slices_clientlib_core.util.status import HTTP_400_BAD_REQUEST

from .config import settings
from .token import OAuthToken

logger = logging.getLogger(__name__)


class ExchangedToken(Struct):
    """Response from token exchange."""

    access_token: str
    expires_in: int
    issued_token_type: str
    scope: str
    token_type: str


class ExchangedIDToken(Struct):
    """Response from token exchange when requesting ID token."""

    access_token: str
    expires_in: int
    issued_token_type: str
    token_type: str


token_decoder = JsonDecoder(ExchangedToken)
id_token_decoder = JsonDecoder(ExchangedIDToken)


def get_id_token_for_audience(auth_token: OAuthToken, audience: str | list[str], project_id: str) -> str:
    """Return JWT ID token of logged in user for target resource."""

    data = {
        "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
        "subject_token": auth_token.access_token,
        "subject_token_type": "urn:ietf:params:oauth:token-type:access_token",
        "client_id": settings.client_id,
        "client_secret": settings.client_secret,
        "requested_token_type": "urn:ietf:params:oauth:token-type:id_token",
        "audience": audience,
        "scope": f"{auth_token.scope} project:{project_id}",
    }

    response = httpx.post(
        url=settings.token_endpoint_url,
        data=data,
        timeout=BASE_HTTP_TIMEOUT_S,
    )
    try:
        response.raise_for_status()
    except HTTPStatusError:
        logger.error("Failed to exchange token: %s", response.content)  # noqa: TRY400 just the error is ok here
        raise
    try:
        return id_token_decoder.decode(response.content).access_token
    except ValidationError:
        logger.error("Failed to decode JWT ID token from: %s", response.content)  # noqa: TRY400 just the error is ok here
        raise


def get_access_token_for_audience(
    auth_token: OAuthToken, audience: str, project_id: str | None = None
) -> ExchangedToken:
    """Return JWT access token of logged in user for target resource."""
    scope = auth_token.scope
    if project_id:
        scope = f"{scope} project:{project_id}"

    data = {
        "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
        "subject_token": auth_token.access_token,
        "subject_token_type": "urn:ietf:params:oauth:token-type:access_token",
        "client_id": settings.client_id,
        "client_secret": settings.client_secret,
        "requested_token_type": "urn:ietf:params:oauth:token-type:access_token",
        "audience": audience,
        "scope": scope,
    }

    response = httpx.post(
        url=settings.token_endpoint_url,
        data=data,
        timeout=BASE_HTTP_TIMEOUT_S,
    )

    if response.status_code == HTTP_400_BAD_REQUEST:
        # extra debug output (because we're running in a bug that happens very infrequently)
        now = datetime.now(UTC).isoformat()
        logger.error("Failed to exchange token due to HTTP_400_BAD_REQUEST, response: %s", response.content)
        logger.error("    endpoint: %s", settings.token_endpoint_url)
        logger.error("         now: %s", now)
        logger.error("   data sent: %s", json.dumps(data, indent=2))
        # also using print to make sure this message always end up somewhere
        print(  # noqa: T201 this is only for debugging nad should be removed later
            f"Failed to exchange token due to HTTP_400_BAD_REQUEST. DEBUG INFO: \n"
            f"   now: {now} \n"
            f"   data sent: {json.dumps(data, indent=2)} \n"
            f"   response: {response.content}"
        )

    response.raise_for_status()
    try:
        return token_decoder.decode(response.content)
    except ValidationError:
        logger.exception("Failed to decode JWT access token from: %s", response.content)
        raise
