import os
import json
import time
from datetime import datetime, timedelta

from authlib.jose import jwt

# Re-use key loading from grant module
from ..models.grant import read_private_key, import_key


class JWKSManager:
    def get_signing_key(self):
        """Load the signing key. Priority: LOCAL_OIDC_JSON_URL → LOCAL_OIDC_PEM_URL → OIDC_JSON_URL"""
        local_json_file = os.environ.get('LOCAL_OIDC_JSON_URL', None)
        local_pem_file = os.environ.get('LOCAL_OIDC_PEM_URL', None)
        oidc_json_url = os.environ.get('OIDC_JSON_URL', None)

        if local_json_file is not None:
            try:
                fp = os.getcwd() + "/" + local_json_file
                key = import_key(fp, True)
                if key is not None:
                    return key
            except Exception as e:
                print(f'JWKSManager: failed to load local JSON key: {e}')

        if local_pem_file is not None:
            try:
                fp = os.getcwd() + "/" + local_pem_file
                key = import_key(fp, False)
                if key is not None:
                    return key
            except Exception as e:
                print(f'JWKSManager: failed to load local PEM key: {e}')

        if oidc_json_url is not None:
            import requests
            try:
                resp = requests.get(oidc_json_url)
                if resp.status_code == 200:
                    return resp.json()
            except Exception as e:
                print(f'JWKSManager: failed to fetch OIDC JSON URL: {e}')

        return read_private_key()

    def get_public_jwks(self) -> dict:
        """Return the public portion of the JWKS."""
        key = self.get_signing_key()
        if key is None:
            return {"keys": []}
        try:
            from joserfc.jwk import RSAKey
            if isinstance(key, dict) and 'kty' in key:
                rsa = RSAKey.import_key(key)
                public = rsa.as_dict(private=False)
                return {"keys": [public]}
            elif isinstance(key, dict) and 'keys' in key:
                public_keys = []
                for k in key['keys']:
                    try:
                        rsa = RSAKey.import_key(k)
                        public_keys.append(rsa.as_dict(private=False))
                    except Exception:
                        pass
                return {"keys": public_keys}
        except Exception as e:
            print(f'JWKSManager: failed to extract public key: {e}')
        return {"keys": []}

    def generate_id_token(self, client, user, scope: str, nonce: str = None,
                          auth_time: int = None, expires_in: int = 3600) -> str:
        """Generate an OIDC ID token (ES256/RS256 JWT)."""
        key = self.get_signing_key()
        if key is None:
            raise ValueError('No signing key available for ID token generation')
        now = int(time.time())
        issuer = os.environ.get('OAUTH2_ISSUER', os.environ.get('JWT_ISSUER', ''))
        payload = {
            'iss': issuer,
            'sub': str(getattr(user, 'uid', getattr(user, 'id', ''))),
            'aud': client.client_id if hasattr(client, 'client_id') else str(client),
            'iat': now,
            'exp': now + expires_in,
        }
        if nonce:
            payload['nonce'] = nonce
        if auth_time:
            payload['auth_time'] = auth_time
        # Include basic user claims when scope includes openid/profile/email
        if hasattr(user, 'email') and 'email' in scope:
            payload['email'] = user.email
        if hasattr(user, 'displayName') and 'profile' in scope:
            payload['name'] = user.displayName
        alg = os.environ.get('OIDC_JWT_ALG', 'RS256')
        header = {'alg': alg}
        token = jwt.encode(header, payload, key, check=False)
        if isinstance(token, bytes):
            return token.decode('utf-8')
        return token

    def openid_configuration(self, issuer: str = None) -> dict:
        """Return the OpenID Connect discovery document."""
        if issuer is None:
            issuer = os.environ.get('OAUTH2_ISSUER', os.environ.get('JWT_ISSUER', ''))
        return {
            "issuer": issuer,
            "authorization_endpoint": issuer + "/o/v1/authorize",
            "token_endpoint": issuer + "/o/v1/token",
            "userinfo_endpoint": issuer + "/o/v1/userinfo",
            "jwks_uri": issuer + "/.well-known/jwks.json",
            "revocation_endpoint": issuer + "/o/v1/revoke",
            "introspection_endpoint": issuer + "/o/v1/introspect",
            "device_authorization_endpoint": issuer + "/o/v1/device",
            "response_types_supported": ["code", "token", "id_token", "code token", "code id_token", "id_token token"],
            "grant_types_supported": [
                "authorization_code", "implicit", "client_credentials",
                "refresh_token", "urn:ietf:params:oauth:grant-type:device_code"
            ],
            "subject_types_supported": ["public"],
            "id_token_signing_alg_values_supported": ["RS256", "ES256"],
            "scopes_supported": ["openid", "profile", "email"],
            "token_endpoint_auth_methods_supported": ["client_secret_basic", "client_secret_post", "none"],
            "claims_supported": ["sub", "iss", "aud", "exp", "iat", "email", "name", "nonce"],
            "code_challenge_methods_supported": ["S256"],
        }


jwks_manager = JWKSManager()
