# build-aux/measurestack/util.py - Analyze stack sizes for compiled objects
#
# Copyright (C) 2024-2025  Luke T. Shumaker <lukeshu@lukeshu.com>
# SPDX-License-Identifier: AGPL-3.0-or-later

import re
import typing

from . import analyze, vcg
from .analyze import BaseName, Node, QName

# pylint: disable=unused-variable
__all__ = [
    "synthetic_node",
    "read_source",
    "get_zero_or_one",
    "re_call_other",
    "Plugin",
    "PluginApplication",
]


def synthetic_node(
    name: str, nstatic: int, calls: typing.Collection[str] = frozenset()
) -> Node:
    n = Node()

    n.funcname = QName(name)

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

    n.calls = dict((QName(c), False) for c in calls)

    return n


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 {location!r}")
    filename = m.group("filename")
    row = int(m.group("row")) - 1
    col = int(m.group("col")) - 1
    with open(filename, "r", encoding="utf-8") as fh:
        return fh.readlines()[row][col:].rstrip()


def get_zero_or_one(
    pred: typing.Callable[[str], bool], fnames: typing.Collection[str]
) -> str | None:
    count = sum(1 for fname in fnames if pred(fname))
    assert count < 2
    if count:
        return next(fname for fname in fnames if pred(fname))
    return None


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


class Plugin(typing.Protocol):
    def is_intrhandler(self, name: QName) -> bool: ...

    # init_array returns a list of functions that are placed in the
    # `.init_array.*` section; AKA functions marked with
    # `__attribute__((constructor))`.
    def init_array(self) -> typing.Collection[QName]: ...

    # extra_includes returns a list of functions that are never
    # called, but are included in the binary anyway.  This may because
    # it is an unused method in a used vtable.  This may be because it
    # is an atexit() callback (we never exit).
    def extra_includes(self) -> typing.Collection[BaseName]: ...

    def extra_nodes(self) -> typing.Collection[Node]: ...
    def indirect_callees(
        self, loc: str, line: str
    ) -> tuple[typing.Collection[QName], bool] | None: ...
    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]: ...


class PluginApplication:
    _location_xform: typing.Callable[[str], str]
    _plugins: list[Plugin]

    def __init__(
        self, location_xform: typing.Callable[[str], str], plugins: list[Plugin]
    ) -> None:
        self._location_xform = location_xform
        self._plugins = plugins

    def extra_nodes(self) -> typing.Collection[Node]:
        ret: list[Node] = []
        for plugin in self._plugins:
            ret.extend(plugin.extra_nodes())
        return ret

    def indirect_callees(
        self, elem: vcg.VCGElem
    ) -> tuple[typing.Collection[QName], bool]:
        loc = elem.attrs.get("label", "")
        line = read_source(loc)

        for plugin in self._plugins:
            ret = plugin.indirect_callees(loc, line)
            if ret is not None:
                return ret

        placeholder = "__indirect_call"
        if m := re_call_other.fullmatch(line):
            placeholder += ":" + m.group("func")
        placeholder += " at " + self._location_xform(elem.attrs.get("label", ""))
        return [QName(placeholder)], False

    def skipmodels(self) -> dict[BaseName, analyze.SkipModel]:
        ret: dict[BaseName, analyze.SkipModel] = {}
        for plugin in self._plugins:
            ret.update(plugin.skipmodels())
        return ret