/*
    * GordianKnot: Security Suite
    * Copyright 2012-2026. Tony Washer
    *
    * Licensed under the Apache License, Version 2.0 (the "License"); you may not
    * use this file except in compliance with the License.  You may obtain a copy
    * of the License at
    *
    *   http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
   * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
   * License for the specific language governing permissions and limitations under
   * the License.
   */
  package io.github.tonywasher.joceanus.gordianknot.impl.bc;
  
  import io.github.tonywasher.joceanus.gordianknot.api.base.GordianException;
  import io.github.tonywasher.joceanus.gordianknot.api.keypair.GordianKeyPair;
  import io.github.tonywasher.joceanus.gordianknot.api.keypair.GordianKeyPairSpec;
  import io.github.tonywasher.joceanus.gordianknot.api.sign.GordianSignParams;
  import io.github.tonywasher.joceanus.gordianknot.api.sign.GordianSignatureSpec;
  import io.github.tonywasher.joceanus.gordianknot.impl.bc.BouncyKeyPair.BouncyPrivateKey;
  import io.github.tonywasher.joceanus.gordianknot.impl.bc.BouncyKeyPair.BouncyPublicKey;
  import io.github.tonywasher.joceanus.gordianknot.impl.core.base.GordianBaseFactory;
  import io.github.tonywasher.joceanus.gordianknot.impl.core.exc.GordianCryptoException;
  import io.github.tonywasher.joceanus.gordianknot.impl.core.keypair.GordianKeyPairValidity;
  import io.github.tonywasher.joceanus.gordianknot.impl.core.sign.GordianCoreSignature;
  import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
  import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
  import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
  import org.bouncycastle.crypto.CipherParameters;
  import org.bouncycastle.crypto.CryptoException;
  import org.bouncycastle.crypto.Signer;
  import org.bouncycastle.crypto.params.AsymmetricKeyParameter;
  import org.bouncycastle.crypto.params.ParametersWithContext;
  import org.bouncycastle.crypto.params.ParametersWithRandom;
  import org.bouncycastle.pqc.crypto.mldsa.HashMLDSASigner;
  import org.bouncycastle.pqc.crypto.mldsa.MLDSAKeyGenerationParameters;
  import org.bouncycastle.pqc.crypto.mldsa.MLDSAKeyPairGenerator;
  import org.bouncycastle.pqc.crypto.mldsa.MLDSAParameters;
  import org.bouncycastle.pqc.crypto.mldsa.MLDSAPrivateKeyParameters;
  import org.bouncycastle.pqc.crypto.mldsa.MLDSAPublicKeyParameters;
  import org.bouncycastle.pqc.crypto.mldsa.MLDSASigner;
  import org.bouncycastle.pqc.crypto.util.PrivateKeyFactory;
  import org.bouncycastle.pqc.crypto.util.PrivateKeyInfoFactory;
  import org.bouncycastle.pqc.crypto.util.PublicKeyFactory;
  import org.bouncycastle.pqc.crypto.util.SubjectPublicKeyInfoFactory;
  
  import java.io.IOException;
  import java.security.spec.PKCS8EncodedKeySpec;
  import java.security.spec.X509EncodedKeySpec;
  import java.util.Arrays;
  
  /**
   * MLDSA KeyPair classes.
   */
  public final class BouncyMLDSAKeyPair {
      /**
       * Private constructor.
       */
      private BouncyMLDSAKeyPair() {
      }
  
      /**
       * Bouncy MLDSA PublicKey.
       */
      public static class BouncyMLDSAPublicKey
              extends BouncyPublicKey<MLDSAPublicKeyParameters> {
          /**
           * Constructor.
           *
           * @param pKeySpec   the keySpec
           * @param pPublicKey the public key
           */
          BouncyMLDSAPublicKey(final GordianKeyPairSpec pKeySpec,
                               final MLDSAPublicKeyParameters pPublicKey) {
              super(pKeySpec, pPublicKey);
          }
  
          @Override
          protected boolean matchKey(final AsymmetricKeyParameter pThat) {
              /* Access keys */
              final MLDSAPublicKeyParameters myThis = getPublicKey();
              final MLDSAPublicKeyParameters myThat = (MLDSAPublicKeyParameters) pThat;
  
              /* Compare keys */
              return compareKeys(myThis, myThat);
          }
  
          /**
           * CompareKeys.
           *
           * @param pFirst  the first key
           * @param pSecond the second key
           * @return true/false
           */
          private static boolean compareKeys(final MLDSAPublicKeyParameters pFirst,
                                            final MLDSAPublicKeyParameters pSecond) {
             return Arrays.equals(pFirst.getEncoded(), pSecond.getEncoded());
         }
     }
 
     /**
      * Bouncy MLDSA PrivateKey.
      */
     public static class BouncyMLDSAPrivateKey
             extends BouncyPrivateKey<MLDSAPrivateKeyParameters> {
         /**
          * Constructor.
          *
          * @param pKeySpec    the keySpec
          * @param pPrivateKey the private key
          */
         BouncyMLDSAPrivateKey(final GordianKeyPairSpec pKeySpec,
                               final MLDSAPrivateKeyParameters pPrivateKey) {
             super(pKeySpec, pPrivateKey);
         }
 
 
         @Override
         protected boolean matchKey(final AsymmetricKeyParameter pThat) {
             /* Access keys */
             final MLDSAPrivateKeyParameters myThis = getPrivateKey();
             final MLDSAPrivateKeyParameters myThat = (MLDSAPrivateKeyParameters) pThat;
 
             /* Compare keys */
             return compareKeys(myThis, myThat);
         }
 
         /**
          * CompareKeys.
          *
          * @param pFirst  the first key
          * @param pSecond the second key
          * @return true/false
          */
         private static boolean compareKeys(final MLDSAPrivateKeyParameters pFirst,
                                            final MLDSAPrivateKeyParameters pSecond) {
             return Arrays.equals(pFirst.getEncoded(), pSecond.getEncoded());
         }
     }
 
     /**
      * BouncyCastle MLDSA KeyPair generator.
      */
     public static class BouncyMLDSAKeyPairGenerator
             extends BouncyKeyPairGenerator {
         /**
          * Generator.
          */
         private final MLDSAKeyPairGenerator theGenerator;
 
         /**
          * Constructor.
          *
          * @param pFactory the Security Factory
          * @param pKeySpec the keySpec
          */
         BouncyMLDSAKeyPairGenerator(final GordianBaseFactory pFactory,
                                     final GordianKeyPairSpec pKeySpec) {
             /* Initialise underlying class */
             super(pFactory, pKeySpec);
 
             /* Determine the parameters */
             final MLDSAParameters myParms = pKeySpec.getMLDSAKeySpec().getParameters();
 
             /* Create and initialise the generator */
             theGenerator = new MLDSAKeyPairGenerator();
             final MLDSAKeyGenerationParameters myParams = new MLDSAKeyGenerationParameters(getRandom(), myParms);
             theGenerator.init(myParams);
         }
 
         @Override
         public BouncyKeyPair generateKeyPair() {
             /* Generate and return the keyPair */
             final AsymmetricCipherKeyPair myPair = theGenerator.generateKeyPair();
             final BouncyMLDSAPublicKey myPublic = new BouncyMLDSAPublicKey(getKeySpec(), (MLDSAPublicKeyParameters) myPair.getPublic());
             final BouncyMLDSAPrivateKey myPrivate = new BouncyMLDSAPrivateKey(getKeySpec(), (MLDSAPrivateKeyParameters) myPair.getPrivate());
             return new BouncyKeyPair(myPublic, myPrivate);
         }
 
         @Override
         public PKCS8EncodedKeySpec getPKCS8Encoding(final GordianKeyPair pKeyPair) throws GordianException {
             /* Protect against exceptions */
             try {
                 /* Check the keyPair type and keySpecs */
                 BouncyKeyPair.checkKeyPair(pKeyPair, getKeySpec());
 
                 /* build and return the encoding */
                 final BouncyMLDSAPrivateKey myPrivateKey = (BouncyMLDSAPrivateKey) getPrivateKey(pKeyPair);
                 final MLDSAPrivateKeyParameters myParms = myPrivateKey.getPrivateKey();
                 final PrivateKeyInfo myInfo = PrivateKeyInfoFactory.createPrivateKeyInfo(myParms, null);
                 return new PKCS8EncodedKeySpec(myInfo.getEncoded());
 
             } catch (IOException e) {
                 throw new GordianCryptoException(ERROR_PARSE, e);
             }
         }
 
         @Override
         public BouncyKeyPair deriveKeyPair(final X509EncodedKeySpec pPublicKey,
                                            final PKCS8EncodedKeySpec pPrivateKey) throws GordianException {
             /* Protect against exceptions */
             try {
                 /* Check the keySpecs */
                 checkKeySpec(pPrivateKey);
 
                 /* derive keyPair */
                 final BouncyMLDSAPublicKey myPublic = derivePublicKey(pPublicKey);
                 final PrivateKeyInfo myInfo = PrivateKeyInfo.getInstance(pPrivateKey.getEncoded());
                 final MLDSAPrivateKeyParameters myParms = (MLDSAPrivateKeyParameters) PrivateKeyFactory.createKey(myInfo);
                 final BouncyMLDSAPrivateKey myPrivate = new BouncyMLDSAPrivateKey(getKeySpec(), myParms);
                 final BouncyKeyPair myPair = new BouncyKeyPair(myPublic, myPrivate);
 
                 /* Check that we have a matching pair */
                 GordianKeyPairValidity.checkValidity(getFactory(), myPair);
 
                 /* Return the keyPair */
                 return myPair;
 
             } catch (IOException e) {
                 throw new GordianCryptoException(ERROR_PARSE, e);
             }
         }
 
         @Override
         public X509EncodedKeySpec getX509Encoding(final GordianKeyPair pKeyPair) throws GordianException {
             /* Protect against exceptions */
             try {
                 /* Check the keyPair type and keySpecs */
                 BouncyKeyPair.checkKeyPair(pKeyPair, getKeySpec());
 
                 /* build and return the encoding */
                 final BouncyMLDSAPublicKey myPublicKey = (BouncyMLDSAPublicKey) getPublicKey(pKeyPair);
                 final MLDSAPublicKeyParameters myParms = myPublicKey.getPublicKey();
                 final SubjectPublicKeyInfo myInfo = SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(myParms);
                 return new X509EncodedKeySpec(myInfo.getEncoded());
 
             } catch (IOException e) {
                 throw new GordianCryptoException(ERROR_PARSE, e);
             }
         }
 
         @Override
         public BouncyKeyPair derivePublicOnlyKeyPair(final X509EncodedKeySpec pEncodedKey) throws GordianException {
             final BouncyMLDSAPublicKey myPublic = derivePublicKey(pEncodedKey);
             return new BouncyKeyPair(myPublic);
         }
 
         /**
          * Derive public key from encoded.
          *
          * @param pEncodedKey the encoded key
          * @return the public key
          * @throws GordianException on error
          */
         private BouncyMLDSAPublicKey derivePublicKey(final X509EncodedKeySpec pEncodedKey) throws GordianException {
             /* Protect against exceptions */
             try {
                 /* Check the keySpecs */
                 checkKeySpec(pEncodedKey);
 
                 /* derive publicKey */
                 final SubjectPublicKeyInfo myInfo = SubjectPublicKeyInfo.getInstance(pEncodedKey.getEncoded());
                 final MLDSAPublicKeyParameters myParms = (MLDSAPublicKeyParameters) PublicKeyFactory.createKey(myInfo);
                 return new BouncyMLDSAPublicKey(getKeySpec(), myParms);
 
             } catch (IOException e) {
                 throw new GordianCryptoException(ERROR_PARSE, e);
             }
         }
     }
 
     /**
      * MLDSA signer.
      */
     public static class BouncyMLDSASignature
             extends GordianCoreSignature {
         /**
          * The MLDSA Signer.
          */
         private Signer theSigner;
 
         /**
          * Constructor.
          *
          * @param pFactory the factory
          * @param pSpec    the signatureSpec.
          */
         BouncyMLDSASignature(final GordianBaseFactory pFactory,
                              final GordianSignatureSpec pSpec) {
             /* Initialise underlying class */
             super(pFactory, pSpec);
         }
 
         /**
          * Create the signer according to the keyPair.
          *
          * @param pKeyPair the keyPair
          * @return the signer
          */
         private static Signer createSigner(final GordianKeyPair pKeyPair) {
             /* Determine whether this is a hashSigner */
             final boolean isHash = pKeyPair.getKeyPairSpec().getMLDSAKeySpec().isHash();
 
             /* Create the internal digests */
             return isHash
                     ? new HashMLDSASigner()
                     : new MLDSASigner();
         }
 
         @Override
         public void initForSigning(final GordianSignParams pParams) throws GordianException {
             /* Initialise detail */
             super.initForSigning(pParams);
             final BouncyKeyPair myPair = getKeyPair();
             final byte[] myContext = getContext();
             BouncyKeyPair.checkKeyPair(myPair);
 
             /* Initialise and set the signer */
             theSigner = createSigner(myPair);
             final BouncyMLDSAPrivateKey myPrivate = (BouncyMLDSAPrivateKey) myPair.getPrivateKey();
             CipherParameters myParms = new ParametersWithRandom(myPrivate.getPrivateKey(), getRandom());
             if (myContext != null) {
                 myParms = new ParametersWithContext(myParms, myContext);
             }
             theSigner.init(true, myParms);
         }
 
         @Override
         public void initForVerify(final GordianSignParams pParams) throws GordianException {
             /* Initialise detail */
             super.initForVerify(pParams);
             final BouncyKeyPair myPair = getKeyPair();
             final byte[] myContext = getContext();
             BouncyKeyPair.checkKeyPair(myPair);
 
             /* Initialise and set the signer */
             theSigner = createSigner(myPair);
             final BouncyMLDSAPublicKey myPublic = (BouncyMLDSAPublicKey) myPair.getPublicKey();
             CipherParameters myParms = myPublic.getPublicKey();
             if (myContext != null) {
                 myParms = new ParametersWithContext(myParms, myContext);
             }
             theSigner.init(false, myParms);
         }
 
         @Override
         public void update(final byte[] pBytes,
                            final int pOffset,
                            final int pLength) {
             theSigner.update(pBytes, pOffset, pLength);
         }
 
         @Override
         public void update(final byte pByte) {
             theSigner.update(pByte);
         }
 
         @Override
         public void update(final byte[] pBytes) {
             theSigner.update(pBytes, 0, pBytes.length);
         }
 
         @Override
         public void reset() {
             theSigner.reset();
         }
 
         @Override
         protected BouncyKeyPair getKeyPair() {
             return (BouncyKeyPair) super.getKeyPair();
         }
 
         @Override
         public GordianBaseFactory getFactory() {
             return (GordianBaseFactory) super.getFactory();
         }
 
         @Override
         public byte[] sign() throws GordianException {
             /* Check that we are in signing mode */
             checkMode(GordianSignatureMode.SIGN);
 
             /* Sign the message */
             try {
                 return theSigner.generateSignature();
             } catch (CryptoException e) {
                 throw new GordianCryptoException("Failed to sign message", e);
             }
         }
 
         @Override
         public boolean verify(final byte[] pSignature) throws GordianException {
             /* Check that we are in verify mode */
             checkMode(GordianSignatureMode.VERIFY);
 
             /* Verify the message */
             return theSigner.verifySignature(pSignature);
         }
     }
 }