#!/usr/bin/env python3

sizes = '348864','460896','6688128','6960119','8192128'

import os
import shutil

def doit(source,target,sort=None,xof=None,kem=None):
  file = {}
  for fn in sorted(os.listdir(source)):
    with open(f'{source}/{fn}') as f:
      file[fn] = f.read()

  linkerdefine = {}
  linkeruse = {}
  wheredefined = {}
  for fn in file:
    linkerdefine[fn] = set()
    linkeruse[fn] = set()
    for line in file[fn].splitlines():
      if line.startswith('// linker define '):
        for x in line.split()[3:]:
          linkerdefine[fn].add(x)
          assert x not in wheredefined
          wheredefined[x] = fn
      if line.startswith('// linker use '):
        for x in line.split()[3:]:
          linkeruse[fn].add(x)
       
  keepinclude = set()

  use = {}
  for fn in sorted(file):
    use[fn] = set()
    if fn.endswith('.h'):
      use[fn].add(fn[:-2]+'.c')
    for x in sorted(linkeruse[fn]):
      use[fn].add(wheredefined[x])
    for line in file[fn].splitlines():
      if line.startswith('#include'):
        x = line.split()[1]
        if x[0] != '"' or x[-1] != '"':
          if x == '<inttypes.h>': x = '<stdint.h>'
          keepinclude.add(x)
          continue
        x = x[1:-1]
        if x in file:
          use[fn].add(x)
        else:
          keepinclude.add(f'"{x}"')
      if line.startswith('#define'):
        x = line.split()[1]
        if x in wheredefined:
          use[fn].add(wheredefined[x])
    if fn.endswith('.c'):
      use[fn].discard(fn[:-2]+'.h')

  def walk(todo,walked,fn):
    if fn not in file: return
    if fn in walked: return
    walked.add(fn)
    for u in sorted(use[fn]):
      walk(todo,walked,u)
    todo += [fn]
  todo = []
  walked = set()
  for fn in reversed(sorted(file)):
    if fn.endswith('.c'):
      walk(todo,walked,fn)

  skipinclude = set()
  if kem:
    skipinclude.add('"crypto_kem.h"')
  if xof:
    skipinclude.add('"crypto_xof.h"')
  if sort:
    skipinclude.add('"crypto_sort.h"')

  def includesort(i):
    if i.startswith('<'):
      return 0,i
    return 1,i

  keepinclude = sorted(keepinclude,key=includesort)

  with open(target,'w') as m:
    m.write(f'''// ===== auto-generated from {source}

''')
    if sort:
      m.write(f'#include "crypto_sort_int{sort}.h"\n')
      m.write('\n')
    if xof:
      m.write(f'#include "crypto_xof_bitwrite16.h"\n')
      m.write('\n')
    if kem:
      m.write('#include <wolfssl/wolfcrypt/libwolfssl_sources.h>\n')
      m.write('#include <wolfssl/wolfcrypt/mceliece.h>\n')
      m.write('\n')
      m.write(f'#define crypto_kem_keypair wc_mceliece{kem}_keypair\n')
      m.write(f'#define crypto_kem_enc wc_mceliece{kem}_enc\n')
      m.write(f'#define crypto_kem_dec wc_mceliece{kem}_dec\n')
      m.write('\n')
      m.write(f'#define CRYPTO_NAMESPACE(x) wc_mceliece{kem}_internal_##x\n')
      m.write(f'#define CRYPTO_SHARED_NAMESPACE(x) wc_mceliece{kem}_internal_##x\n')
      m.write('\n')

    for x in keepinclude:
      if x in skipinclude: continue
      m.write(f'#include {x}\n')
      skipinclude.add(x)

    skipinclude.add('<inttypes.h>')

    blankline = True
    for fn in todo:
      if fn.endswith('.h'):
        m.write(f'\n// ===== integrating {fn}\n\n')
        for line in file[fn].splitlines():
          if line.startswith('#include'):
            if line.split()[1] in skipinclude:
              continue
          if line.split() == []:
            if not blankline: m.write('\n')
            blankline = True
          else:
            m.write(line+'\n')
            blankline = False
        skipinclude.add(f'"{fn}"')
        m.write(f'\n// ===== done integrating {fn}\n')
    for fn in todo:
      if fn.endswith('.c'):
        m.write(f'\n// ===== integrating {fn}\n\n')
        for line in file[fn].splitlines():
          if line.startswith('// linker define '): continue
          if line.startswith('// linker use '): continue
          if line.startswith('#include'):
            x = line.split()[1]
            if x in skipinclude:
              continue
            if x[0] == '"' and x[-1] == '"' and x[1:-1] in file:
              m.write(f'// ===== level 1: integrating {x[1:-1]}\n')
              m.write(file[x[1:-1]])
              m.write(f'// ===== level 1: done integrating {x[1:-1]}\n')
              continue
          if sort:
            if line.strip() == 'void crypto_sort(void *array,long long n)':
              m.write(f'void wc_mceliece_internal_sort_int{sort}(int{sort} *x,long long n)\n')
              continue
            if line.strip() == f'int{sort} *x = array;': continue
            if line.strip() == f'int{sort} *x = (int{sort} *) array;': continue
          if xof:
            if line.strip() == 'int crypto_xof(':
              m.write('void wc_mceliece_internal_bitwrite16(\n')
              continue
            if line.strip() == 'return 0;': continue
          if kem:
            line = line.replace('crypto_sort_int16(ind,','crypto_sort_int16((crypto_int16 *) ind,')
            line = line.replace('crypto_sort_int64(list,','crypto_sort_int64((crypto_int64 *) list,')
            if fn == 'fft_tr.c': line = line.replace('postprocess','fft_tr_postprocess')
            if fn == 'controlbits.c': line = line.replace('layer','controlbits_layer')
            if line.strip() == 'long long x = 2*j;':
              m.write('#define x xnew\n')
            if line.startswith('int crypto_kem_'):
              m.write('WOLFSSL_API\n')
          if line.split() == []:
            if not blankline: m.write('\n')
            blankline = True
          else:
            m.write(line+'\n')
            blankline = False
        m.write(f'\n// ===== done integrating {fn}\n')

  return file['api.h']

def top():
  source='supercop-20250415'
  target_c='wolfssl/wolfcrypt/src'
  target_h='wolfssl/wolfssl/wolfcrypt'
  target_test='wolfssl/wolfcrypt/test'
  target_bench='wolfssl/wolfcrypt/benchmark'

  os.makedirs(f'{target_c}/mceliece',exist_ok=True)

  for b in 'int8','uint8','int16','uint16','int32','uint32','int64','uint64':
    with open(f'{source}/cryptoint/crypto_{b}.h') as f:
      with open(f'{target_c}/mceliece/crypto_{b}.h','w') as g:
        for line in f:
          g.write(line.replace('supercop_','wc_mceliece_internal_'))
    shutil.copy(f'{source}/cryptoint/{b}_optblocker.c',f'{target_c}/mceliece')

  with open(f'{target_c}/mceliece/crypto_declassify.h','w') as f:
    f.write('#define crypto_declassify(x,xlen) do {} while(0)')

  with open(f'{target_c}/mceliece/randombytes.h','w') as f:
    f.write('''#ifndef randombytes_h
#define randombytes_h

#define randombytes wc_mceliece_internal_randombytes

extern void randombytes(unsigned char *,unsigned long long);

#endif /* randombytes_h */
''')

  with open(f'{target_c}/mceliece/randombytes.c','w') as f:
    f.write('''#include <wolfssl/wolfcrypt/random.h>
#include "randombytes.h"

void randombytes(unsigned char *x,unsigned long long xlen)
{
    WC_RNG rng;
    while (wc_InitRng(&rng) != 0)
        ; // XXX
    // XXX: assuming xlen is small enough for GenerateBlock
    while (wc_RNG_GenerateBlock(&rng,x,xlen) != 0)
        ; // XXX
    wc_FreeRng(&rng);
}
''')

  doit(f'{source}/crypto_xof/bitwrite16/ref',f'{target_c}/mceliece/wc_bitwrite16.c',xof=True)
  with open(f'{target_c}/mceliece/crypto_xof_bitwrite16.h','w') as f:
    f.write('''#ifndef crypto_xof_bitwrite16_h
#define crypto_xof_bitwrite16_h

#define crypto_xof_bitwrite16 wc_mceliece_internal_bitwrite16

extern void crypto_xof_bitwrite16(
  unsigned char *,long long,
  const unsigned char *,long long
);

#endif /* crypto_xof_bitwrite16_h */
''')

  with open(f'{target_c}/mceliece/crypto_xof_shake256.h','w') as f:
    f.write(r'''#ifndef crypto_xof_shake256_h
#define crypto_xof_shake256_h

#include <wolfssl/wolfcrypt/hash.h>

#define crypto_xof_shake256(out,outlen,in,inlen) \
  wc_Shake256Hash(in,inlen,out,outlen)

#endif /* crypto_xof_shake256_h */
''')

  for b in '16','32','64':
    doit(f'{source}/crypto_sort/int{b}/portable4',f'{target_c}/mceliece/wc_sort_int{b}.c',sort=b)
    with open(f'{target_c}/mceliece/crypto_sort_int{b}.h','w') as f:
      f.write(f'''#ifndef crypto_sort_int{b}_h
#define crypto_sort_int{b}_h

#define crypto_sort_int{b} wc_mceliece_internal_sort_int{b}

#include "crypto_int{b}.h"

extern void crypto_sort_int{b}(crypto_int{b} *,long long);

#endif /* crypto_sort_int{b}_h */
''')

  api = {}

  for size in sizes:
    for kem in size,size+'pc':
      api[kem] = doit(f'{source}/crypto_kem/mceliece{kem}f/vec',f'{target_c}/mceliece/wc_mceliece{kem}.c',kem=kem)

  with open(f'{target_h}/mceliece.h','w') as f:
    f.write('''#ifndef WOLF_CRYPT_MCELIECE_H
#define WOLF_CRYPT_MCELIECE_H

#include <wolfssl/wolfcrypt/types.h>

#ifdef __cplusplus
    extern "C" {
#endif
''')

    for size in sizes:
      for kem in size,size+'pc':
        KEM = kem.upper()
        pkbytes = None
        skbytes = None
        ctbytes = None
        ssbytes = None
        for line in api[kem].splitlines():
          line = line.split()
          if line[:1] != ['#define']: continue
          if line[1] == 'CRYPTO_PUBLICKEYBYTES': pkbytes = line[2]
          if line[1] == 'CRYPTO_SECRETKEYBYTES': skbytes = line[2]
          if line[1] == 'CRYPTO_CIPHERTEXTBYTES': ctbytes = line[2]
          if line[1] == 'CRYPTO_BYTES': ssbytes = line[2]

        f.write(f'''
#ifdef WOLFSSL_HAVE_MCELIECE{KEM}

#define WC_MCELIECE{KEM}_PUBLIC_KEY_SIZE {pkbytes}
#define WC_MCELIECE{KEM}_PRIVATE_KEY_SIZE {skbytes}
#define WC_MCELIECE{KEM}_CIPHER_TEXT_SIZE {ctbytes}
#define WC_MCELIECE{KEM}_SS_SIZE {ssbytes}

WOLFSSL_API int wc_mceliece{kem}_keypair(unsigned char* pk,
    unsigned char *sk);
WOLFSSL_API int wc_mceliece{kem}_enc(unsigned char* c,
    unsigned char* ss, const unsigned char* pk);
WOLFSSL_API int wc_mceliece{kem}_dec(unsigned char* ss,
    const unsigned char* c, const unsigned char* sk);

#endif /* WOLFSSL_HAVE_MCELIECE{KEM} */
''')

    f.write('''

#ifdef __cplusplus
    } /* extern "C" */
#endif

#endif /* WOLF_CRYPT_MCELIECE_H */
''')

  with open(f'{target_test}/mceliece.i','w') as f:
    for size in sizes:
      for kem in size,size+'pc':
        f.write(f'#include "mceliece/{kem}.i"\n')

  os.makedirs(f'{target_test}/mceliece',exist_ok=True)

  for size in sizes:
    for kem in size,size+'pc':
      KEM = kem.upper()
      with open(f'{target_test}/mceliece/{kem}.i','w') as f:
        f.write(f'''#ifdef WOLFSSL_HAVE_MCELIECE{KEM}
WOLFSSL_TEST_SUBROUTINE wc_test_ret_t mceliece{kem}_test(void)
{{
    wc_test_ret_t ret;
    int i;
    unsigned char *pk = 0;
    unsigned char *sk = 0;
    unsigned char *ct = 0;
    unsigned char *ss = 0;
    unsigned char *pk2 = 0;
    unsigned char *sk2 = 0;
    unsigned char *ct2 = 0;
    unsigned char *ss2 = 0;
    WOLFSSL_ENTER("mceliece{kem}_test");

    pk = (unsigned char *)XMALLOC(WC_MCELIECE{KEM}_PUBLIC_KEY_SIZE, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
    if (!pk) ERROR_OUT(WC_TEST_RET_ENC_ERRNO, out);
    sk = (unsigned char *)XMALLOC(WC_MCELIECE{KEM}_PRIVATE_KEY_SIZE, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
    if (!sk) ERROR_OUT(WC_TEST_RET_ENC_ERRNO, out);
    ct = (unsigned char *)XMALLOC(WC_MCELIECE{KEM}_CIPHER_TEXT_SIZE, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
    if (!ct) ERROR_OUT(WC_TEST_RET_ENC_ERRNO, out);
    ss = (unsigned char *)XMALLOC(WC_MCELIECE{KEM}_SS_SIZE, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
    if (!ss) ERROR_OUT(WC_TEST_RET_ENC_ERRNO, out);
    pk2 = (unsigned char *)XMALLOC(WC_MCELIECE{KEM}_PUBLIC_KEY_SIZE, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
    if (!pk2) ERROR_OUT(WC_TEST_RET_ENC_ERRNO, out);
    sk2 = (unsigned char *)XMALLOC(WC_MCELIECE{KEM}_PRIVATE_KEY_SIZE, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
    if (!sk2) ERROR_OUT(WC_TEST_RET_ENC_ERRNO, out);
    ct2 = (unsigned char *)XMALLOC(WC_MCELIECE{KEM}_CIPHER_TEXT_SIZE, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
    if (!ct2) ERROR_OUT(WC_TEST_RET_ENC_ERRNO, out);
    ss2 = (unsigned char *)XMALLOC(WC_MCELIECE{KEM}_SS_SIZE, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
    if (!ss2) ERROR_OUT(WC_TEST_RET_ENC_ERRNO, out);

    for (i = 0;i < 10;++i) {{
        // test: keypair succeeds
        ret = wc_mceliece{kem}_keypair(pk, sk);
        if (ret != 0)
            ERROR_OUT(WC_TEST_RET_ENC_I(i), out);
        ret = wc_mceliece{kem}_keypair(pk2, sk2);
        if (ret != 0)
            ERROR_OUT(WC_TEST_RET_ENC_I(i), out);

        // test: multiple keypair calls produce different results
        if (XMEMCMP(pk, pk2, WC_MCELIECE{KEM}_PUBLIC_KEY_SIZE) == 0)
            ERROR_OUT(WC_TEST_RET_ENC_NC, out);
        if (XMEMCMP(sk, sk2, WC_MCELIECE{KEM}_PRIVATE_KEY_SIZE) == 0)
            ERROR_OUT(WC_TEST_RET_ENC_NC, out);

        // test: enc succeeds
        ret = wc_mceliece{kem}_enc(ct, ss, pk);
        if (ret != 0)
            ERROR_OUT(WC_TEST_RET_ENC_I(i), out);
        ret = wc_mceliece{kem}_enc(ct2, ss2, pk);
        if (ret != 0)
            ERROR_OUT(WC_TEST_RET_ENC_I(i), out);

        // test: multiple enc calls with same pk produce different results
        if (XMEMCMP(ct, ct2, WC_MCELIECE{KEM}_CIPHER_TEXT_SIZE) == 0)
            ERROR_OUT(WC_TEST_RET_ENC_NC, out);
        if (XMEMCMP(ss, ss2, WC_MCELIECE{KEM}_SS_SIZE) == 0)
            ERROR_OUT(WC_TEST_RET_ENC_NC, out);

        // test: dec succeeds
        ret = wc_mceliece{kem}_dec(ss2, ct, sk);
        if (ret != 0)
            ERROR_OUT(WC_TEST_RET_ENC_I(i), out);

        // test: dec gives same result as enc
        if (XMEMCMP(ss, ss2, WC_MCELIECE{KEM}_SS_SIZE) != 0)
            ERROR_OUT(WC_TEST_RET_ENC_NC, out);

        // test: dec succeeds if ct is tweaked
        ct[0] ^= 1;
        ret = wc_mceliece{kem}_dec(ss2, ct, sk);
        if (ret != 0)
            ERROR_OUT(WC_TEST_RET_ENC_I(i), out);

        // test: dec gives different ss if ct is tweaked
        if (XMEMCMP(ss, ss2, WC_MCELIECE{KEM}_SS_SIZE) == 0)
            ERROR_OUT(WC_TEST_RET_ENC_NC, out);
    }}

out:
    if (ss2) {{ XFREE(ss2, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); ss2 = 0; }}
    if (ct2) {{ XFREE(ct2, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); ct2 = 0; }}
    if (sk2) {{ XFREE(sk2, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); sk2 = 0; }}
    if (pk2) {{ XFREE(pk2, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); pk2 = 0; }}
    if (ss) {{ XFREE(ss, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); ss = 0; }}
    if (ct) {{ XFREE(ct, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); ct = 0; }}
    if (sk) {{ XFREE(sk, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); sk = 0; }}
    if (pk) {{ XFREE(pk, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); pk = 0; }}

    return ret;
}}
#endif /* WOLFSSL_HAVE_MCELIECE{KEM} */
''')

  with open(f'{target_bench}/mceliece.i','w') as f:
    for size in sizes:
      for kem in size,size+'pc':
        f.write(f'#include "mceliece/{kem}.i"\n')

  os.makedirs(f'{target_bench}/mceliece',exist_ok=True)

  for size in sizes:
    for kem in size,size+'pc':
      KEM = kem.upper()
      with open(f'{target_bench}/mceliece/{kem}.i','w') as f:
        f.write(fr'''#ifdef WOLFSSL_HAVE_MCELIECE{KEM}
void bench_mceliece{kem}(void)
{{
    const char* name = "mceliece";
    int keySize = {size};
    int ret = 0, times, count, pending = 0;
    double start;
    WC_RNG rng;
    const char**desc = bench_desc_words[lng_index];
    unsigned char *pk = 0;
    unsigned char *sk = 0;
    unsigned char *ct = 0;
    unsigned char *ss = 0;
    DECLARE_MULTI_VALUE_STATS_VARS()

    ret = wc_InitRng(&rng);
    if (ret < 0) {{
        printf("InitRng failed %d\n", ret);
        return;
    }}

    pk = (unsigned char *)XMALLOC(WC_MCELIECE{KEM}_PUBLIC_KEY_SIZE, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
    if (!pk) goto cleanup;
    sk = (unsigned char *)XMALLOC(WC_MCELIECE{KEM}_PRIVATE_KEY_SIZE, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
    if (!sk) goto cleanup;
    ct = (unsigned char *)XMALLOC(WC_MCELIECE{KEM}_CIPHER_TEXT_SIZE, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
    if (!ct) goto cleanup;
    ss = (unsigned char *)XMALLOC(WC_MCELIECE{KEM}_SS_SIZE, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
    if (!ss) goto cleanup;

    bench_stats_start(&count, &start);
    do {{
        for (times = 0; times < agreeTimes || pending > 0; times++) {{
            wc_mceliece{kem}_keypair(pk, sk);
            RECORD_MULTI_VALUE_STATS();
        }} /* for times */
        count += times;
    }} while (bench_stats_check(start)
#ifdef MULTI_VALUE_STATISTICS
       || runs < minimum_runs
#endif
       );
    bench_stats_asym_finish(name, keySize, desc[2], 0, count, start, ret);
#ifdef MULTI_VALUE_STATISTICS
    bench_multi_value_stats(max, min, sum, squareSum, runs);
#endif

    RESET_MULTI_VALUE_STATS_VARS();
    bench_stats_start(&count, &start);
    do {{
        for (times = 0; times < agreeTimes || pending > 0; times++) {{
            wc_mceliece{kem}_enc(ct, ss, pk);
            RECORD_MULTI_VALUE_STATS();
        }} /* for times */
        count += times;
    }} while (bench_stats_check(start)
#ifdef MULTI_VALUE_STATISTICS
       || runs < minimum_runs
#endif
       );
    bench_stats_asym_finish(name, keySize, desc[9], 0, count, start, ret);
#ifdef MULTI_VALUE_STATISTICS
    bench_multi_value_stats(max, min, sum, squareSum, runs);
#endif

    RESET_MULTI_VALUE_STATS_VARS();
    bench_stats_start(&count, &start);
    do {{
        for (times = 0; times < agreeTimes || pending > 0; times++) {{
            if (wc_mceliece{kem}_dec(ss, ct, sk) != 0) goto dec_failed;
            RECORD_MULTI_VALUE_STATS();
        }} /* for times */
        count += times;
    }} while (bench_stats_check(start)
#ifdef MULTI_VALUE_STATISTICS
       || runs < minimum_runs
#endif
       );
dec_failed:
    bench_stats_asym_finish(name, keySize, desc[13], 0, count, start, ret);
#ifdef MULTI_VALUE_STATISTICS
    bench_multi_value_stats(max, min, sum, squareSum, runs);
#endif

cleanup:
    if (ss) {{ XFREE(ss, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); ss = 0; }}
    if (ct) {{ XFREE(ct, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); ct = 0; }}
    if (sk) {{ XFREE(sk, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); sk = 0; }}
    if (pk) {{ XFREE(pk, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); pk = 0; }}
}}
#endif /* WOLFSSL_HAVE_MCELIECE{KEM} */
''')

top()