#!/usr/bin/env python3
# Copyright (c) 2024 The Bitcoin developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.

import json
import os
import re
import sys
import zlib


def main():
    """
    Output structure
    [                             <packs_vector>
      { "name": pack_name         <pack>
        "tests" : [               <pack_tests>
          { "name": test_name,    <pack_test>
            "tests": [
              test_vector1,
              test_vector2,
            ],
            "scriptonly": [ "id1", "id2", ... ]
          }
        ]
      }
    ]
    """
    encode = False
    header_only = False
    vector_files = sys.argv[1:]
    packs_vector = []
    files_used = []
    for vec_file in vector_files:
        if vec_file == "--encode":
            encode = True
            continue
        if vec_file == '--header-only':
            assert encode, "--header-only may only be used with --encode"
            header_only = True
            continue
        assert os.path.isfile(vec_file)
        m = re.match(r'.*/bch_([^_]*)_([^_]*)/.*\.(\w*)\.json', vec_file)
        assert m
        with open(vec_file, "rt", encoding="utf8") as f:
            contents = f.read()
            json_cont = json.loads(contents)
            pack_name = m[1]
            standardness = m[2]
            results = m[3].endswith("_results")
            scriptonly = m[3].endswith("scriptonly")
            if results:
                continue  # Skip "results" files for now
            elif scriptonly:
                map_key = "scriptonly"
            else:
                map_key = "tests"
            test_name = standardness

            # Get/set pack_tests vector within packs_vector
            pack_tests = None
            for pack in packs_vector:
                if pack["name"] == pack_name:
                    pack_tests = pack["tests"]
                    break
            else:
                pack_tests = []
                packs_vector.append({"name": pack_name, "tests": pack_tests})
            assert pack_tests is not None

            # Update or create the named test within pack_test, and set the "scriptonly" or "tests" content
            for pack_test in pack_tests:
                if pack_test["name"] == test_name:
                    if map_key not in pack_test:
                        pack_test[map_key] = json_cont
                    else:
                        pack_test[map_key].extend(json_cont)
                    break
            else:
                pack_tests.append({"name": test_name,
                                   "tests": json_cont if map_key == "tests" else [],
                                   "scriptonly": json_cont if map_key == "scriptonly" else []})

            # Update file list, used in comments in generated file.
            parent_dir = os.path.split(os.path.split(vec_file)[0])[-1]
            fname = os.path.split(vec_file)[-1]
            files_used.append(os.path.sep.join([parent_dir, fname]))

    if encode:
        uncompressed_bytes = json.dumps(packs_vector).encode(encoding="utf8")
        uncompressed_size = len(uncompressed_bytes)
        compressed_bytes = zlib.compress(uncompressed_bytes)
        del uncompressed_bytes
        print("// This file is auto-generated by " + sys.argv[0])
        files_str = '\n    '.join(files_used)
        print(f"/*\nIncluded files: \n    {files_str}\n*/")
        print("#include <array>")
        print("#include <cstdint>\n")
        print("namespace json_tests {")
        # Note: Even in non-header-only mode this should appear in the .cpp file to ensure compiler emits the symbol
        print("extern const std::array<const uint8_t, " + str(len(compressed_bytes)) + "> libauth_test_vectors;")
        if header_only:
            print(f"inline constexpr size_t libauth_test_vectors_uncompressed_size = {uncompressed_size};")
        else:
            # Note: Initializer syntax with double-braces '{{' is preferred here to avoid some compiler warnings on some
            # compilers.
            print("const std::array<const uint8_t, " + str(len(compressed_bytes)) + "> libauth_test_vectors{{")

            ctr = 0

            def formatter(bb: int) -> str:
                nonlocal ctr
                ret = f"0x{bb:02x}"
                # Allow for up to 20 items per line
                if ctr >= 20:
                    ctr = 0
                    ret = "\n" + ret
                ctr += 1
                return ret

            print(", ".join(map(formatter, compressed_bytes)))
            print("}};")
        print("} // namespace json_tests")
    else:
        print(json.dumps(packs_vector, indent=2))


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: generate_libauth_test_vectors.py [--encode [--header-only]] [json_test_vector_file, ...]")
        sys.exit(1)

    main()
