#!/usr/bin/env python3
#
# Copyright (c) 2024 Meta Platforms
#
# SPDX-License-Identifier: Apache-2.0

import argparse
import os
import re
import sys

from elftools.elf.descriptions import (
    describe_symbol_type,
)
from elftools.elf.elffile import ELFFile


class gen_symtab_log:
    def __init__(self, debug=False):
        self.__debug = debug

    def debug(self, text):
        """Print debug message if debugging is enabled.

        Note - this function requires config global variable to be initialized.
        """
        if self.__debug:
            sys.stdout.write(os.path.basename(sys.argv[0]) + ": " + text + "\n")

    @staticmethod
    def error(text):
        sys.exit(os.path.basename(sys.argv[0]) + ": error: " + text + "\n")

    def set_debug(self, state):
        self.__debug = state


log = gen_symtab_log()


def parse_args():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
        allow_abbrev=False,
    )

    parser.add_argument("-k", "--kernel", required=True, help="Zephyr kernel image")
    parser.add_argument("-o", "--output", required=True, help="Output source file")
    parser.add_argument(
        "-d", "--debug", action="store_true", help="Print additional debugging information"
    )

    return parser.parse_args()


class symtab_entry:
    def __init__(self, addr, size, offset, name):
        self.addr = addr
        self.size = size
        self.offset = offset
        self.name = name

    def __eq__(self, other):
        return self.addr == other.addr


first_addr = 0
symtab_list = []


def sanitize_func_name(name):
    pattern = r'(^[a-zA-Z_][a-zA-Z0-9_]*)'
    match = re.match(pattern, name)
    if match:
        return match.group(0)
    else:
        log.error(f"Failed to sanitize function name: {name}")

    return name


def main():
    args = parse_args()
    log.set_debug(args.debug)

    with open(args.kernel, "rb") as rf:
        elf = ELFFile(rf)

        # Find the symbol table.
        symtab = elf.get_section_by_name('.symtab')

        for symbol in symtab.iter_symbols():
            symbol_type = describe_symbol_type(symbol['st_info']['type'])
            symbol_addr = symbol['st_value']
            symbol_size = symbol['st_size']

            if symbol_type == 'FUNC' and symbol_addr != 0:
                symbol_name = sanitize_func_name(symbol.name)
                dummy_offset = 0  # offsets will be calculated later after we know the first address
                entry = symtab_entry(symbol_addr, symbol_size, dummy_offset, symbol_name)
                # Prevent entries with duplicated addresses
                if entry not in symtab_list:
                    symtab_list.append(entry)

        # Sort the address in ascending order
        symtab_list.sort(key=lambda x: x.addr, reverse=False)

        # Get the address of the first symbol
        first_addr = symtab_list[0].addr

        for i, entry in enumerate(symtab_list):
            # Offset is calculated here
            entry.offset = entry.addr - first_addr

            # Debug print
            log.debug(f'{i:6d}: {hex(entry.addr)} {hex(entry.size)} {entry.name:.25s}')

    with open(args.output, 'w') as wf:
        print("/* AUTO-GENERATED by gen_symtab.py, do not edit! */", file=wf)
        print("", file=wf)
        print("#include <zephyr/linker/sections.h>", file=wf)
        print("#include <zephyr/debug/symtab.h>", file=wf)
        print("", file=wf)
        num = len(symtab_list) + 1
        print(
            f"const struct z_symtab_entry __symtab_entry z_symtab_entries[{num}] = {{",
            file=wf,
        )
        for i, entry in enumerate(symtab_list):
            print(f"\t/* ADDR: {hex(entry.addr)} SIZE: {hex(entry.size)} */", file=wf)
            print(
                f"\t[{i}] = {{.offset = {hex(entry.offset)}, .name = \"{entry.name}\"}},", file=wf
            )

        # Append a dummy entry at the end to facilitate the binary search
        if symtab_list[-1].size == 0:
            dummy_offset = f"{hex(symtab_list[-1].offset)} + sizeof(uintptr_t)"
        else:
            dummy_offset = f"{hex(symtab_list[-1].offset + symtab_list[-1].size)}"
        print("\t/* dummy entry */", file=wf)
        print(f"\t[{len(symtab_list)}] = {{.offset = {dummy_offset}, .name = \"?\"}},", file=wf)
        print("};\n", file=wf)

        print("const struct symtab_info __symtab_info z_symtab = {", file=wf)
        print(f"\t.first_addr = {hex(first_addr)},", file=wf)
        print(f"\t.length = {len(symtab_list)},", file=wf)
        print("\t.entries = z_symtab_entries,", file=wf)
        print("};\n", file=wf)


if __name__ == "__main__":
    main()