#!/usr/bin/env python3
# build-aux/stack.c.gen - Analyze stack sizes for compiled objects
#
# Copyright (C) 2024  Luke T. Shumaker <lukeshu@lukeshu.com>
# SPDX-License-Identifier: AGPL-3.0-or-later

import os.path
import re
import sys
import typing

################################################################################
#
# Parse the "VCG" language
#
# https://www.rw.cdl.uni-saarland.de/people/sander/private/html/gsvcg1.html
#
# The formal syntax is found at
# ftp://ftp.cs.uni-sb.de/pub/graphics/vcg/vcg.tgz `doc/grammar.txt`.


class VCGElem:
    typ: str
    lineno: int
    attrs: dict[str, str]


def parse_vcg(reader: typing.TextIO) -> typing.Iterator[VCGElem]:
    re_beg = re.compile(r"(edge|node):\s*\{\s*")
    _re_tok = r"[a-zA-Z_][a-zA-Z0-9_]*"
    _re_str = r'"(?:[^\"]|\\.)*"'
    re_attr = re.compile(
        "(" + _re_tok + r")\s*:\s*(" + _re_tok + "|" + _re_str + r")\s*"
    )
    re_end = re.compile(r"\}\s*$")
    re_skip = re.compile(r"(graph:\s*\{\s*title\s*:\s*" + _re_str + r"\s*|\})\s*")
    re_esc = re.compile(r"\\.")

    for lineno, line in enumerate(reader):
        pos = 0

        def _raise(msg: str) -> typing.NoReturn:
            nonlocal lineno
            nonlocal line
            nonlocal pos
            e = SyntaxError(msg)
            e.lineno = lineno
            e.offset = pos
            e.text = line
            raise e

        if re_skip.fullmatch(line):
            continue

        elem = VCGElem()
        elem.lineno = lineno

        m = re_beg.match(line, pos=pos)
        if not m:
            _raise("does not look like a VCG line")
        elem.typ = m.group(1)
        pos = m.end()

        elem.attrs = {}
        while True:
            if re_end.match(line, pos=pos):
                break
            m = re_attr.match(line, pos=pos)
            if not m:
                _raise("unexpected character")
            k = m.group(1)
            v = m.group(2)
            if k in elem.attrs:
                _raise(f"duplicate key: {repr(k)}")
            if v.startswith('"'):

                def unesc(esc: re.Match[str]) -> str:
                    match esc.group(0)[1:]:
                        case "n":
                            return "\n"
                        case '"':
                            return '"'
                        case "\\":
                            return "\\"
                        case _:
                            _raise(f"invalid escape code {repr(esc.group(0))}")

                v = re_esc.sub(unesc, v[1:-1])
            elem.attrs[k] = v
            pos = m.end()

        yield elem


################################################################################
# Main analysis


class Node:
    # from .title (`static` and `__weak` functions are prefixed with
    # the compilation unit .c file.  For static functions that's fine,
    # but we'll have to handle it specially for __weak.).
    funcname: str
    # .label is "{funcname}\n{location}\n{nstatic} bytes (static}\n{ndynamic} dynamic objects"
    location: str
    nstatic: int
    ndynamic: int

    # edges with .sourcename set to this node
    calls: set[str]


def synthetic_node(name: str, nstatic: int, calls: set[str] = set()) -> Node:
    n = Node()

    n.funcname = name

    n.location = "<synthetic>"
    n.nstatic = nstatic
    n.ndynamic = 0

    n.calls = calls

    return n


def analyze(
    *,
    ci_fnames: list[str],
    extra_nodes: list[Node] = [],
    app_func_filters: dict[str, typing.Callable[[str], bool]],
    app_location_xform: typing.Callable[[str], str],
    app_indirect_callees: typing.Callable[[VCGElem], list[str]],
    app_skip_call: typing.Callable[[list[str], str], bool],
    cfg_max_call_depth: int,
) -> None:
    re_node_label = re.compile(
        r"(?P<funcname>[^\n]+)\n"
        + r"(?P<location>[^\n]+:[0-9]+:[0-9]+)\n"
        + r"(?P<nstatic>[0-9]+) bytes \(static\)\n"
        + r"(?P<ndynamic>[0-9]+) dynamic objects",
        flags=re.MULTILINE,
    )

    graph: dict[str, Node] = dict()
    qualified: dict[str, set[str]] = dict()

    def handle_elem(elem: VCGElem) -> None:
        match elem.typ:
            case "node":
                node = Node()
                node.calls = set()
                skip = False
                for k, v in elem.attrs.items():
                    match k:
                        case "title":
                            node.funcname = v
                        case "label":
                            if elem.attrs.get("shape", "") != "ellipse":
                                m = re_node_label.fullmatch(v)
                                if not m:
                                    raise ValueError(
                                        f"unexpected label value {repr(v)}"
                                    )
                                node.location = m.group("location")
                                node.nstatic = int(m.group("nstatic"))
                                node.ndynamic = int(m.group("ndynamic"))
                        case "shape":
                            if v != "ellipse":
                                raise ValueError(f"unexpected shape value {repr(v)}")
                            skip = True
                        case _:
                            raise ValueError(f"unknown edge key {repr(k)}")
                if not skip:
                    if node.funcname in graph:
                        raise ValueError(f"duplicate node {repr(node.funcname)}")
                    graph[node.funcname] = node
                    if ":" in node.funcname:
                        _, shortname = node.funcname.rsplit(":", 1)
                        if shortname not in qualified:
                            qualified[shortname] = set()
                        qualified[shortname].add(node.funcname)
            case "edge":
                caller: str | None = None
                callee: str | None = None
                for k, v in elem.attrs.items():
                    match k:
                        case "sourcename":
                            caller = v
                        case "targetname":
                            callee = v
                        case "label":
                            pass
                        case _:
                            raise ValueError(f"unknown edge key {repr(k)}")
                if caller is None or callee is None:
                    raise ValueError(f"incomplete edge: {repr(elem.attrs)}")
                if caller not in graph:
                    raise ValueError(f"unknown caller: {caller}")
                if callee == "__indirect_call":
                    for callee in app_indirect_callees(elem):
                        graph[caller].calls.add(callee)
                else:
                    graph[caller].calls.add(callee)
            case _:
                raise ValueError(f"unknown elem type {repr(elem.typ)}")

    for ci_fname in ci_fnames:
        with open(ci_fname, "r") as fh:
            for elem in parse_vcg(fh):
                handle_elem(elem)

    for node in extra_nodes:
        if node.funcname in graph:
            raise ValueError(f"duplicate node {repr(node.funcname)}")
        graph[node.funcname] = node

    missing: set[str] = set()

    print("/*")

    dbg = False

    def resolve_funcname(funcname: str) -> str | None:
        # Handle `ld --wrap` functions
        if f"__wrap_{funcname}" in graph:
            return f"__wrap_{funcname}"
        if funcname.startswith("__real_") and funcname[len("__real_") :] in graph:
            funcname = funcname[len("__real_") :]

        # Usual case
        if funcname in graph:
            return funcname

        # Handle `__weak` functions
        if funcname in qualified and len(qualified[funcname]) == 1:
            return sorted(qualified[funcname])[0]

        return None

    def nstatic(orig_funcname: str, chain: list[str] = []) -> int:
        nonlocal dbg
        funcname = resolve_funcname(orig_funcname)
        if not funcname:
            if app_skip_call(chain, orig_funcname):
                return 0
            missing.add(orig_funcname)
            return 0
        if app_skip_call(chain, funcname):
            return 0

        if len(chain) == cfg_max_call_depth:
            raise ValueError(f"max call depth exceeded: {chain+[funcname]}")

        node = graph[funcname]
        if dbg:
            print(f"//dbg: {funcname}\t{node.nstatic}")
        return node.nstatic + max(
            [0, *[nstatic(call, chain + [funcname]) for call in node.calls]]
        )

    for grp_name, grp_filter in app_func_filters.items():
        namelen = max(
            [len(app_location_xform(name)) for name in graph if grp_filter(name)]
            + [len(grp_name) + 4]
        )
        numlen = max(len(str(nstatic(name))) for name in graph if name.endswith("_cr"))
        sep1 = ("=" * namelen) + " " + "=" * numlen
        sep2 = ("-" * namelen) + " " + "-" * numlen

        print("= " + grp_name + " " + sep1[len(grp_name) + 3 :])

        nmax = 0
        nsum = 0
        for funcname in graph:
            if grp_filter(funcname):
                n = nstatic(funcname)
                print(
                    f"{app_location_xform(funcname).ljust(namelen)} {str(n).rjust(numlen)}"
                )
                if n > nmax:
                    nmax = n
                nsum += n

        print(sep2)
        print(f"{'Total'.ljust(namelen)} {str(nsum).rjust(numlen)}")
        print(f"{'Maximum'.ljust(namelen)} {str(nmax).rjust(numlen)}")

        print(sep1)

    for funcname in sorted(missing):
        print(f"warning: missing: {funcname}")

    print("*/")


################################################################################
# Application-specific code

re_location = re.compile(r"(?P<filename>.+):(?P<row>[0-9]+):(?P<col>[0-9]+)")


def read_source(location: str) -> str:
    m = re_location.fullmatch(location)
    if not m:
        raise ValueError(f"unexpected label value {repr(location)}")
    filename = m.group("filename")
    row = int(m.group("row")) - 1
    col = int(m.group("col")) - 1
    with open(m.group("filename"), "r") as fh:
        return fh.readlines()[row][col:].rstrip()


def main(
    *, arg_base_dir: str, arg_ci_fnames: list[str], arg_c_fnames: list[str]
) -> None:

    re_call_other = re.compile(r"(?P<func>[^(]+)\(.*")

    # The sbc-harness codebase #######################################

    vcalls: dict[str, set[str]] = {}
    re_vtable_start = re.compile(r"_vtable\s*=\s*\{")
    re_vtable_entry = re.compile(r"^\s+\.(?P<meth>\S+)\s*=\s*(?P<impl>\S+),.*")
    for fname in c_fnames:
        with open(fname, "r") as fh:
            in_vtable = False
            for line in fh:
                line = line.rstrip()
                if in_vtable:
                    if m := re_vtable_entry.fullmatch(line):
                        meth = m.group("meth")
                        impl = m.group("impl")
                        if impl == "NULL":
                            continue
                        if m.group("meth") not in vcalls:
                            vcalls[meth] = set()
                        vcalls[meth].add(impl)
                    if "}" in line:
                        in_vtable = False
                elif re_vtable_start.search(line):
                    in_vtable = True

    re_call_vcall = re.compile(r"VCALL\((?P<obj>[^,]+), (?P<meth>[^,)]+)[,)].*")

    def sbc_indirect_callees(loc: str, line: str) -> list[str] | None:
        if "/3rd-party/" in loc:
            return None
        if m := re_call_vcall.fullmatch(line):
            if m.group("meth") in vcalls:
                return sorted(vcalls[m.group("meth")])
            return [f"__indirect_call:{m.group('obj')}->vtable->{m.group('meth')}"]
        if "trigger->cb(trigger->cb_arg)" in line:
            return [
                "alarmclock_sleep_intrhandler",
                "w5500_tcp_alarm_handler",
                "w5500_udp_alarm_handler",
            ]
        if "/chan.h:" in loc and "front->dequeue(" in line:
            return [
                "_cr_chan_dequeue",
                "_cr_select_dequeue",
            ]
        return None

    def sbc_is_thread(name: str) -> bool:
        return name.endswith("_cr") or name == "main"

    def sbc_is_intrhandler(name: str) -> bool:
        return name in [
            "rp2040_hwtimer_intrhandler",
            "_cr_gdb_intrhandler",
            "hostclock_handle_sig_alarm",
            "hostnet_handle_sig_io",
        ]

    sbc_gpio_handlers = [
        "w5500_intrhandler",
    ]

    def sbc_skip_call(chain: list[str], call: str) -> bool:
        if (
            len(chain) > 1
            and chain[-1].endswith(":__assert_msg_fail")
            and call == "_log_printf"
            and any(c.endswith(":__assert_msg_fail") for c in chain[:-1])
        ):
            return True
        if call == "_cr_select_dequeue":
            return True
        return False

    # pico-sdk #######################################################

    def pico_is_intrhandler(name: str) -> bool:
        return name in [
            "gpio_default_irq_handler",
        ]

    def pico_indirect_callees(loc: str, line: str) -> list[str] | None:
        if "/3rd-party/pico-sdk/" not in loc or "/3rd-party/pico-sdk/lib/" in loc:
            return None
        m = re_call_other.fullmatch(line)
        call: str | None = m.group("func") if m else None

        match call:
            case "connect_internal_flash_func":
                return ["rom_func_lookup(ROM_FUNC_CONNECT_INTERNAL_FLASH)"]
            case "flash_exit_xip_func":
                return ["rom_func_lookup(ROM_FUNC_FLASH_EXIT_XIP)"]
            case "flash_range_erase_func":
                return ["rom_func_lookup(ROM_FUNC_FLASH_RANGE_ERASE)"]
            case "flash_flush_cache_func":
                return ["rom_func_lookup(ROM_FUNC_FLASH_FLUSH_CACHE)"]
            case "rom_table_lookup":
                return ["rom_hword_as_ptr(BOOTROM_TABLE_LOOKUP_OFFSET)"]
        if "/flash.c:" in loc and "boot2_copyout" in line:
            return ["_stage2_boot"]
        if "/gpio.c:" in loc and call == "callback":
            return sbc_gpio_handlers
        if "/printf.c:" in loc:
            if call == "out":
                return [
                    "_out_buffer",
                    "_out_null",
                    "_out_fct",
                ]
            if "->fct(" in line:
                return ["stdio_buffered_printer"]
        if "/stdio.c:" in loc:
            if call == "out_func":
                return [
                    "stdio_out_chars_crlf",
                    "stdio_out_chars_no_crlf",
                ]
            if call and (call.startswith("d->") or call.startswith("driver->")):
                _, meth = call.split("->", 1)
                match meth:
                    case "out_chars":
                        return ["stdio_uart_out_chars"]
                    case "out_flush":
                        return ["stdio_uart_out_flush"]
                    case "in_chars":
                        return ["stdio_uart_in_chars"]
        return None

    def pico_skip_call(chain: list[str], call: str) -> bool:
        if call == "_out_buffer" or call == "_out_fct":
            last = ""
            for pcall in chain:
                if pcall in [
                    "__wrap_sprintf",
                    "__wrap_snprintf",
                    "__wrap_vsnprintf",
                    "vfctprintf",
                ]:
                    last = pcall
            if last == "vfctprintf":
                return call != "_out_fct"
            else:
                return call == "_out_buffer"
        return False

    pico_nodes: list[Node] = [
        # src/rp2_common/pico_int64_ops/pico_int64_ops_aeabi.S
        synthetic_node("__aeabi_lmul", 4),
        # src/rp2_common/pico_divider/divider_hardware.S
        synthetic_node("__aeabi_uidiv", 0, {"__aeabi_uidivmod"}),
        synthetic_node("__aeabi_uidivmod", 0),  # TODO
        synthetic_node("__aeabi_uldivmod", 0),  # TODO
        # src/rp2040/boot_stage2/boot2_${name,,}.S for name=W25Q080,
        # controlled by `#define PICO_BOOT_STAGE2_{name} 1` in
        # src/boards/include/boards/pico.h
        synthetic_node("_stage2_boot", 0),  # TODO
        # https://github.com/raspberrypi/pico-bootrom-rp2040
        synthetic_node("rom_func_lookup(ROM_FUNC_CONNECT_INTERNAL_FLASH)", 0),  # TODO
        synthetic_node("rom_func_lookup(ROM_FUNC_FLASH_EXIT_XIP)", 0),  # TODO
        synthetic_node("rom_func_lookup(ROM_FUNC_FLASH_FLUSH_CACHE)", 0),  # TODO
        synthetic_node("rom_hword_as_ptr(BOOTROM_TABLE_LOOKUP_OFFSET)", 0),  # TODO
    ]

    # TinyUSB ########################################################

    tusb_config_fname = (
        arg_base_dir + "/cmd/sbc_harness/config/tusb_config.h"
    )  # TODO: FIXME
    re_tud_class = re.compile(
        r"^\s*#\s*define\s+(?P<k>CFG_TUD_(?:\S{3}|AUDIO|VIDEO|MIDI|VENDOR|USBTMC|DFU_RUNTIME|ECM_RNDIS))\s+(?P<v>\S+).*"
    )
    tusb_config: dict[str, bool] = {}
    with open(tusb_config_fname, "r") as fh:
        in_table = False
        for line in fh:
            line = line.rstrip()
            if m := re_tud_class.fullmatch(line):
                k = m.group("k")
                v = m.group("v")
                tusb_config[k] = bool(int(v))

    usbd_fname = next(
        fname for fname in c_fnames if fname.endswith("/tinyusb/src/device/usbd.c")
    )
    tud_drivers: dict[str, set[str]] = {}
    re_tud_entry = re.compile(
        r"^\s+\.(?P<meth>\S+)\s*=\s*(?P<impl>[a-zA-Z0-9_]+)(?:,.*)?"
    )
    re_tud_if1 = re.compile(r"^\s*#\s*if (\S+)\s*")
    re_tud_if2 = re.compile(r"^\s*#\s*if (\S+)\s*\|\|\s*(\S+)\s*")
    re_tud_endif = re.compile(r"^\s*#\s*endif\s*")
    with open(usbd_fname, "r") as fh:
        in_table = False
        enabled = True
        for line in fh:
            line = line.rstrip()
            if in_table:
                if m := re_tud_if1.fullmatch(line):
                    enabled = tusb_config[m.group(1)]
                elif m := re_tud_if2.fullmatch(line):
                    enabled = tusb_config[m.group(1)] or tusb_config[m.group(2)]
                elif re_tud_endif.fullmatch(line):
                    enabled = True
                if m := re_tud_entry.fullmatch(line):
                    meth = m.group("meth")
                    impl = m.group("impl")
                    if meth == "name" or not enabled:
                        continue
                    if meth not in tud_drivers:
                        tud_drivers[meth] = set()
                    if impl != "NULL":
                        tud_drivers[meth].add(impl)
                if line.startswith("}"):
                    in_table = False
            elif " _usbd_driver[] = {" in line:
                in_table = True

    def tud_indirect_callees(loc: str, line: str) -> list[str] | None:
        if "/tinyusb/" not in loc or "/tinyusb/src/host/" in loc or "_host.c:" in loc:
            return None
        m = re_call_other.fullmatch(line)
        assert m
        call = m.group("func")
        if call == "_ctrl_xfer.complete_cb":
            return [
                # "process_test_mode_cb",
                "tud_vendor_control_xfer_cb",
                *sorted(tud_drivers["control_xfer_cb"]),
            ]
        elif call.startswith("driver->"):
            return sorted(tud_drivers[call[len("driver->") :]])
        elif call == "event.func_call.func":
            # callback from usb_defer_func()
            return []

        return None

    def tud_skip_call(chain: list[str], call: str) -> bool:
        if call == "usbd_app_driver_get_cb":
            return True
        return False

    # newlib #########################################################

    newlib_nodes: list[Node] = [
        # malloc
        synthetic_node("free", 0),  # TODO
        synthetic_node("malloc", 0),  # TODO
        synthetic_node("realloc", 0),  # TODO
        synthetic_node("aligned_alloc", 0),  # TODO
        synthetic_node("reallocarray", 0),  # TODO
        # execution
        synthetic_node("abort", 0),  # TODO
        synthetic_node("longjmp", 0),  # TODO
        synthetic_node("setjmp", 0),  # TODO
        # <strings.h>
        synthetic_node("memcmp", 0),  # TODO
        synthetic_node("memcpy", 0),  # TODO
        synthetic_node("memset", 0),  # TODO
        synthetic_node("strlen", 0),  # TODO
        synthetic_node("strncpy", 0),  # TODO
        # other
        synthetic_node("random", 0),  # TODO
    ]

    # main ###########################################################

    def thread_filter(name: str) -> bool:
        return sbc_is_thread(name)

    def intrhandler_filter(name: str) -> bool:
        name = name.rsplit(":", 1)[-1]
        return sbc_is_intrhandler(name) or pico_is_intrhandler(name)

    def location_xform(loc: str) -> str:
        if not loc.startswith("/"):
            return loc
        parts = loc.split(":", 1)
        parts[0] = "./" + os.path.relpath(parts[0], arg_base_dir)
        return ":".join(parts)

    def indirect_callees(elem: VCGElem) -> list[str]:
        loc = elem.attrs.get("label", "")
        line = read_source(loc)

        ret = sbc_indirect_callees(loc, line)
        if ret is not None:
            return ret

        ret = pico_indirect_callees(loc, line)
        if ret is not None:
            return ret

        ret = tud_indirect_callees(loc, line)
        if ret is not None:
            return ret

        return [f"__indirect_call:" + location_xform(elem.attrs.get("label", ""))]

    def skip_call(chain: list[str], call: str) -> bool:
        return (
            sbc_skip_call(chain, call)
            or pico_skip_call(chain, call)
            or tud_skip_call(chain, call)
        )

    analyze(
        ci_fnames=arg_ci_fnames,
        extra_nodes=pico_nodes + newlib_nodes,
        app_func_filters={
            "Threads": thread_filter,
            "Interrupt handlers": intrhandler_filter,
        },
        app_location_xform=location_xform,
        app_indirect_callees=indirect_callees,
        app_skip_call=skip_call,
        cfg_max_call_depth=100,
    )


if __name__ == "__main__":
    base_dir = sys.argv[1]

    re_suffix = re.compile(r"\.c\.o(bj)?$")

    ci_fnames = [
        re_suffix.sub(".c.ci", fname)
        for fname in sys.argv[2:]
        if re_suffix.search(fname)
    ]

    c_fnames = [fname for fname in sys.argv[2:] if fname.endswith(".c")]

    main(
        arg_base_dir=base_dir,
        arg_ci_fnames=ci_fnames,
        arg_c_fnames=c_fnames,
    )