from __future__ import annotations
from jose import jwt, ExpiredSignatureError, JWTError
from src.application.contracts import ILogger, IJwtService
from src.application.domain.dto import AccessTokenPayload
from src.application.domain.exceptions import ApplicationException,InternalServerException,JwtException
from src.infrastructure.config.settings import settings
from src.infrastructure.vault import JwtKeyStore


class JwtService(IJwtService):
    def __init__(self, logger: ILogger, key_store: JwtKeyStore) -> None:
        self._logger = logger
        self._key_store = key_store

    async def decode_access_token(self, token: str) -> AccessTokenPayload:
        payload = await self._decode_and_verify(token)

        if payload.get('type') != 'access':
            self._logger.warning(f'Access token invalid type received_type={payload.get('type')}')
            raise JwtException(message='Invalid token type')

        try:
            return AccessTokenPayload(
                sub=str(payload['sub']),
                type='access',
                sid=str(payload['sid']),
                iat=int(payload['iat']),
                nbf=int(payload['nbf']),
                exp=int(payload['exp']),
                iss=payload.get('iss'),
                aud=payload.get('aud'),
            )
        except KeyError as exception:
            self._logger.warning(f'Access token missing claim error={str(exception)}')
            raise JwtException(message=f'Missing token claim: {exception}')

    async def _decode_and_verify(self, token: str) -> dict:
        kid: str | None = None
        try:
            header = jwt.get_unverified_header(token)

            kid = header.get('kid')
            if not kid:
                self._logger.warning(f'JWT header missing kid header={header}')
                raise JwtException(message='Missing token header: kid')

            received_alg = header.get('alg')
            if received_alg != settings.JWT_ALGORITHM:
                self._logger.warning(f'JWT invalid algorithm kid={kid} received_alg={received_alg} expected_alg={settings.JWT_ALGORITHM}')
                raise JwtException(message='Invalid token algorithm')

            public_pem = await self._key_store.get_public_key_for_kid(str(kid))

            if not public_pem:
                self._logger.info(f'JWT kid miss kid={kid} forcing keystore refresh')
                await self._key_store.refresh()
                public_pem = await self._key_store.get_public_key_for_kid(str(kid))

            if not public_pem:
                self._logger.warning(f'JWT unknown kid kid={kid}')
                raise JwtException(message='Unknown token kid')

            options = {
                'verify_signature': True,
                'verify_exp': True,
                'verify_nbf': True,
                'verify_iat': True,
                'verify_aud': bool(settings.JWT_AUDIENCE),
                'verify_iss': bool(settings.JWT_ISSUER),
                'require_exp': True,
                'require_iat': True,
                'require_nbf': True,
                'require_sub': True,
                'leeway': 10,
            }

            payload = jwt.decode(
                token,
                public_pem,
                algorithms=[settings.JWT_ALGORITHM],
                audience=settings.JWT_AUDIENCE or None,
                issuer=settings.JWT_ISSUER or None,
                options=options,
            )

            if 'sid' not in payload:
                self._logger.warning(f'JWT missing sid claim kid={kid}')
                raise JwtException(message='Missing token claim: sid')

            if 'type' not in payload:
                self._logger.warning(f'JWT missing type claim kid={kid}')
                raise JwtException(message='Missing token claim: type')

            return payload

        except ExpiredSignatureError as exception:
            self._logger.info(f'JWT expired kid={kid} error={str(exception)}')
            raise JwtException(message='Token expired')

        except ApplicationException:
            raise

        except JWTError as exception:
            self._logger.warning(f'JWT decode failed kid={kid} error={str(exception)}')
            raise JwtException(message='Invalid token')

        except Exception as exception:
            self._logger.error(f'Unexpected JWT decode error kid={kid} error={str(exception)}')
            raise InternalServerException(message='JWT decode failed')