from abc import ABC, abstractmethod
import requests
import jwt

from enum import Enum
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from locknessie.settings import ConfigSettings

class AuthType(str, Enum):
    user = "user"
    daemon = "daemon"


class AuthBase(ABC):
    _auth_type: AuthType
    settings: "ConfigSettings"
    auth_roles_claim: str = "groups"
    algorithm: str = "RS256"

    def __init__(self,
                 settings: "ConfigSettings",
                 auth_type: Optional[AuthType] = AuthType.user):
        self._auth_type = auth_type
        self.settings = settings

    @property
    def auth_type(self) -> AuthType:
        return self._auth_type

    @abstractmethod
    def get_token(self) -> str:
        pass

    def validate_token(self, token: Optional[str] = None) -> bool:
        token = token or self.get_token()
        return self._validate_token(token)

    def _validate_token(self, token: str) -> None:
        """ensure the token is correctly generated by the auth provider

        Raises:
            ValueError: if the token is not correctly generated by the auth provider
            InvalidSignatureError: if the token signature is invalid
            InvalidTokenError: if the token is invalid
            InvalidAlgorithmError: if the token algorithm is invalid
        """

        key = self._get_jwk(token)
        token = jwt.decode(token, key=key, algorithms=[self.algorithm], options={"verify_signature": True})
        if self.auth_roles_claim not in token:
            raise ValueError(f"Token does not contain {self.auth_roles_claim} claim, required for {self.auth_type} auth")

    def _get_jwk(self, token: str) -> str:
        """get the correct jwk from the auth provider"""
        header = jwt.get_unverified_header(token)
        try:
            jwks_uri = header["jwks_uri"]
        except KeyError:
            try:
                issuer = jwt.decode(token, algorithms=[self.algorithm], options={"verify_signature": False})["iss"]
            except KeyError:
                raise ValueError(f"Token does not contain {self.auth_type} issuer OR jwks_uri, one is required for {self.auth_type} auth")
            explorer_uri = f"{issuer}/.well-known/openid-configuration"
            server_config = self._get_json_or_raise(requests.get(explorer_uri))
            try:
                jwks_uri = server_config["jwks_uri"]
            except KeyError:
                raise ValueError("Server explorer endpoint does not return a jwks_uri")

        jwks = self._get_json_or_raise(requests.get(jwks_uri))
        for key in jwks["keys"]:
            if key["kid"] == header["kid"]:
                return jwt.algorithms.RSAAlgorithm.from_jwk(key)
        raise ValueError(f"No matching jwk found for token")

    def _get_json_or_raise(self, response: requests.Response) -> dict:
        assert response.status_code == 200, f"Failed to get json from {response.url}, status code: {response.status_code}"
        try:
            return response.json()
        except requests.exceptions.JSONDecodeError as e:
            raise ValueError(f"unable to decode json from {response.url}: {e}. This is likely due to a bad url or a malformed response from the auth provider.")