require_relative '../data_conversion'

java_import org.bouncycastle.pqc.crypto.crystals.kyber.KyberParameters
java_import org.bouncycastle.pqc.crypto.crystals.kyber.KyberKeyGenerationParameters
java_import org.bouncycastle.pqc.crypto.crystals.kyber.KyberPublicKeyParameters
java_import org.bouncycastle.pqc.crypto.crystals.kyber.KyberKEMGenerator
java_import org.bouncycastle.pqc.crypto.crystals.kyber.KyberKEMExtractor

require_relative '../data_conversion'

module Ccrypto
  module Java
    
    class CrystalKyberPublicKey < Ccrypto::CrystalKyberPublicKey
      include DataConversion 
      include TR::CondUtils

      def initialize(kp, param)
        super(kp) 
        @keyParam = param
      end

      # Public API
      def key_param
        @keyParam
      end
      alias_method :param, :key_param

      # Public API
      def to_bin
        @native_pubKey.encoded
      end

      # Public API
      def encoded
        to_bin
      end

      def to_s
        "Crystal Kyber - #{@keyParam} : #{native.inspect}"
      end

      # Public API
      def self.to_key(params, bin)
        bin = to_java_bytes(bin) if not bin.is_a?(::Java::byte[])

        case params
        when Ccrypto::CrystalKyberConfig
          pa = params.provider_config[:params]
        when Symbol
          pa = CrystalKyberEngine.supported_configs[params]
          if not_empty?(pa)
            pa = pa.provider_config[:params]
          end
        else
          raise CrystalKyberEngineError, "Unsupported params type '#{params.class}'"
        end

        raise CrystalKyberEngineError, "Unknown param '#{param}'" if is_empty?(pa)

        pubKey = KyberPublicKeyParameters.new(pa, bin)
        CrystalKyberPublicKey.new(pubKey, params)
      end

      def key_equals?(pubKey)
        if not @native_pubKey.nil?
          case pubKey
          when CrystalKyberPublicKey
            @native_pubKey.encoded == pubKey.to_bin
          else
            logger.warn "Unmatched public key : (native) #{@native_pubKey} vs. (subject) #{pubKey}"

            false
          end
        else
          logger.warn "CrystalKyberPublicKey equals? returned false because native_pubKey is nil"
          false
        end
      end
      alias_method :equals?, :key_equals?

    end # class CrystalKyberPublicKey

    class CrystalKyberPrivateKey < Ccrypto::PrivateKey

      def initialize(kp, param)
        super(kp) 
        @keyParam = param
      end

      def key_param
        @keyParam
      end
      alias_method :param, :key_param

      def to_bin
        @native_privKey.encoded
      end

      def self.to_key(params = {})
        param = params[:param]
        bcParam = CrystalKyberEngine.find_config(param)
        bcPrivKey = org.bouncycastle.pqc.crypto.crystals.kyber::KyberPrivateKeyParameters.new(bcParam.provider_config[:params], params[:s].to_java_bytes, params[:hpk].to_java_bytes, params[:nonce].to_java_bytes, params[:t].to_java_bytes, params[:rho].to_java_bytes)
        CrystalKyberPrivateKey.new(bcPrivKey, param)
      end

      def equals?(privKey)
        if not @native_privKey.nil?
          case privKey
          when CrystalKyberPrivateKey
            @native_privKey.encoded == privKey.to_bin
          else
            logger.warn "Unmatched private key : (native) #{@native_privKey} vs. (subject) #{privKey}"
            false
          end
        else
          logger.warn "CrystalKyberPrivateKey equals? returned false because native_privKey is nil"
          false
        end
      end
      alias_method :key_equals?, :equals?

    end # class CrystalKyberPrivateKey

    class CrystalKyberKEM
      include TR::CondUtils
      include DataConversion

      def initialize(kem)
        @kem = kem
        raise KeypairEngineException, "KEM cannot be empty" if is_empty?(@kem)
      end

      def for_recipient
        @kem.getEncapsulation() 
      end

      def for_cipher
        @kem.getSecret()
      end

      def native
        @kem
      end

    end # class CrystalKyberKEM

    # 
    # Crystal-Kyber mainly for encryption via derivation
    #
    class CrystalKyberKeyBundle
      include Ccrypto::KeyBundle
      include TR::CondUtils
      include DataConversion

      def initialize(kp, params)
        @nativeKeypair = kp
        @nativeParam = params
      end

      def enc_derive(pubKey, &block)
        raise KeypairEngineException, "Given public key to derive cannot be empty" if is_empty?(pubKey) 
        raise KeypairEngineException, "Given public key need to be a Crystal Kyber public key" if not pubKey.is_a?(CrystalKyberPublicKey)

        rand = java.security.SecureRandom.getInstanceStrong()
        kemGen = KyberKEMGenerator.new(rand)
        CrystalKyberKEM.new(kemGen.generateEncapsulated(pubKey.native))
      end
      alias_method :derive_enc_shared_secret, :enc_derive

      def dec_derive(kem, &block)
        raise KeypairEngineException, "Native Keypair not available" if is_empty?(@nativeKeypair)

        kemExt = KyberKEMExtractor.new(@nativeKeypair.getPrivate)
        kemExt.extractSecret(to_java_bytes(kem))
      end
      alias_method :derive_dec_shared_secret, :dec_derive

      def public_key
        if @_pubKey.nil?
          @_pubKey = CrystalKyberPublicKey.new(@nativeKeypair.getPublic(), @nativeParam.param)
        end
        @_pubKey
      end

      def private_key
        if @_privKey.nil?
          @_privKey = CrystalKyberPrivateKey.new(@nativeKeypair.getPrivate(), @nativeParam.param)
        end
        @_privKey
      end

    end # class CrystalKyberKeyBundle

    class CrystalKyberEngine

      def self.supported_params
        supported_configs.keys
      end

      def self.supported_configs
        if @supportedConfig.nil?
          @supportedConfig = {}
          {
            kyber512: KyberParameters::kyber512,
            kyber768: KyberParameters::kyber768,
            kyber1024: KyberParameters::kyber1024,
            kyber512_aes: KyberParameters::kyber512_aes,
            kyber768_aes: KyberParameters::kyber768_aes,
            kyber1024_aes: KyberParameters::kyber1024_aes,
          }.each do |k,kp|
            conf = Ccrypto::CrystalKyberConfig.new(k)
            conf.provider_config = { params: kp }
            @supportedConfig[k] = conf
          end
        end 
        @supportedConfig.freeze
      end

       def self.find_config(conf)
        supported_configs.select { |k,v| k == conf }.first[1]
      end

     def self.get_available_session_keysize
        [512, 768, 1024].freeze
      end

      def self.get_session_keysize_param(keysize, with_aes = false)
        case keysize.to_i
        when 512, 768, 1024
          if with_aes
            supported_configs["kyber#{keysize.to_i}_aes".to_sym]
          else
            supported_configs["kyber#{keysize.to_i}".to_sym]
          end
        else
          raise KeypairEngineException, "Crystal Kyber only has keysizes of 512,768 and 1024"
        end
      end

      def initialize(*args, &block)
        @config = args.first
        raise KeypairEngineException, "1st parameter must be a #{Ccrypto::KeypairConfig.class} object" if not @config.is_a?(Ccrypto::KeypairConfig)
      end

      def generate_keypair(&block)
      
        rand = java.security.SecureRandom.getInstanceStrong
        kpg = org.bouncycastle.pqc.crypto.crystals.kyber.KyberKeyPairGenerator.new
        kpg.init(KyberKeyGenerationParameters.new(rand, @config.provider_config[:params]))

        CrystalKyberKeyBundle.new(kpg.generateKeyPair(), @config)

      end

    end # class CrystalKyberEngine

  end
end