#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
# SPDX-License-Identifier: LGPL-2.1-or-later


import argparse
import keyword
import re
import sys
from typing import Dict, List, NamedTuple, Sequence, TextIO, cast


class DwarfConstant(NamedTuple):
    name: str
    value: int


class DwarfConstantType(NamedTuple):
    name: str
    constants: Sequence[DwarfConstant]


def parse_dwarf_constants(f: TextIO) -> Sequence[DwarfConstantType]:
    types: Dict[str, Dict[str, int]] = {
        type_name: {}
        for type_name in (
            "DW_ACCESS",
            "DW_ADDR",
            "DW_AT",
            "DW_ATE",
            "DW_CC",
            "DW_CFA",
            "DW_CHILDREN",
            "DW_DEFAULTED",
            "DW_DS",
            "DW_DSC",
            "DW_EH_PE",
            "DW_END",
            "DW_FORM",
            "DW_ID",
            "DW_IDX",
            "DW_INL",
            "DW_LANG",
            "DW_LLE",
            "DW_LNCT",
            "DW_LNE",
            "DW_LNS",
            "DW_MACINFO",
            "DW_MACRO",
            "DW_OP",
            "DW_ORD",
            "DW_RLE",
            "DW_SECT",
            "DW_TAG",
            "DW_UT",
            "DW_VIRTUALITY",
            "DW_VIS",
        )
    }

    for match in re.finditer(
        r"^\s*#\s*define\s+(" + "|".join(types) + r")_(\w+)\s+(\S+)",
        sys.stdin.read(),
        flags=re.MULTILINE,
    ):
        type_name = match.group(1)
        name = match.group(2)
        value = int(match.group(3), 0)
        if (type_name, name) in {
            # Typos in the wild that libdwarf includes but we don't want.
            ("DW_AT", "stride"),  # "DWARF3 (do not use)"
            ("DW_CFA", "low_user"),  # "Incorrect spelling, do not use"
            ("DW_TAG", "namelist_items"),  # "SGI misspelling/typo"
            ("DW_TAG", "template_type_param"),  # "DWARF2 inconsistent"
            ("DW_TAG", "template_value_param"),  # "DWARF2 inconsistent"
            # libdwarf probably included this one to be consistent with the
            # standard DWARF template_foo_parameter names, but it's called
            # DW_TAG_GNU_template_template_param everywhere else.
            ("DW_TAG", "GNU_template_template_parameter"),
            # This name isn't mentioned in any version of the DWARF standard.
            ("DW_CFA", "extended"),
        }:
            continue
        # Typos in libdwarf itself.
        elif (type_name, name) == ("DW_CFA", "high_user"):
            name = "hi_user"
        elif (type_name, name) == ("DW_IDX", "hi_user"):
            value = 0x3FFF
        elif (type_name, name) == ("DW_LANG", "Haskel"):
            name = "Haskell"
        if types[type_name].setdefault(name, value) != value:
            raise ValueError(f"{type_name}_{name} redefined with different value")

    result = [
        DwarfConstantType(
            name=type_name,
            constants=[DwarfConstant(name, value) for name, value in constants.items()],
        )
        for type_name, constants in types.items()
    ]

    def insert_after(
        type_name: str, after_name: str, insert_constant: DwarfConstant
    ) -> None:
        for constant_type in result:
            if constant_type.name == type_name:
                break
        else:
            raise ValueError()
        constants = cast(List[DwarfConstant], constant_type.constants)
        for i, constant in enumerate(constants):
            if constant.name == after_name:
                break
        else:
            raise ValueError()
        constants.insert(i + 1, insert_constant)

    insert_after("DW_EH_PE", "sdata8", DwarfConstant("signed", 0x8))
    insert_after("DW_EH_PE", "aligned", DwarfConstant("indirect", 0x80))

    return result


_DWARF_CONSTANTS_WANT_STR = {"DW_OP", "DW_TAG"}


def gen_dwarf_constants_h(
    dwarf_constants: Sequence[DwarfConstantType], f: TextIO
) -> None:
    f.write(
        """\
// Copyright (c) Meta Platforms, Inc. and affiliates.
// SPDX-License-Identifier: LGPL-2.1-or-later
// Generated by scripts/gen_dwarf_constants.py.

/**
 * @file
 *
 * DWARF constant definitions.
 *
 * This file defines the following for each known DWARF constant type:
 *
 * 1. An X macro defining all of the known names and values of the type:
 *    `DW_FOO_DEFINITIONS`.
 * 2. Enumerators defining the constants: `DW_FOO_a`, `DW_FOO_b`, etc.
 * 3. For select types, a function to translate a value to its name:
 *   `dw_foo_str()`.
 */

#ifndef DWARF_CONSTANTS_H
#define DWARF_CONSTANTS_H

#define X(name, value) name = value,
"""
    )
    for constant_type in dwarf_constants:
        f.write(
            f"""
#define {constant_type.name}_DEFINITIONS \\
"""
        )
        for i, constant in enumerate(constant_type.constants):
            end = " \\" if i < len(constant_type.constants) - 1 else ""
            f.write(
                f"\tX({constant_type.name}_{constant.name}, 0x{constant.value:x}){end}\n"
            )
        f.write(f"enum {{ {constant_type.name}_DEFINITIONS }};\n")
        if constant_type.name in _DWARF_CONSTANTS_WANT_STR:
            f.write(
                f"""\
#define {constant_type.name}_STR_UNKNOWN_FORMAT "{constant_type.name}_<0x%x>"
#define {constant_type.name}_STR_BUF_LEN (sizeof({constant_type.name}_STR_UNKNOWN_FORMAT) - 2 + 2 * sizeof(int))
/**
 * Get the name of a `{constant_type.name}` value.
 *
 * @return Static string if the value is known or @p buf if the value is
 * unknown.
 */
const char *{constant_type.name.lower()}_str(int value, char buf[static {constant_type.name}_STR_BUF_LEN]);
"""
            )

    f.write(
        """
#undef X

#endif /* DWARF_CONSTANTS_H */
"""
    )


def gen_dwarf_constants_c(
    dwarf_constants: Sequence[DwarfConstantType], f: TextIO
) -> None:
    f.write(
        """\
// Copyright (c) Meta Platforms, Inc. and affiliates.
// SPDX-License-Identifier: LGPL-2.1-or-later
// Generated by scripts/gen_dwarf_constants.py.

#include <stdio.h>

#include "dwarf_constants.h"

#define X(name, _) if (value == name) return #name;
"""
    )
    for constant_type in dwarf_constants:
        if constant_type.name in _DWARF_CONSTANTS_WANT_STR:
            f.write(
                f"""
const char *{constant_type.name.lower()}_str(int value, char buf[static {constant_type.name}_STR_BUF_LEN])
{{
	{constant_type.name}_DEFINITIONS
	snprintf(buf, {constant_type.name}_STR_BUF_LEN, {constant_type.name}_STR_UNKNOWN_FORMAT, value);
	return buf;
}}
"""
            )

    f.write(
        """
#undef X
"""
    )


def gen_tests_dwarf_py(dwarf_constants: Sequence[DwarfConstantType], f: TextIO) -> None:
    f.write(
        """\
# Copyright (c) Meta Platforms, Inc. and affiliates.
# SPDX-License-Identifier: LGPL-2.1-or-later
# Generated by scripts/gen_dwarf_constants.py.

import enum
"""
    )
    for constant_type in dwarf_constants:
        f.write(f"\n\nclass {constant_type.name}(enum.IntEnum):\n")
        for constant in constant_type.constants:
            name = constant.name
            if keyword.iskeyword(name):
                name += "_"
            f.write(f"    {name} = 0x{constant.value:X}")
            if name == "name":
                f.write("  # type: ignore")
            f.write("\n")


def main() -> None:
    argparse.ArgumentParser(
        description="Generate libdrgn/dwarf_constants.h, libdrgn/dwarf_constants.c, and tests/dwarf.py from libdwarf/src/lib/libdwarf/dwarf.h (read from standard input)"
    ).parse_args()

    dwarf_constants = parse_dwarf_constants(sys.stdin)
    with open("libdrgn/dwarf_constants.h", "w") as f:
        gen_dwarf_constants_h(dwarf_constants, f)
    with open("libdrgn/dwarf_constants.c", "w") as f:
        gen_dwarf_constants_c(dwarf_constants, f)
    with open("tests/dwarf.py", "w") as f:
        gen_tests_dwarf_py(dwarf_constants, f)


if __name__ == "__main__":
    main()