import logging
import sys
import signal
from json import dumps as json_dumps

from joserfc.jwk import RSAKey, ECKey

from .classes import DeviceProperties
from .config import KEY_DIR
from .util import split_key_type

log = logging.getLogger(__name__)


class TimeoutException(Exception):
    pass

def timeout_handler(signum, frame):
    raise TimeoutException()

class JWK:
    '''Working with JWK(S)'''

    def __init__(self, device_properties: 'DeviceProperties', key_dir=KEY_DIR):
        self.key_security = device_properties.key_security
        self.key_type = device_properties.key_type
        self.key_dir = key_dir
        self.jwk_path = key_dir + "/" + "public_jwk.json"
        self.public_key_path = key_dir + "/" + "public.pem"
        '''The key is generated automatically'''
        self.key = self._generate()

    def _generate(self):
        key_cryptography = split_key_type(self.key_type)[0]

        with open(self.public_key_path, 'rb', buffering=0) as reader:
            file_content = reader.read()
            
            try:
                signal.signal(signal.SIGALRM, timeout_handler)
                signal.alarm(10)  # Set a 10-second timeout

                if key_cryptography == 'RSA':
                    print('Importing RSA key')
                    key = RSAKey.import_key(file_content)
                elif key_cryptography == 'ECC':
                    print('Importing ECC key')
                    key = ECKey.import_key(file_content)
                else:
                    raise Exception(
                        'Unexpected key cryptography: {}. Supported: RSA, ECC'.format(
                            key_cryptography
                        )
                    )
                signal.alarm(0)  # Disable the alarm

                jwk = key.as_dict()

                if self.key_type == 'RSA:2048':
                    jwk['alg'] = 'RS256'
                elif self.key_type == 'ECC:p256':
                    jwk['alg'] = 'ES256'
                elif self.key_type == 'ECC:p384':
                    jwk['alg'] = 'ES384'
                elif self.key_type == 'ECC:p521':
                    jwk['alg'] = 'ES521'
                else:
                    log.error('Unsupported key type.')
                    sys.exit(1)
                    
                print('Generated JWK')
                return jwk

            except TimeoutException:
                log.error('Importing RSA key timed out.')
                return None
            except Exception as e:
                log.error('Failed to import key: %s', e)
                return None

    def save_to_key_path(self) -> bool:
        '''Generates and saves JWK to the default path'''
        formatted_key = bytearray(json_dumps(
            self.key, ensure_ascii=True).encode('utf8'))
        try:
            with open(self.jwk_path, 'wb') as writer:
                writer.write(formatted_key)
            return True
        except:
            log.error('Could not save JWK to {}'.format(self.jwk_path))
            sys.exit(1)

    def get(self):
        '''Generates and returns JWK'''
        return self.key

    def get_jwks(self) -> dict:
        '''Generates JWK and returns JWKS'''
        jwks = {
            'keys': [self.key]
        }
        return jwks
