/*
 * Copyright (c) The mldsa-native project authors
 * SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
 */

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* Import public mldsa-native API
 *
 * This requires specifying the parameter set and namespace prefix
 * used for the build.
 */
#include <mldsa_native.h>
#include "expected_signatures.h"

#define CHECK(x)                                              \
  do                                                          \
  {                                                           \
    int rc;                                                   \
    rc = (x);                                                 \
    if (!rc)                                                  \
    {                                                         \
      fprintf(stderr, "ERROR (%s,%d)\n", __FILE__, __LINE__); \
      return 1;                                               \
    }                                                         \
  } while (0)

#define TEST_MSG \
  "This is a test message for ML-DSA digital signature algorithm!"
#define TEST_MSG_LEN (sizeof(TEST_MSG) - 1)

#define TEST_CTX "test_context_123"
#define TEST_CTX_LEN (sizeof(TEST_CTX) - 1)

int main(void)
{
  const char test_msg[] = TEST_MSG;
  const char test_ctx[] = TEST_CTX;

  uint8_t pk[CRYPTO_PUBLICKEYBYTES];
  uint8_t sk[CRYPTO_SECRETKEYBYTES];
  uint8_t sig[CRYPTO_BYTES];
  /* For testing, we use the same randomness that would be generated by our
   * deterministic test rng such that we can produce the same keys and
   * signatures as in the basic example.
   */
  /* # !!! WARNING !!!
   * This is for TESTING ONLY.
   * For key generation, you MUST USE cryptographic randomness outside of
   * testing.
   * For signing, you SHOULD provide cryptographic randomness. If you require
   * deterministic signing, set sign_rnd to all zeros.
   */

  uint8_t keypair_rnd[MLDSA_SEEDBYTES] = {
      0x93, 0x4d, 0x60, 0xb3, 0x56, 0x24, 0xd7, 0x40, 0xb3, 0x0a, 0x7f,
      0x22, 0x7a, 0xf2, 0xae, 0x7c, 0x67, 0x8e, 0x4e, 0x04, 0xe1, 0x3c,
      0x5f, 0x50, 0x9e, 0xad, 0xe2, 0xb7, 0x9a, 0xea, 0x77, 0xe2};

  uint8_t sign_rnd[MLDSA_RNDBYTES] = {
      0x3e, 0x2a, 0x2e, 0xa6, 0xc9, 0xc4, 0x76, 0xfc, 0x49, 0x37, 0xb0,
      0x13, 0xc9, 0x93, 0xa7, 0x93, 0xd6, 0xc0, 0xab, 0x99, 0x60, 0x69,
      0x5b, 0xa8, 0x38, 0xf6, 0x49, 0xda, 0x53, 0x9c, 0xa3, 0xd0};
  uint8_t
      pre[TEST_CTX_LEN + 2]; /* prefix string (0, ctxlen, ctx) for signing */
  size_t siglen;

  printf("ML-DSA-%d Deterministic Example\n", MLD_CONFIG_PARAMETER_SET);
  printf("======================\n\n");

  printf("Message: %s\n", test_msg);
  printf("Context: %s\n\n", test_ctx);

  printf("Generating keypair ... ");

  /* Alice generates a public/private key pair */
  CHECK(mldsa_keypair_internal(pk, sk, keypair_rnd) == 0);

  printf("DONE\n");
  printf("Signing message... ");

  /* Prepare pre = (0, ctxlen, ctx) */
  pre[0] = 0;
  pre[1] = TEST_CTX_LEN;
  memcpy(pre + 2, test_ctx, TEST_CTX_LEN);


  /* Alice signs the message */
  CHECK(mldsa_signature_internal(sig, &siglen, (const uint8_t *)test_msg,
                                 TEST_MSG_LEN, pre, sizeof(pre), sign_rnd, sk,
                                 0) == 0);

  printf("DONE\n");
  printf("Verifying signature... ");

  /* Bob verifies Alice's signature */
  CHECK(mldsa_verify(sig, siglen, (const uint8_t *)test_msg, TEST_MSG_LEN,
                     (const uint8_t *)test_ctx, TEST_CTX_LEN, pk) == 0);

  printf("DONE\n");

  printf("Results:\n");
  printf("--------\n");
  printf("Public key size:  %d bytes\n", CRYPTO_PUBLICKEYBYTES);
  printf("Secret key size:  %d bytes\n", CRYPTO_SECRETKEYBYTES);
  printf("Signature size:   %d bytes\n", CRYPTO_BYTES);
  printf("Message length:   %lu bytes\n", (unsigned long)TEST_MSG_LEN);
  printf("Signature length: %lu bytes\n", (unsigned long)siglen);

  /* Check against expected signature to make sure that
   * we integrated the library correctly */
  printf("Checking deterministic signature... ");
  {
    /* Compare the generated signature directly against the expected signature
     */
    CHECK(siglen == sizeof(expected_signature));
    CHECK(memcmp(sig, expected_signature, siglen) == 0);
  }
  printf("DONE\n");

  printf("Signature verification completed successfully!\n");

  printf("\nAll tests passed! ML-DSA signature verification successful.\n");
  return 0;
}
