#!/usr/bin/env python3
"""Converts the ChillDKG test vectors of the bip-frost-dkg reference
repository into a C header used by the tests in src/modules/chilldkg/tests_impl.h.

Usage: tools/test_vectors_chilldkg_generate.py <dir>

<dir> must contain the vector files of the bip-frost-dkg repository
(https://github.com/BlockstreamResearch/bip-frost-dkg, directory vectors/)
pinned at commit a91896883f85b159415ecf298d5e844879af112d.
"""

import sys
import json
import textwrap

if len(sys.argv) < 2:
    print(__doc__)
    sys.exit(1)

VEC_DIR = sys.argv[1]

FILES = {
    "hostpubkey_gen": "hostpubkey_gen_vectors.json",
    "params_hash": "params_hash_vectors.json",
    "participant_step1": "participant_step1_vectors.json",
    "participant_step2": "participant_step2_vectors.json",
    "participant_finalize": "participant_finalize_vectors.json",
    "participant_investigate": "participant_investigate_vectors.json",
    "coordinator_step1": "coordinator_step1_vectors.json",
    "coordinator_finalize": "coordinator_finalize_vectors.json",
    "coordinator_investigate": "coordinator_investigate_vectors.json",
    "recover": "recover_vectors.json",
}

data = {k: json.load(open(VEC_DIR + "/" + f)) for k, f in FILES.items()}

skipped_cases = []


def hexstr_to_intarray(s):
    # Always emitted inside brace initializers. C90 forbids empty initializer
    # braces, so an empty byte string becomes "0" (the associated length field
    # is 0, so the padding byte is never read).
    return ", ".join([f"0x{b:02X}" for b in bytes.fromhex(s)]) or "0"


def byte_array(hex_str):
    return "{ %s }" % hexstr_to_intarray(hex_str)


def indent(s, level=1):
    return textwrap.indent(s, 4 * level * " ")


def map_error(case, step_name):
    """Maps the reference expectedError of an error case to a
    chilldkg_vec_error code and fault index. Returns None if the case is not
    expressible in the C API (wrong-length or wrong-count inputs; the C API
    uses fixed-size buffers whose lengths derive from the session
    parameters)."""
    err = case["expectedError"]
    etype = err["type"]
    if etype == "ValueError":
        skipped_cases.append(
            "%s tcId %d (%s): not expressible in the C API (wrong input length or count)"
            % (step_name, case["tcId"], case.get("comment", "").strip())
        )
        return None
    fault_index = "UINT32_MAX"
    if "participantId" in err:
        assert etype in ("FaultyParticipantError", "FaultyParticipantOrCoordinatorError", "InvalidHostPubkeyError", "DuplicateHostPubkeyError")
        fault_index = str(err["participantId"])
    if etype in ("HostSeckeyError", "ThresholdOrCountError", "InvalidHostPubkeyError",
                 "DuplicateHostPubkeyError", "RandomnessError", "RecoveryDataError"):
        return "CHILLDKG_VEC_INVALID_INPUT", "UINT32_MAX"
    if etype == "FaultyCoordinatorError":
        return "CHILLDKG_VEC_FAULTY_COORDINATOR", fault_index
    if etype == "FaultyParticipantError":
        return "CHILLDKG_VEC_FAULTY_PARTICIPANT", fault_index
    if etype == "FaultyParticipantOrCoordinatorError":
        return "CHILLDKG_VEC_FAULTY_PARTICIPANT_OR_COORDINATOR", fault_index
    if etype == "UnknownFaultyParticipantOrCoordinatorError":
        return "CHILLDKG_VEC_UNKNOWN_FAULTY_PARTICIPANT_OR_COORDINATOR", fault_index
    sys.exit("Unknown error type %s in %s tcId %d" % (etype, step_name, case["tcId"]))


def groups(d):
    return d["testGroups"] if "testGroups" in d else [d]


def all_cases(d):
    for g in groups(d):
        for c in g.get("validTestCases", []) + g.get("errorTestCases", []):
            yield g, c


def get_params(case, group):
    return case.get("params", group.get("params"))


# Compute the shared maximum sizes.
max_n = 0
max_t = 0
max_pmsg1 = 0
max_cmsg1 = 0
max_cmsg2 = 0
max_recovery = 0
max_cinv = 0
max_pmsg1_pool = 0
max_pmsg2_pool = 0
max_cmsg1_pool = 0
max_indices = 0
for name, d in data.items():
    for g, c in all_cases(d):
        p = get_params(c, g)
        if p is not None:
            max_n = max(max_n, len(p["hostpubkeys"]))
            max_t = max(max_t, p["t"])
        for key, cur in (("expectedPmsg1", "pmsg1"), ("cmsg1", "cmsg1"), ("cmsg2", "cmsg2"),
                         ("recoveryData", "recovery"), ("cinvMsg", "cinv")):
            v = c.get(key)
            if isinstance(v, str):
                if cur == "pmsg1":
                    max_pmsg1 = max(max_pmsg1, len(v) // 2)
                elif cur == "cmsg1":
                    max_cmsg1 = max(max_cmsg1, len(v) // 2)
                elif cur == "cmsg2":
                    max_cmsg2 = max(max_cmsg2, len(v) // 2)
                elif cur == "recovery":
                    max_recovery = max(max_recovery, len(v) // 2)
                elif cur == "cinv":
                    max_cinv = max(max_cinv, len(v) // 2)
        if "pmsg1Indices" in c:
            max_indices = max(max_indices, len(c["pmsg1Indices"]))
        if "pmsg2Indices" in c:
            max_indices = max(max_indices, len(c["pmsg2Indices"]))
    for g in groups(d):
        if "pmsg1" in g and isinstance(g["pmsg1"], str):
            max_pmsg1 = max(max_pmsg1, len(g["pmsg1"]) // 2)
        if "cmsg1" in g and isinstance(g["cmsg1"], str):
            max_cmsg1 = max(max_cmsg1, len(g["cmsg1"]) // 2)
        if "pmsgs1" in g:
            max_n = max(max_n, len(g["pmsgs1"]))
            for m in g["pmsgs1"]:
                max_pmsg1 = max(max_pmsg1, len(m) // 2)
        if "pmsg1Pool" in g:
            max_pmsg1_pool = max(max_pmsg1_pool, len(g["pmsg1Pool"]))
            for m in g["pmsg1Pool"]:
                max_pmsg1 = max(max_pmsg1, len(m) // 2)
        if "pmsg2Pool" in g:
            max_pmsg2_pool = max(max_pmsg2_pool, len(g["pmsg2Pool"]))
        if "cmsg1Pool" in g:
            max_cmsg1_pool = max(max_cmsg1_pool, len(g["cmsg1Pool"]))
            for m in g["cmsg1Pool"]:
                max_cmsg1 = max(max_cmsg1, len(m) // 2)
        for c in g.get("validTestCases", []):
            if "expectedCinvMsgs" in c:
                for m in c["expectedCinvMsgs"]:
                    max_cinv = max(max_cinv, len(m) // 2)
            eo = c.get("expectedOutput")
            if eo is not None:
                if "recoveryData" in eo:
                    max_recovery = max(max_recovery, len(eo["recoveryData"]) // 2)
                if "cmsg2" in eo:
                    max_cmsg2 = max(max_cmsg2, len(eo["cmsg2"]) // 2)
        for c in g.get("validTestCases", []) + g.get("errorTestCases", []):
            if isinstance(c.get("recoveryData"), str):
                max_recovery = max(max_recovery, len(c["recoveryData"]) // 2)

s = """/**
 * Automatically generated by %s.
 *
 * The test vectors are from the bip-frost-dkg reference repository
 * https://github.com/BlockstreamResearch/bip-frost-dkg, pinned at commit
 * a91896883f85b159415ecf298d5e844879af112d (BIP version 0.3.0-dev). They are
 * used by the tests in src/modules/chilldkg/tests_impl.h.
 *
 * Regenerate with:
 *   tools/test_vectors_chilldkg_generate.py <dir> > src/modules/chilldkg/vectors.h
 * where <dir> is the vectors/ directory of the reference repository at the
 * commit above. */
#ifndef SECP256K1_MODULE_CHILLDKG_VECTORS_H
#define SECP256K1_MODULE_CHILLDKG_VECTORS_H

#include <stddef.h>
#include <stdint.h>

""" % sys.argv[0]

s += """/* The expected result of an error test case, mapped from the reference
 * exception taxonomy; the values coincide with secp256k1_chilldkg_fault.
 * CHILLDKG_VEC_INVALID_INPUT covers the local input errors (ValueError
 * subclasses that are not protocol faults: HostSeckeyError,
 * ThresholdOrCountError, InvalidHostPubkeyError, DuplicateHostPubkeyError,
 * RandomnessError, RecoveryDataError). */
enum chilldkg_vec_error {
    CHILLDKG_VEC_OK = 0,
    CHILLDKG_VEC_FAULTY_COORDINATOR = 1,
    CHILLDKG_VEC_FAULTY_PARTICIPANT = 2,
    CHILLDKG_VEC_FAULTY_PARTICIPANT_OR_COORDINATOR = 3,
    CHILLDKG_VEC_UNKNOWN_FAULTY_PARTICIPANT_OR_COORDINATOR = 4,
    CHILLDKG_VEC_INVALID_INPUT = 5
};

enum {
    CHILLDKG_VEC_MAX_PARTICIPANTS = %d,
    CHILLDKG_VEC_MAX_PMSG1_LEN = %d,
    CHILLDKG_VEC_MAX_CMSG1_LEN = %d,
    CHILLDKG_VEC_MAX_CMSG2_LEN = %d,
    CHILLDKG_VEC_MAX_RECOVERY_LEN = %d,
    CHILLDKG_VEC_MAX_CINV_LEN = %d,
    CHILLDKG_VEC_MAX_PMSG1_POOL = %d,
    CHILLDKG_VEC_MAX_PMSG2_POOL = %d,
    CHILLDKG_VEC_MAX_CMSG1_POOL = %d,
    CHILLDKG_VEC_MAX_INDICES = %d
};

/* Session parameters (hostpubkeys || t). */
struct chilldkg_vec_params {
    uint32_t t;
    size_t n;
    unsigned char hostpubkeys[CHILLDKG_VEC_MAX_PARTICIPANTS][33];
};

/* The DKG output of a participant or the coordinator (has_secshare == 0). */
struct chilldkg_vec_dkg_output {
    int has_secshare;
    unsigned char secshare[32];
    unsigned char thresh_pk[33];
    unsigned char pubshares[CHILLDKG_VEC_MAX_PARTICIPANTS][33];
};
""" % (max_n, max_pmsg1, max_cmsg1, max_cmsg2, max_recovery, max_cinv,
       max_pmsg1_pool, max_pmsg2_pool, max_cmsg1_pool, max_indices)


def init_params(p):
    inner = "\n" + indent(",\n".join(byte_array(h) for h in p["hostpubkeys"]), 2) + "\n"
    return "{ %d, %d, {%s} }" % (p["t"], len(p["hostpubkeys"]), inner)


def init_dkg_output(o):
    return "{ %d, %s, %s, {\n%s\n} }" % (
        0 if o["secshare"] is None else 1,
        byte_array(o["secshare"]) if o["secshare"] is not None else "{ 0 }",
        byte_array(o["threshPk"]),
        indent(",\n".join(byte_array(p) for p in o["pubshares"]), 2),
    )


def case_list(d):
    """Returns the (is_valid, case) list of a vector file (flat per-group
    files) with inexpressible error cases removed."""
    out = []
    for g in groups(d):
        for c in g.get("validTestCases", []):
            out.append((1, c))
        for c in g.get("errorTestCases", []):
            out.append((0, c))
    return out


# hostpubkey_gen vectors
cases = []
for is_valid, c in case_list(data["hostpubkey_gen"]):
    if not is_valid and map_error(c, "hostpubkey_gen") is None:
        continue
    cases.append((is_valid, c))
s += """
struct chilldkg_vec_hostpubkey_gen_case {
    int is_valid;
    unsigned char hostseckey[32];
    unsigned char expected_hostpubkey[33];
};

static const struct chilldkg_vec_hostpubkey_gen_case chilldkg_vec_hostpubkey_gen_cases[%d] = {
""" % len(cases)
for is_valid, c in cases:
    s += indent("{ %d, %s, %s },\n" % (
        is_valid,
        byte_array(c["hostseckey"]),
        byte_array(c.get("expectedHostpubkey") or ""),
    ))
s += "};\n"

# params_hash vectors
cases = []
for is_valid, c in case_list(data["params_hash"]):
    if not is_valid and map_error(c, "params_hash") is None:
        continue
    cases.append((is_valid, c))
s += """
struct chilldkg_vec_params_hash_case {
    int is_valid;
    struct chilldkg_vec_params params;
    unsigned char expected_hash[32];
};

static const struct chilldkg_vec_params_hash_case chilldkg_vec_params_hash_cases[%d] = {
""" % len(cases)
for is_valid, c in cases:
    s += indent("{ %d,\n%s,\n%s },\n" % (
        is_valid,
        indent(init_params(c["params"]) + ",", 1)[:-1],
        byte_array(c.get("expectedParamsHash") or ""),
    ))
s += "};\n"

# participant_step1 vectors
cases = []
for is_valid, c in case_list(data["participant_step1"]):
    if not is_valid and map_error(c, "participant_step1") is None:
        continue
    cases.append((is_valid, c))
s += """
struct chilldkg_vec_step1_case {
    int is_valid;
    unsigned char hostseckey[32];
    struct chilldkg_vec_params params;
    unsigned char random[32];
    unsigned char expected_pmsg1[CHILLDKG_VEC_MAX_PMSG1_LEN];
};

static const struct chilldkg_vec_step1_case chilldkg_vec_step1_cases[%d] = {
""" % len(cases)
for is_valid, c in cases:
    s += indent("{ %d, %s,\n%s,\n%s, %s },\n" % (
        is_valid,
        byte_array(c["hostseckey"]),
        indent(init_params(c["params"]) + ",", 1)[:-1],
        byte_array(c["random"]),
        byte_array(c.get("expectedPmsg1") or ""),
    ))
s += "};\n"

# participant_step2 vectors
s += """
struct chilldkg_vec_step2_group {
    struct chilldkg_vec_params params;
    unsigned char hostseckey[32];
    unsigned char random[32];
    unsigned char aux_rand[32];
    unsigned char pmsg1[CHILLDKG_VEC_MAX_PMSG1_LEN];
};

struct chilldkg_vec_step2_case {
    int is_valid;
    /* hostseckey override (error cases only); has_hostseckey == 0 means the
     * group's hostseckey is used. */
    int has_hostseckey;
    unsigned char hostseckey[32];
    unsigned char cmsg1[CHILLDKG_VEC_MAX_CMSG1_LEN];
    /* Expected result: CHILLDKG_VEC_OK or an enum chilldkg_vec_error value,
     * mapped to secp256k1_chilldkg_fault / return-0 as in the tests. */
    int expected_error;
    uint32_t expected_fault_index;
    unsigned char expected_pmsg2[64];
};

"""
s += "static const struct chilldkg_vec_step2_group chilldkg_vec_step2_groups[%d] = {\n" % len(groups(data["participant_step2"]))
for g in groups(data["participant_step2"]):
    s += indent("{\n%s,\n%s, %s, %s, %s },\n" % (
        indent(init_params(g["params"]) + ",", 1)[:-1],
        byte_array(g["hostseckey"]),
        byte_array(g["random"]),
        byte_array(g["auxRand"]),
        byte_array(g["pmsg1"]),
    ))
s += "};\n"
step2_cases = []
for gi, g in enumerate(groups(data["participant_step2"])):
    for c in g.get("validTestCases", []):
        step2_cases.append((gi, 1, c))
    for c in g.get("errorTestCases", []):
        if map_error(c, "participant_step2") is None:
            continue
        step2_cases.append((gi, 0, c))
s += "static const struct chilldkg_vec_step2_case chilldkg_vec_step2_cases[%d] = {\n" % len(step2_cases)
for gi, is_valid, c in step2_cases:
    if is_valid:
        error_code, fault_index = "CHILLDKG_VEC_OK", "UINT32_MAX"
        expected = byte_array(c["expectedPmsg2"])
    else:
        error_code, fault_index = map_error(c, "participant_step2")
        expected = "{ 0 }"
    has_hsk = 1 if "hostseckey" in c else 0
    s += indent("{ %d, %d, %s,\n%s,\n%s, %s, %s },\n" % (
        is_valid,
        has_hsk,
        byte_array(c.get("hostseckey") or "00" * 32),
        indent(byte_array(c["cmsg1"]) + ",", 1)[:-1],
        error_code, fault_index, expected,
    ))
# Note: the group index is implicit (cases are emitted group by group).
s += "};\n"
s += "static const size_t chilldkg_vec_step2_group_case_counts[%d] = {\n" % len(groups(data["participant_step2"]))
gi = 0
for g in groups(data["participant_step2"]):
    cnt = sum(1 for x in step2_cases if x[0] == gi)
    s += indent("%d,\n" % cnt)
    gi += 1
s += "};\n"

# participant_finalize vectors
s += """
struct chilldkg_vec_finalize_group {
    struct chilldkg_vec_params params;
    unsigned char hostseckey[32];
    unsigned char random[32];
    unsigned char aux_rand[32];
    unsigned char pmsg1[CHILLDKG_VEC_MAX_PMSG1_LEN];
    unsigned char cmsg1[CHILLDKG_VEC_MAX_CMSG1_LEN];
    unsigned char pmsg2[64];
};

struct chilldkg_vec_finalize_case {
    int is_valid;
    unsigned char cmsg2[CHILLDKG_VEC_MAX_CMSG2_LEN];
    /* Expected result: CHILLDKG_VEC_OK or CHILLDKG_VEC_FAULTY_COORDINATOR. */
    int expected_error;
    struct chilldkg_vec_dkg_output expected_output;
    unsigned char expected_recovery[CHILLDKG_VEC_MAX_RECOVERY_LEN];
};

"""
s += "static const struct chilldkg_vec_finalize_group chilldkg_vec_finalize_groups[%d] = {\n" % len(groups(data["participant_finalize"]))
for g in groups(data["participant_finalize"]):
    s += indent("{\n%s,\n%s, %s, %s, %s,\n%s,\n%s },\n" % (
        indent(init_params(g["params"]) + ",", 1)[:-1],
        byte_array(g["hostseckey"]),
        byte_array(g["random"]),
        byte_array(g["auxRand"]),
        byte_array(g["pmsg1"]),
        indent(byte_array(g["cmsg1"]) + ",", 1)[:-1],
        byte_array(g["pmsg2"]),
    ))
s += "};\n"
finalize_cases = []
for gi, g in enumerate(groups(data["participant_finalize"])):
    for c in g.get("validTestCases", []):
        finalize_cases.append((gi, 1, c))
    for c in g.get("errorTestCases", []):
        if map_error(c, "participant_finalize") is None:
            continue
        finalize_cases.append((gi, 0, c))
s += "static const struct chilldkg_vec_finalize_case chilldkg_vec_finalize_cases[%d] = {\n" % len(finalize_cases)
for gi, is_valid, c in finalize_cases:
    if is_valid:
        error_code = "CHILLDKG_VEC_OK"
        eo = init_dkg_output(c["expectedOutput"]["dkgOutput"])
        rec = byte_array(c["expectedOutput"]["recoveryData"])
    else:
        error_code, _ = map_error(c, "participant_finalize")
        eo = "{ 0, { 0 }, { 0 }, { { 0 } } }"
        rec = "{ 0 }"
    s += indent("{ %d,\n%s,\n%s,\n%s,\n%s },\n" % (
        is_valid,
        indent(byte_array(c["cmsg2"]) + ",", 1)[:-1],
        error_code,
        indent(eo + ",", 1)[:-1],
        indent(rec, 1),
    ))
s += "};\n"
s += "static const size_t chilldkg_vec_finalize_group_case_counts[%d] = {\n" % len(groups(data["participant_finalize"]))
for gi in range(len(groups(data["participant_finalize"]))):
    cnt = sum(1 for x in finalize_cases if x[0] == gi)
    s += indent("%d,\n" % cnt)
s += "};\n"

# participant_investigate vectors
s += """
struct chilldkg_vec_investigate_group {
    struct chilldkg_vec_params params;
    unsigned char hostseckey[32];
    unsigned char random[32];
    unsigned char aux_rand[32];
    unsigned char pmsg1[CHILLDKG_VEC_MAX_PMSG1_LEN];
    size_t n_cmsg1;
    unsigned char cmsg1_pool[CHILLDKG_VEC_MAX_CMSG1_POOL][CHILLDKG_VEC_MAX_CMSG1_LEN];
};

struct chilldkg_vec_investigate_case {
    size_t cmsg1_index;
    unsigned char cinv[CHILLDKG_VEC_MAX_CINV_LEN];
    /* Expected result of participant_investigate. */
    int expected_error;
    uint32_t expected_fault_index;
};

"""
s += "static const struct chilldkg_vec_investigate_group chilldkg_vec_investigate_groups[%d] = {\n" % len(groups(data["participant_investigate"]))
for g in groups(data["participant_investigate"]):
    pool = g["cmsg1Pool"]
    inner = "\n" + indent(",\n".join(byte_array(m) for m in pool), 2) + "\n"
    s += indent("{\n%s,\n%s, %s, %s, %s,\n%s, {%s} },\n" % (
        indent(init_params(g["params"]) + ",", 1)[:-1],
        byte_array(g["hostseckey"]),
        byte_array(g["random"]),
        byte_array(g["auxRand"]),
        byte_array(g["pmsg1"]),
        len(pool),
        inner,
    ))
s += "};\n"
investigate_cases = []
for gi, g in enumerate(groups(data["participant_investigate"])):
    for c in g.get("errorTestCases", []):
        investigate_cases.append((gi, c))
s += "static const struct chilldkg_vec_investigate_case chilldkg_vec_investigate_cases[%d] = {\n" % len(investigate_cases)
for gi, c in investigate_cases:
    error_code, fault_index = map_error(c, "participant_investigate")
    s += indent("{ %d,\n%s,\n%s, %s },\n" % (
        c["cmsg1Index"],
        indent(byte_array(c["cinvMsg"]) + ",", 1)[:-1],
        error_code, fault_index,
    ))
s += "};\n"
s += "static const size_t chilldkg_vec_investigate_group_case_counts[%d] = {\n" % len(groups(data["participant_investigate"]))
for gi in range(len(groups(data["participant_investigate"]))):
    cnt = sum(1 for x in investigate_cases if x[0] == gi)
    s += indent("%d,\n" % cnt)
s += "};\n"

# coordinator_step1 vectors
s += """
struct chilldkg_vec_coord_step1_group {
    size_t n_pmsg1_pool;
    unsigned char pmsg1_pool[CHILLDKG_VEC_MAX_PMSG1_POOL][CHILLDKG_VEC_MAX_PMSG1_LEN];
};

struct chilldkg_vec_coord_step1_case {
    int is_valid;
    size_t n_pmsgs1;
    size_t pmsg1_indices[CHILLDKG_VEC_MAX_INDICES];
    struct chilldkg_vec_params params;
    /* Expected result: CHILLDKG_VEC_OK or an enum chilldkg_vec_error value. */
    int expected_error;
    uint32_t expected_fault_index;
    unsigned char expected_cmsg1[CHILLDKG_VEC_MAX_CMSG1_LEN];
};

"""
s += "static const struct chilldkg_vec_coord_step1_group chilldkg_vec_coord_step1_groups[%d] = {\n" % len(groups(data["coordinator_step1"]))
for g in groups(data["coordinator_step1"]):
    pool = g["pmsg1Pool"]
    inner = "\n" + indent(",\n".join(byte_array(m) for m in pool), 2) + "\n"
    s += indent("{ %d, {%s} },\n" % (len(pool), inner))
s += "};\n"
coord_step1_cases = []
for gi, g in enumerate(groups(data["coordinator_step1"])):
    for c in g.get("validTestCases", []):
        coord_step1_cases.append((gi, 1, c))
    for c in g.get("errorTestCases", []):
        if map_error(c, "coordinator_step1") is None:
            continue
        coord_step1_cases.append((gi, 0, c))
s += "static const struct chilldkg_vec_coord_step1_case chilldkg_vec_coord_step1_cases[%d] = {\n" % len(coord_step1_cases)
for gi, is_valid, c in coord_step1_cases:
    if is_valid:
        error_code, fault_index = "CHILLDKG_VEC_OK", "UINT32_MAX"
        expected = byte_array(c["expectedCmsg1"])
    else:
        error_code, fault_index = map_error(c, "coordinator_step1")
        expected = "{ 0 }"
    s += indent("{ %d, %d, { %s },\n%s,\n%s, %s,\n%s },\n" % (
        is_valid,
        len(c["pmsg1Indices"]),
        ", ".join(map(str, c["pmsg1Indices"])),
        indent(init_params(c["params"]) + ",", 1)[:-1],
        error_code, fault_index,
        indent(expected, 1),
    ))
s += "};\n"
s += "static const size_t chilldkg_vec_coord_step1_group_case_counts[%d] = {\n" % len(groups(data["coordinator_step1"]))
for gi in range(len(groups(data["coordinator_step1"]))):
    cnt = sum(1 for x in coord_step1_cases if x[0] == gi)
    s += indent("%d,\n" % cnt)
s += "};\n"

# coordinator_finalize vectors
s += """
struct chilldkg_vec_coord_finalize_group {
    struct chilldkg_vec_params params;
    unsigned char pmsgs1[CHILLDKG_VEC_MAX_PARTICIPANTS][CHILLDKG_VEC_MAX_PMSG1_LEN];
    unsigned char cmsg1[CHILLDKG_VEC_MAX_CMSG1_LEN];
    size_t n_pmsg2_pool;
    unsigned char pmsg2_pool[CHILLDKG_VEC_MAX_PMSG2_POOL][64];
};

struct chilldkg_vec_coord_finalize_case {
    int is_valid;
    size_t n_pmsgs2;
    size_t pmsg2_indices[CHILLDKG_VEC_MAX_INDICES];
    /* Expected result: CHILLDKG_VEC_OK or CHILLDKG_VEC_FAULTY_PARTICIPANT. */
    int expected_error;
    uint32_t expected_fault_index;
    unsigned char expected_cmsg2[CHILLDKG_VEC_MAX_CMSG2_LEN];
    struct chilldkg_vec_dkg_output expected_output;
    unsigned char expected_recovery[CHILLDKG_VEC_MAX_RECOVERY_LEN];
};

"""
s += "static const struct chilldkg_vec_coord_finalize_group chilldkg_vec_coord_finalize_groups[%d] = {\n" % len(groups(data["coordinator_finalize"]))
for g in groups(data["coordinator_finalize"]):
    pool = g["pmsg2Pool"]
    p1_inner = "\n" + indent(",\n".join(byte_array(m) for m in g["pmsgs1"]), 2) + "\n"
    p2_inner = "\n" + indent(",\n".join(byte_array(m) for m in pool), 2) + "\n"
    s += indent("{\n%s,\n{%s},\n%s,\n%s, {%s} },\n" % (
        indent(init_params(g["params"]) + ",", 1)[:-1],
        p1_inner,
        indent(byte_array(g["cmsg1"]) + ",", 1)[:-1],
        len(pool),
        p2_inner,
    ))
s += "};\n"
coord_finalize_cases = []
for gi, g in enumerate(groups(data["coordinator_finalize"])):
    for c in g.get("validTestCases", []):
        coord_finalize_cases.append((gi, 1, c))
    for c in g.get("errorTestCases", []):
        if map_error(c, "coordinator_finalize") is None:
            continue
        coord_finalize_cases.append((gi, 0, c))
s += "static const struct chilldkg_vec_coord_finalize_case chilldkg_vec_coord_finalize_cases[%d] = {\n" % len(coord_finalize_cases)
for gi, is_valid, c in coord_finalize_cases:
    if is_valid:
        error_code, fault_index = "CHILLDKG_VEC_OK", "UINT32_MAX"
        expected_cmsg2 = byte_array(c["expectedOutput"]["cmsg2"])
        eo = init_dkg_output(c["expectedOutput"]["dkgOutput"])
        rec = byte_array(c["expectedOutput"]["recoveryData"])
    else:
        error_code, fault_index = map_error(c, "coordinator_finalize")
        expected_cmsg2 = "{ 0 }"
        eo = "{ 0, { 0 }, { 0 }, { { 0 } } }"
        rec = "{ 0 }"
    s += indent("{ %d, %d, { %s }, %s, %s,\n%s,\n%s,\n%s },\n" % (
        is_valid,
        len(c["pmsg2Indices"]),
        ", ".join(map(str, c["pmsg2Indices"])),
        error_code, fault_index,
        indent(expected_cmsg2 + ",", 1)[:-1],
        indent(eo + ",", 1)[:-1],
        indent(rec, 1),
    ))
s += "};\n"
s += "static const size_t chilldkg_vec_coord_finalize_group_case_counts[%d] = {\n" % len(groups(data["coordinator_finalize"]))
for gi in range(len(groups(data["coordinator_finalize"]))):
    cnt = sum(1 for x in coord_finalize_cases if x[0] == gi)
    s += indent("%d,\n" % cnt)
s += "};\n"

# coordinator_investigate vectors
s += """
struct chilldkg_vec_coord_investigate_group {
    struct chilldkg_vec_params params;
    unsigned char pmsgs1[CHILLDKG_VEC_MAX_PARTICIPANTS][CHILLDKG_VEC_MAX_PMSG1_LEN];
    unsigned char expected_cinv[CHILLDKG_VEC_MAX_PARTICIPANTS][CHILLDKG_VEC_MAX_CINV_LEN];
};

"""
s += "static const struct chilldkg_vec_coord_investigate_group chilldkg_vec_coord_investigate_groups[%d] = {\n" % len(groups(data["coordinator_investigate"]))
for g in groups(data["coordinator_investigate"]):
    p1_inner = "\n" + indent(",\n".join(byte_array(m) for m in g["pmsgs1"]), 2) + "\n"
    cinv_inner = "\n" + indent(",\n".join(byte_array(m) for m in g["validTestCases"][0]["expectedCinvMsgs"]), 2) + "\n"
    s += indent("{\n%s,\n{%s},\n{%s} },\n" % (
        indent(init_params(g["params"]) + ",", 1)[:-1],
        p1_inner,
        cinv_inner,
    ))
s += "};\n"

# recover vectors
s += """
struct chilldkg_vec_recover_case {
    int is_valid;
    /* Coordinator recovery if has_hostseckey == 0. */
    int has_hostseckey;
    unsigned char hostseckey[32];
    size_t recovery_len;
    unsigned char recovery[CHILLDKG_VEC_MAX_RECOVERY_LEN];
    /* Expected result: CHILLDKG_VEC_OK or CHILLDKG_VEC_INVALID_INPUT. */
    int expected_error;
    struct chilldkg_vec_dkg_output expected_output;
    struct chilldkg_vec_params expected_params;
};

"""
recover_cases = []
for c in data["recover"]["validTestCases"]:
    recover_cases.append((1, c))
for c in data["recover"]["errorTestCases"]:
    if map_error(c, "recover") is None:
        continue
    recover_cases.append((0, c))
s += "static const struct chilldkg_vec_recover_case chilldkg_vec_recover_cases[%d] = {\n" % len(recover_cases)
for is_valid, c in recover_cases:
    if is_valid:
        error_code = "CHILLDKG_VEC_OK"
        eo = init_dkg_output(c["expectedOutput"]["dkgOutput"])
        ep = init_params(c["expectedOutput"]["params"])
    else:
        error_code, _ = map_error(c, "recover")
        eo = "{ 0, { 0 }, { 0 }, { { 0 } } }"
        ep = "{ 0, 0, { { 0 } } }"
    s += indent("{ %d, %d, %s, %d,\n%s,\n%s,\n%s,\n%s },\n" % (
        is_valid,
        1 if c["hostseckey"] else 0,
        byte_array(c["hostseckey"] or "00" * 32),
        len(bytes.fromhex(c["recoveryData"])),
        indent(byte_array(c["recoveryData"]) + ",", 1)[:-1],
        error_code,
        indent(eo + ",", 1)[:-1],
        indent(ep, 1),
    ))
s += "};\n"

if skipped_cases:
    s += "\n/* Skipped test cases (not expressible in the C API):\n"
    for case in skipped_cases:
        s += " * - %s\n" % case
    s += " */\n"

s += "\n#endif\n"
print(s)