#! /usr/bin/env python3
import json
import sys

# Note: this script does not run as part of the build process.
# It is used to generate structs from the TDX global_metadata.json
# file, and functions to fill in said structs.  Rerun it if
# you need more fields.

TDX_STRUCTS = {
    "tdx_sys_info_version": [
        "BUILD_DATE",
        "BUILD_NUM",
        "MINOR_VERSION",
        "MAJOR_VERSION",
        "UPDATE_VERSION",
        "INTERNAL_VERSION",
    ],
    "tdx_sys_info_features": [
        "TDX_FEATURES0"
    ],
    "tdx_sys_info_tdmr": [
        "MAX_TDMRS",
        "MAX_RESERVED_PER_TDMR",
        "PAMT_4K_ENTRY_SIZE",
        "PAMT_2M_ENTRY_SIZE",
        "PAMT_1G_ENTRY_SIZE",
    ],
    "tdx_sys_info_cmr": [
        "NUM_CMRS", "CMR_BASE", "CMR_SIZE"
    ],
}


def print_struct(name, fields, file):
    print("struct %s {" % (name,), file=file)
    for f in fields:
        fname = f["Field Name"]
        element_bytes = int(f["Element Size (Bytes)"])
        element_bits = element_bytes * 8
        num_fields = int(f["Num Fields"])
        if num_fields == 1:
            print("\tu%d %s;" % (element_bits, fname.lower()), file=file)
        else:
            print(
                "\tu%d %s[%d];" % (element_bits, fname.lower(), num_fields), file=file
            )
    print("};", file=file)


def print_field(number, member, indent, file):
    print(
        "%sif (!ret && !(ret = read_sys_metadata_field(%s, &val))) out->%s = val;"
        % (indent, number, member),
        file=file,
    )


def print_function(name, fields, file):
    print("static int get_%s(struct %s *out)" % (name, name), file=file)
    print("{", file=file)

    print("\tint ret = 0;", file=file)
    print("\tu64 val;", file=file)
    for f in fields:
        num_fields = int(f["Num Fields"])
        if num_fields > 1:
            print("\tint i;", file=file)
            break

    print(file=file)
    for f in fields:
        num_fields = int(f["Num Fields"])
        if num_fields == 1:
            print_field(
                f["Base FIELD_ID (Hex)"],
                f["Field Name"].lower(),
                indent="\t",
                file=file,
            )
        else:
            print("\tfor (i = 0; i < %d; i++)" % (num_fields,), file=file)
            print_field(
                f["Base FIELD_ID (Hex)"] + " + i",
                f["Field Name"].lower() + "[i]",
                indent="\t\t",
                file=file,
            )

    print(file=file)
    print("\treturn ret;", file=file)
    print("}", file=file)


jsonfile = sys.argv[1]
hfile = sys.argv[2]
cfile = sys.argv[3]

with open(jsonfile, "r") as f:
    json_in = json.load(f)
    fields = {x["Field Name"]: x for x in json_in["Fields"]}

with open(hfile, "w") as f:
    print("/* Automatically generated by tdx.py */", file=f)
    print("#ifndef TDX_DATA_GENERATED_H", file=f)
    print("#define TDX_DATA_GENERATED_H 1", file=f)
    for name, field_names in TDX_STRUCTS.items():
        print(file=f)
        print_struct(name, [fields[x] for x in field_names], file=f)
    print(file=f)
    print("#endif", file=f)

with open(cfile, "w") as f:
    print("/* Automatically generated by tdx.py */", file=f)
    for name, field_names in TDX_STRUCTS.items():
        print(file=f)
        print_function(name, [fields[x] for x in field_names], file=f)