import sys
import re
import argparse


def main():
    parser = argparse.ArgumentParser(
        description="A script to generate Go bindings for the CRIU magic values provided by magic.h"
    )
    parser.add_argument("src", help="Input path (magic.h)", type=str)
    parser.add_argument("dest", help="Output path (magic.go)", type=str)

    args = parser.parse_args()

    try:
        with open(args.src, "r") as src:
            magics = read_magics(src)
    except IOError:
        sys.exit("Failed to open magic.h")

    try:
        with open(args.dest, "w") as dest:
            write_magics(dest, magics)
    except IOError:
        sys.exit("Failed to open magic.py")


def read_magics(src):
    magics = {}
    # The magic.h header file contains lines like below:
    # #define MAGIC_NAME MAGIC_VALUE
    # where the magic value can be a hexadecimal or regular
    # integer value. The regex below matches lines in the
    # header and extracts the name and value as regex groups.
    pattern = re.compile(r"#define\s+(\S+)\s+(0x[0-9A-Fa-f]+|\d+)")

    for line in src:
        match = pattern.match(line)
        if match:
            name = match.group(1)
            value_str = match.group(2)
            # If the magic is defined as another magic, skip it
            if value_str in magics:
                continue
            try:
                value = int(value_str, 16)
                magics[name] = value
            except ValueError:
                print(f"Failed to parse magic {name}")
                continue

    return magics


def write_magics(dest, magics):
    dest.write(
        """// Code generated by magicgen. DO NOT EDIT.

package magic

type MagicMap struct {
\tByName  map[string]uint64
\tByValue map[uint64]string
}

func LoadMagic() MagicMap {
\tmagicMap := MagicMap{
\t\tByName:  make(map[string]uint64),
\t\tByValue: make(map[uint64]string),
\t}"""
    )

    for name in sorted(magics.keys()):
        value = magics[name]
        # Ignore RAW_IMAGE_MAGIC and V1 magic
        if value == 0 or value == 1:
            continue
        name = name.replace("_MAGIC", "")
        dest.write(f'\n\tmagicMap.ByName["{name}"] = {value}')
        dest.write(f'\n\tmagicMap.ByValue[{value}] = "{name}"')

    dest.write("\n\treturn magicMap\n}\n")


if __name__ == "__main__":
    main()
