import base64
from enum import Enum

from cryptojwt.jwk.jwk import key_from_jwk_dict
from pydantic import BaseModel, ValidationError
from satosa.context import Context

from pyeudiw.jwt.exceptions import JWSVerificationError
from pyeudiw.jwt.jws_helper import JWSHelper
from pyeudiw.oauth2.dpop.verifier import DPoPVerifier
from pyeudiw.satosa.exceptions import InvalidRequestException
from pyeudiw.satosa.frontends.openid4vci.endpoints.vci_base_endpoint import VCIBaseEndpoint, POST_ACCEPTED_METHODS
from pyeudiw.satosa.frontends.openid4vci.models.openid4vci_basemodel import CONFIG_CTX
from pyeudiw.satosa.frontends.openid4vci.models.token import AccessToken, RefreshToken
from pyeudiw.satosa.frontends.openid4vci.models.token_request import TokenRequest, REDIRECT_URI_CTX, CODE_CHALLENGE_CTX, CODE_CHALLENGE_METHOD_CTX, SCOPE_CTX
from pyeudiw.satosa.frontends.openid4vci.models.token_response import TokenResponse
from pyeudiw.satosa.frontends.openid4vci.storage.engine import OpenId4VciDBEngineHandler
from pyeudiw.satosa.frontends.openid4vci.storage.entity import OpenId4VCIEntity
from pyeudiw.satosa.frontends.openid4vci.tools.exceptions import InvalidScopeException
from pyeudiw.satosa.utils.session import get_session_id
from pyeudiw.satosa.utils.validation import (
    validate_content_type,
    validate_request_method,
    validate_oauth_client_attestation,
    validate_oauth_client_attestation_pop,
    OAUTH_CLIENT_ATTESTATION_POP_HEADER,
)
from pyeudiw.tools.content_type import HTTP_CONTENT_TYPE_HEADER, FORM_URLENCODED
from pyeudiw.tools.utils import iat_now


class TokenTypsEnum(Enum):
    REFRESH_TOKEN_TYP = "rt+jwt"  # nosec B105
    ACCESS_TOKEN_TYP = "at+jwt"  # nosec B105


class TokenHandler(VCIBaseEndpoint):

    def __init__(self, config: dict, internal_attributes: dict[str, dict[str, str | list[str]]], base_url: str, name: str, *args):
        """
        Initialize the token endpoint class.

        Args:
            config (dict): The configuration dictionary.
            internal_attributes (dict): The internal attributes config.
            base_url (str): The base URL of the service.
            name (str): The name of the SATOSA module to append to the URL.
        """

        super().__init__(config, internal_attributes, base_url, name)
        self.jws_helper = JWSHelper(self.config["metadata_jwks"])
        self.db_engine = OpenId4VciDBEngineHandler(config).db_engine

    def endpoint(self, context: Context):
        """
        Handle a POST request to the token endpoint.

        Args:
            context (Context): The SATOSA context.

        Returns:
            A Response object.
        """

        try:
            validate_request_method(context.request_method, POST_ACCEPTED_METHODS)
            validate_content_type(context.http_headers[HTTP_CONTENT_TYPE_HEADER], FORM_URLENCODED)

            if self.wallet_attestation_required:
                try:
                    validate_oauth_client_attestation_pop(context)
                    validate_oauth_client_attestation(context, self.dpop_signing_alg_values_supported)
                except InvalidRequestException as e:
                    self._log_error(e.__class__.__name__, f"Error during OAuth client attestation validation in `par` endpoint: {e}")
                    return self._handle_400(context, str(e), e)

            dpop_verifier = None
            if self.dpop_required:
                if not context.http_headers or ("DPoP" not in context.http_headers):
                    raise InvalidRequestException("Missing DPoP header")

                dpop = context.http_headers.get("DPoP")

                try:
                    dpop_verifier = DPoPVerifier(http_header_dpop=dpop)
                    if not dpop_verifier.is_valid:
                        raise InvalidRequestException("Invalid DPoP proof")
                except ValueError as e:
                    self._log_error(e.__class__.__name__, f"Error during DPoP validation in `token` endpoint: {e}")
                    return self._handle_400(context, str(e), e)

            oauth_client_attestation = self._get_oauth_client_attestation(context, self.wallet_attestation_required)
            if oauth_client_attestation:
                self.jws_helper.verify(oauth_client_attestation)

            entity = self.db_engine.get_by_session_id(get_session_id(context))

            if not entity:
                raise InvalidRequestException("session not found")

            vci_entity = OpenId4VCIEntity(**entity)

            TokenRequest.model_validate(
                self._get_body(context),
                context={
                    CONFIG_CTX: self.config,
                    REDIRECT_URI_CTX: vci_entity.redirect_uri,
                    CODE_CHALLENGE_METHOD_CTX: vci_entity.code_challenge_method,
                    CODE_CHALLENGE_CTX: vci_entity.code_challenge,
                    SCOPE_CTX: vci_entity.scope,
                },
            )
            iat = iat_now()
            authorization_details = vci_entity.authorization_details
            if authorization_details or len(authorization_details) > 0:
                for ad in authorization_details:
                    ad.credential_identifiers = self.config_utils.get_credential_configurations_supported(ad.credential_configuration_id).scope

            cnf = self._build_dpop_cnf(dpop_verifier) if dpop_verifier else {}
            return TokenResponse.to_created_response(
                self._to_token(iat, entity, TokenTypsEnum.ACCESS_TOKEN_TYP, cnf),
                self._to_token(iat, entity, TokenTypsEnum.REFRESH_TOKEN_TYP, cnf),
                iat + self.config_utils.get_jwt().access_token_exp,
                authorization_details,
            )
        except (InvalidRequestException, InvalidScopeException, JWSVerificationError, ValidationError, TypeError) as e:
            return self._handle_400(context, self._handle_validate_request_error(e, "token"), e)
        except Exception as e:
            self._log_error(e.__class__.__name__, f"Error during invoke token endpoint: {e}")
            return self._handle_500(context, "error during invoke token endpoint", e)

    def _build_dpop_cnf(self, dpop_verifier: DPoPVerifier) -> dict:
        """Build cnf with jkt (RFC 9449) to bind token to DPoP key."""
        jwk = key_from_jwk_dict(dpop_verifier.public_jwk)
        thumbprint = jwk.thumbprint("SHA-256")
        jkt = base64.urlsafe_b64encode(thumbprint).rstrip(b"=").decode()
        return {"jkt": jkt}

    def _to_token(self, iat: int, entity: OpenId4VCIEntity, typ: TokenTypsEnum, cnf: dict = None) -> str:

        if isinstance(entity, dict):
            entity = OpenId4VCIEntity(**entity)
        cnf = cnf or {}

        match typ:
            case TokenTypsEnum.ACCESS_TOKEN_TYP:
                exp = iat + self.config_utils.get_jwt().access_token_exp
            case TokenTypsEnum.REFRESH_TOKEN_TYP:
                exp = iat + self.config_utils.get_jwt().refresh_token_exp
            case _:
                self._log_error(self.__class__.__name__, f"unexpected typ {typ} for token ")
                raise Exception(f"Invalid token typ {typ}")

        token = AccessToken(
            iss=self.entity_id,
            aud=self.entity_id,
            exp=exp,
            iat=iat,
            client_id=entity.client_id,
            sub=entity.client_id,
            cnf=cnf,
        )
        if typ == TokenTypsEnum.REFRESH_TOKEN_TYP:
            token = RefreshToken(**token.model_dump())

        return self._sign_token(token, typ.value)

    def _sign_token(self, token: BaseModel, typ: str) -> str:
        jws_headers = {
            "typ": typ,
        }
        return self.jws_helper.sign(protected=jws_headers, plain_dict=token.model_dump())

    @staticmethod
    def _get_oauth_client_attestation(context: Context, required: bool = True) -> str | None:
        """
        Retrieve oauth client attestation pop header
        """

        if (
            not context.http_headers
            or (OAUTH_CLIENT_ATTESTATION_POP_HEADER not in context.http_headers)
            or (context.http_headers.get(OAUTH_CLIENT_ATTESTATION_POP_HEADER) is None)
        ):
            if required:
                raise InvalidRequestException("Missing OAuth-Client-Attestation header")
            else:
                return None
        return context.http_headers.get(OAUTH_CLIENT_ATTESTATION_POP_HEADER)

    def _validate_configs(self):
        self._validate_required_configs(
            [
                ("jwt.access_token_exp", self.config_utils.get_jwt().access_token_exp),
                ("jwt.refresh_token_exp", self.config_utils.get_jwt().refresh_token_exp),
            ]
        )
        oauth_authorization_server = self.config_utils.get_oauth_authorization_server()
        if not oauth_authorization_server:
            self._validate_required_configs(
                [
                    ("metadata.oauth_authorization_server", oauth_authorization_server),
                ]
            )
        self._validate_required_configs(
            [
                ("metadata.oauth_authorization_server.scopes_supported", oauth_authorization_server.scopes_supported),
            ]
        )
