package com.alttd.altitudeweb.controllers.login;

import com.alttd.altitudeweb.database.Databases;
import com.alttd.altitudeweb.database.web_db.KeyPairEntity;
import com.alttd.altitudeweb.database.web_db.KeyPairMapper;
import com.alttd.altitudeweb.setup.Connection;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.time.Instant;
import java.util.Base64;
import java.util.concurrent.CompletableFuture;

@Service
@RequiredArgsConstructor
@Slf4j
public class KeyPairService {

    private KeyPair cachedKeyPair = null;
    private static final String RSA_ALGORITHM = "RSA";
    private static final int RSA_KEY_SIZE = 2048;

    public KeyPair getJwtSigningKeyPair() {
        if (cachedKeyPair != null) {
            return cachedKeyPair;
        }
        KeyPair keyPair = getOrCreateKeyPair();
        if (keyPair != null) {
            cachedKeyPair = keyPair;
            return cachedKeyPair;
        }
        throw new IllegalStateException("Failed to generate or load key pair");
    }

    public KeyPair getOrCreateKeyPair() {
        CompletableFuture<KeyPairEntity> keyPairFuture = new CompletableFuture<>();
        Connection.getConnection(Databases.DEFAULT)
                .runQuery(sqlSession -> {
                    log.debug("Loading key pair");
                    try {
                        KeyPairEntity entity = sqlSession.getMapper(KeyPairMapper.class).getKeyPair();
                        keyPairFuture.complete(entity);
                    } catch (Exception e) {
                        log.error("Failed to key pair", e);
                        keyPairFuture.completeExceptionally(e);
                    }
                });
        KeyPairEntity keyPairEntity = keyPairFuture.join();
        if (keyPairEntity != null) {
            try {
                byte[] privateKeyBytes = Base64.getDecoder().decode(keyPairEntity.getPrivateKey());
                byte[] publicKeyBytes = Base64.getDecoder().decode(keyPairEntity.getPublicKey());

                KeyFactory keyFactory = KeyFactory.getInstance(RSA_ALGORITHM);
                PrivateKey privateKey = keyFactory.generatePrivate(new PKCS8EncodedKeySpec(privateKeyBytes));
                PublicKey publicKey = keyFactory.generatePublic(new X509EncodedKeySpec(publicKeyBytes));

                return new KeyPair(publicKey, privateKey);
            } catch (Exception e) {
                log.error("Failed to load key pair from database", e);
            }
        }

        KeyPair keyPair = generateKeyPair();

        try {
            KeyPairEntity entity = new KeyPairEntity();
            entity.setPrivateKey(Base64.getEncoder().encodeToString(keyPair.getPrivate().getEncoded()));
            entity.setPublicKey(Base64.getEncoder().encodeToString(keyPair.getPublic().getEncoded()));
            entity.setCreatedAt(Instant.now());

            Connection.getConnection(Databases.DEFAULT)
                    .runQuery(sqlSession -> {
                        log.debug("Saving key pair");
                        try {
                            sqlSession.getMapper(KeyPairMapper.class).save(entity);
                            log.info("Generated and saved new key pair");
                        } catch (Exception e) {
                            log.error("Failed to key pair", e);
                        }
                    });
        } catch (Exception e) {
            log.error("Failed to save key pair to database", e);
        }

        return keyPair;
    }

    private KeyPair generateKeyPair() {
        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(RSA_ALGORITHM);
            keyPairGenerator.initialize(RSA_KEY_SIZE);
            return keyPairGenerator.generateKeyPair();
        } catch (NoSuchAlgorithmException e) {
            throw new IllegalStateException("Error generating key pair", e);
        }
    }
}
