from functools import wraps
import json
import os
from urllib.request import urlopen
from dotenv import find_dotenv, load_dotenv
from jose import jwt
#import jwt
from rest_framework.request import Request
from rest_framework.viewsets import ViewSet
from management.utils.error_handling import GenericError
#load_dotenv(find_dotenv())
from typing import Union


def get_token_auth_header(request: Request):
    """Obtains the Access Token from the Authorization Header
    """
    auth = request._request.headers.get("Authorization", None)
    if not auth:
        raise GenericError({"code": "authorization_header_missing",
                         "description":
                         "Authorization header is expected"}, 401)

    parts = auth.split()

    if parts[0].lower() != "bearer":
        raise GenericError({"code": "invalid_header",
                         "description":
                         "Authorization header must start with"
                         " Bearer"}, 401)
    elif len(parts) == 1:
        raise GenericError({"code": "invalid_header",
                         "description": "Token not found"}, 401)
    elif len(parts) > 2:
        raise GenericError({"code": "invalid_header",
                         "description":
                         "Authorization header must be"
                         " Bearer token"}, 401)

    token = parts[1]
    return token

def requires_auth(permission: Union[str, list] = None):
    """Determines if the Access Token is valid
    """
    def auth_decorator(f):

        @wraps(f)
        def decorated(view, request, *args, **kwargs):

            tenant_id = os.getenv('TENANT_ID')
            client_id = os.getenv('CLIENT_ID')

            try:

                token = get_token_auth_header(request)
              
                jsonurl = urlopen("https://login.microsoftonline.com/" + tenant_id + "/discovery/v2.0/keys")
                
                jwks = json.loads(jsonurl.read())
                
                unverified_header = jwt.get_unverified_header(token)

                rsa_key = {}
                for key in jwks["keys"]:
                    if key["kid"] == unverified_header["kid"]:
                        rsa_key = {
                            "kty": key["kty"],
                            "kid": key["kid"],
                            "use": key["use"],
                            "n": key["n"],
                            "e": key["e"]
                        }
            except Exception:
                raise GenericError({"code": "invalid_header",
                                    "description":
                                    "Unable to parse authentication"
                                    " token."}, 401)
                
            if rsa_key:
                payload = None
                try:
                    payload = jwt.decode(
                        token,
                        rsa_key,
                        algorithms=["RS256"],
                        audience=client_id,
                        issuer=f"https://login.microsoftonline.com/{tenant_id}/v2.0"
                    )
                    request.user_id = payload['oid']

                except jwt.ExpiredSignatureError:
                    
                    raise GenericError({"code": "token_expired",
                                    "description": "token is expired"}, 401)
                except jwt.JWTClaimsError:
                    raise GenericError({"code": "invalid_claims",
                                    "description":
                                    "incorrect claims,"
                                    "please check the audience and issuer"}, 401)
                except Exception:
                    raise GenericError({"code": "invalid_header",
                                    "description":
                                    "Unable to parse authentication"
                                    " token."}, 401)
                
                if isinstance(permission, list):
                    roles = payload.get("roles", [])
                    # for value in permission:
                    #     print(f"Checking permission value: {value}")
                    #     print(f"Is '{value}' in roles? {value in roles}")
                    
                    if any(value in roles for value in permission):
                        return f(view, request, *args, **kwargs)
                    
                    else:                        
                        raise GenericError({"code": "forbidden",
                            "description": "No valid permission for access the resource"}, 403)
                else:    
                    if permission not in payload["roles"]:
                   
                        raise GenericError({"code": "forbidden",
                            "description": "No valid permission for access the resource"}, 403)
                    
                    return f(view, request, *args, **kwargs)
            raise GenericError({"code": "invalid_header",
                            "description": "Unable to find appropriate key"}, 401)
        return decorated
    return auth_decorator