# build-aux/measurestack/analyze.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 random
import re
import sys
import typing

from . import vcg

# Whether to print "//dbg-cache:" on cache writes
dbg_cache = False
# Whether to print the graph in a /* comment */ before processing it
dbg_dumpgraph = False
# Whether to print "//dbg-nstatic:" lines that trace nstatic() execution
dbg_nstatic = False
# Whether to disable nstatic() caching (but does NOT disable any cache-related debug logging)
dbg_nocache = False
# Whether to sort things for consistently-ordered execution, or shuffle things to detect bugs
dbg_sort: typing.Literal["unsorted", "sorted", "shuffled"] = "unsorted"

# pylint: disable=unused-variable
__all__ = [
    "BaseName",
    "QName",
    "UsageKind",
    "Node",
    "NodeHandler",
    "NodeHandleCB",
    "maybe_sorted",
    "handle_simple_node",
    "AnalyzeResultVal",
    "AnalyzeResultGroup",
    "AnalyzeResult",
    "analyze",
]


def dumps(x: typing.Any, depth: int = 0, compact: bool = False) -> str:
    match x:
        case int() | str() | None:
            return repr(x)
        case dict():
            if len(x) == 0:
                return "{}"
            ret = "{"
            if not compact:
                ret += "\n"
            for k, v in x.items():
                if not compact:
                    ret += "\t" * (depth + 1)
                ret += dumps(k, depth + 1, True)
                ret += ":"
                if not compact:
                    ret += " "
                ret += dumps(v, depth + 1, compact)
                ret += ","
                if not compact:
                    ret += "\n"
            if not compact:
                ret += "\t" * depth
            ret += "}"
            return ret
        case list():
            if len(x) == 0:
                return "[]"
            ret = "["
            if not compact:
                ret += "\n"
            for v in x:
                if not compact:
                    ret += "\t" * (depth + 1)
                ret += dumps(v, depth + 1, compact)
                ret += ","
                if not compact:
                    ret += "\n"
            if not compact:
                ret += "\t" * depth
            ret += "]"
            return ret
        case set():
            if len(x) == 0:
                return "set()"
            ret = "{"
            if not compact:
                ret += "\n"
            for v in x:
                if not compact:
                    ret += "\t" * (depth + 1)
                ret += dumps(v, depth + 1, compact)
                ret += ","
                if not compact:
                    ret += "\n"
            if not compact:
                ret += "\t" * depth
            ret += "}"
            return ret
        case _:
            if hasattr(x, "__dict__"):
                return f"{x.__class__.__name__}(*{dumps(x.__dict__, depth, compact)})"
            return f"TODO({x.__class__.__name__})"


# types ########################################################################


class BaseName:
    # class ##########################################################

    _interned: dict[str, "BaseName"] = {}

    def __new__(cls, content: str) -> "BaseName":
        if ":" in content:
            raise ValueError(f"invalid non-qualified name: {content!r}")
        content = sys.intern(content)
        if content not in cls._interned:
            self = super().__new__(cls)
            self._content = content
            cls._interned[content] = self
        return cls._interned[content]

    # instance #######################################################

    _content: str

    def __str__(self) -> str:
        return self._content

    def __repr__(self) -> str:
        return f"BaseName({self._content!r})"

    def __format__(self, fmt_spec: str, /) -> str:
        return repr(self)

    def __eq__(self, other: typing.Any) -> bool:
        assert isinstance(
            other, BaseName
        ), f"comparing BaseName with {other.__class__.__name__}"
        return self._content == other._content

    def __lt__(self, other: "BaseName") -> bool:
        return self._content < other._content

    def __hash__(self) -> int:
        return self._content.__hash__()

    def as_qname(self) -> "QName":
        return QName(self._content)


class QName:
    # class ##########################################################

    _interned: dict[str, "QName"] = {}

    def __new__(cls, content: str) -> "QName":
        content = sys.intern(content)
        if content not in cls._interned:
            self = super().__new__(cls)
            self._content = content
            self._base = None
            cls._interned[content] = self
        return cls._interned[content]

    # instance #######################################################

    _content: str
    _base: BaseName | None

    def __str__(self) -> str:
        return self._content

    def __repr__(self) -> str:
        return f"QName({self._content!r})"

    def __format__(self, fmt_spec: str, /) -> str:
        return repr(self)

    def __eq__(self, other: typing.Any) -> bool:
        assert isinstance(
            other, QName
        ), f"comparing QName with {other.__class__.__name__}"
        return self._content == other._content

    def __lt__(self, other: "QName") -> bool:
        return self._content < other._content

    def __hash__(self) -> int:
        return self._content.__hash__()

    def base(self) -> BaseName:
        if self._base is None:
            self._base = BaseName(self._content.rsplit(":", 1)[-1].split(".", 1)[0])
        return self._base


UsageKind: typing.TypeAlias = typing.Literal["static", "dynamic", "dynamic,bounded"]


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: QName
    # .label is "{funcname}\n{location}\n{nstatic} bytes (static}\n{ndynamic} dynamic objects"
    location: str
    usage_kind: UsageKind
    nstatic: int
    ndynamic: int

    # edges with .sourcename set to this node, val is if it's
    # OK/expected that the function be missing.
    calls: dict[QName, bool]


class AnalyzeResultVal(typing.NamedTuple):
    nstatic: int
    cnt: int


class AnalyzeResultGroup(typing.NamedTuple):
    rows: dict[QName, AnalyzeResultVal]


class AnalyzeResult(typing.NamedTuple):
    groups: dict[str, AnalyzeResultGroup]
    missing: set[QName]
    dynamic: set[QName]

    included_funcs: set[QName]


class NodeHandleCB(typing.Protocol):
    def __call__(
        self, chain: typing.Sequence[QName], missing_ok: bool = False
    ) -> tuple[int, bool]: ...


class NodeHandler(typing.Protocol):
    def __call__(
        self, handle: NodeHandleCB, chain: typing.Sequence[QName], node: Node
    ) -> tuple[int, bool]: ...


class Application(typing.Protocol):
    def extra_nodes(self) -> typing.Collection[Node]: ...
    def indirect_callees(
        self, elem: vcg.VCGElem
    ) -> tuple[typing.Collection[QName], bool]: ...
    def node_handlers(self) -> dict[BaseName, NodeHandler]: ...


# code #########################################################################

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 \((?P<usage_kind>static|dynamic|dynamic,bounded)\)\n"
    + r"(?P<ndynamic>[0-9]+) dynamic objects"
    + r"(?:\n.*)*",
    flags=re.MULTILINE,
)


class _Graph:
    graph: dict[QName, Node]
    qualified: dict[BaseName, QName]

    _resolve_cache: dict[QName, QName | None]

    def __init__(self) -> None:
        self._resolve_cache = {}

    def _resolve_funcname(self, funcname: QName) -> QName | None:
        s = str(funcname)
        is_qualified = ":" in s

        # Handle `ld --wrap` functions
        if not is_qualified:
            with_wrap = QName(f"__wrap_{s}")
            if with_wrap in self.graph:
                return with_wrap
            if s.startswith("__real_"):
                without_real = QName(s[len("__real_") :])
                if without_real in self.graph:
                    funcname = without_real

        # Usual case
        if funcname in self.graph:
            return funcname

        # Handle `__weak`/`[[gnu::weak]]` functions
        if not is_qualified:
            return self.qualified.get(BaseName(s))

        return None

    def resolve_funcname(self, funcname: QName) -> QName | None:
        if funcname not in self._resolve_cache:
            self._resolve_cache[funcname] = self._resolve_funcname(funcname)
        return self._resolve_cache[funcname]


if typing.TYPE_CHECKING:
    from _typeshed import SupportsRichComparisonT as _T_sortable

_T = typing.TypeVar("_T")


@typing.overload
def maybe_sorted(
    unsorted: typing.Iterable["_T_sortable"], /, *, key: None = None
) -> typing.Iterable["_T_sortable"]: ...
@typing.overload
def maybe_sorted(
    unsorted: typing.Iterable[_T], /, *, key: typing.Callable[[_T], "_T_sortable"]
) -> typing.Iterable[_T]: ...


def maybe_sorted(
    unsorted: typing.Iterable[_T],
    /,
    *,
    key: typing.Callable[[_T], "_T_sortable"] | None = None,
) -> typing.Iterable[_T]:
    match dbg_sort:
        case "unsorted":
            return unsorted
        case "sorted":
            return sorted(unsorted, key=key)  # type: ignore
        case "shuffled":
            ret = [*unsorted]
            random.shuffle(ret)
            return ret


def _make_graph(
    ci_fnames: typing.Collection[str],
    app: Application,
) -> _Graph:
    graph: dict[QName, Node] = {}
    qualified: dict[BaseName, set[QName]] = {}

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

    for ci_fname in maybe_sorted(ci_fnames):
        with open(ci_fname, "r", encoding="utf-8") as fh:
            for elem in vcg.parse_vcg(fh):
                handle_elem(elem)

    def sort_key(node: Node) -> QName:
        return node.funcname

    for node in maybe_sorted(app.extra_nodes(), key=sort_key):
        if node.funcname in graph:
            raise ValueError(f"duplicate node {node.funcname}")
        graph[node.funcname] = node

    ret = _Graph()
    ret.graph = graph
    ret.qualified = {}
    for bname, qnames in qualified.items():
        if len(qnames) == 1:
            ret.qualified[bname] = next(name for name in qnames)
    return ret


def handle_simple_node(
    handle: NodeHandleCB,
    chain: typing.Sequence[QName],
    node: Node,
    skip_p: typing.Callable[[QName], bool] | None = None,
) -> tuple[int, bool]:
    cacheable = True
    max_call_nstatic = 0
    for call_qname, call_missing_ok in maybe_sorted(node.calls.items()):
        if skip_p and skip_p(call_qname):
            if dbg_nstatic:
                print(f"//dbg-nstatic: {'- '*(len(chain)+1)}{call_qname}\tskip")
            continue
        call_nstatic, call_cacheable = handle(
            [*chain, node.funcname, call_qname], call_missing_ok
        )
        max_call_nstatic = max(max_call_nstatic, call_nstatic)
        if not call_cacheable:
            cacheable = False
    return node.nstatic + max_call_nstatic, cacheable


def analyze(
    *,
    ci_fnames: typing.Collection[str],
    app_func_filters: dict[str, typing.Callable[[QName], tuple[int, bool]]],
    app: Application,
    cfg_max_call_depth: int,
) -> AnalyzeResult:
    graphdata = _make_graph(ci_fnames, app)
    if dbg_dumpgraph:
        print(f"/* {dumps(graphdata)} */")

    missing: set[QName] = set()
    dynamic: set[QName] = set()
    included_funcs: set[QName] = set()

    track_inclusion: bool = True

    node_handlers = app.node_handlers()

    def default_node_handler(
        handle: NodeHandleCB, chain: typing.Sequence[QName], node: Node
    ) -> tuple[int, bool]:
        return handle_simple_node(handle, chain, node)

    nstatic_cache: dict[QName, int] = {}

    def node_handle_cb(
        chain: typing.Sequence[QName], missing_ok: bool = False
    ) -> tuple[int, bool]:
        nonlocal nstatic_cache
        assert len(chain) > 0

        if len(chain) > cfg_max_call_depth:
            raise ValueError(f"max call depth exceeded: {chain}")

        call_orig_qname = chain[-1]
        call_qname = graphdata.resolve_funcname(call_orig_qname)

        def dbglog(msg: str) -> None:
            if dbg_nstatic:
                print(
                    f"//dbg-nstatic: {'- '*(len(chain)-1)}{call_qname or call_orig_qname}\t{msg}"
                )

        if not call_qname:
            if not missing_ok:
                missing.add(call_orig_qname)
            dbglog("missing")
            nstatic_cache[call_orig_qname] = 0
            return 0, True

        assert call_qname in graphdata.graph
        if (not dbg_nocache) and call_qname in nstatic_cache:
            nstatic = nstatic_cache[call_qname]
            dbglog(f"total={nstatic} (cache-read)")
            return nstatic, True
        node = graphdata.graph[call_qname]
        dbglog(str(node.nstatic))
        if node.usage_kind == "dynamic" or node.ndynamic > 0:
            dynamic.add(call_qname)
        if track_inclusion:
            included_funcs.add(call_qname)

        handler = node_handlers.get(call_qname.base(), default_node_handler)
        nstatic, cacheable = handler(node_handle_cb, chain[:-1], node)
        if cacheable:
            dbglog(f"total={nstatic} (cache-write)")
            if call_qname not in nstatic_cache:
                if dbg_cache:
                    print(f"//dbg-cache: {call_qname} = {nstatic}")
                nstatic_cache[call_qname] = nstatic
            else:
                assert dbg_nocache
                assert nstatic_cache[call_qname] == nstatic
        else:
            dbglog(f"total={nstatic} (do-not-cache)")
        return nstatic, cacheable

    def nstatic(funcname: QName) -> int:
        return node_handle_cb([funcname])[0]

    groups: dict[str, AnalyzeResultGroup] = {}
    for grp_name, grp_filter in app_func_filters.items():
        rows: dict[QName, AnalyzeResultVal] = {}
        for funcname in graphdata.graph:
            cnt, track_inclusion = grp_filter(funcname)
            if cnt:
                rows[funcname] = AnalyzeResultVal(nstatic=nstatic(funcname), cnt=cnt)
        groups[grp_name] = AnalyzeResultGroup(rows=rows)

    return AnalyzeResult(
        groups=groups, missing=missing, dynamic=dynamic, included_funcs=included_funcs
    )