# Copyright 2026 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Utilities for viewing and diffing Sum Trees of memory usage."""

import copy
import jinja2
import json
import os
import pandas as pd
import pathlib
import string
import sys
from dataclasses import dataclass, field

_HAS_IPYTHON = False
try:
    if not 'unittest' in sys.modules:
        # Do not import IPython when running unit tests to avoid module name
        # conflicts. There is `import cProfile as profile` in IPython.
        from IPython.display import HTML, display

        _HAS_IPYTHON = True
except ImportError:
    _HAS_IPYTHON = False

from . import demangler

_MEMORY_USAGE_DIR = pathlib.Path(__file__).resolve().parent
_SRC_PATH = _MEMORY_USAGE_DIR.parents[3]
sys.path.append(str(_SRC_PATH / 'third_party/perfetto/python'))
from perfetto.trace_processor import TraceProcessor


@dataclass
class TreeNode:
    """Node in a memory usage hierarchy."""

    name: str
    value: int = 0
    delta: int = 0
    children: list['TreeNode'] = field(default_factory=list)

    @classmethod
    def from_dict(cls, data: dict) -> 'TreeNode':
        raw_children = data.get('children', [])
        children = [TreeNode.from_dict(child) for child in raw_children]
        params = {k: v for k, v in data.items() if k != 'children'}
        return cls(**params, children=children)


class MemoryUsageView:
    """Holds memory usage information possibly with multiple root nodes."""

    def __init__(self, roots: list[TreeNode]):
        self.roots = roots

    @classmethod
    def from_json(cls, json_data: str) -> 'MemoryUsageView':
        """Creates a MemoryUsageView from a JSON string.

        Provides a simple way to initialize the structure for testing.
        """
        data = json.loads(json_data)
        roots = [TreeNode.from_dict(item) for item in data]
        return cls(roots)

    @classmethod
    def from_heap_dump(
        cls, trace_file_name: str, demangle: bool = True, aggregate: bool = True
    ) -> 'MemoryUsageView':
        """Creates a MemoryUsageView from a trace file with a heap dump."""
        abs_path = os.path.abspath(trace_file_name)
        with TraceProcessor(trace=abs_path) as trace_processor:
            df = _extract_heap_dump(trace_processor)
            return cls.from_df(df, demangle=demangle, aggregate=aggregate)
        return None

    def to_json(self, **kwargs) -> str:
        """Converts the MemoryUsageView to a JSON string."""
        return json.dumps(self.roots, cls=TreeNodeEncoder, indent=0, **kwargs)

    @classmethod
    def from_df(
        cls, df: pd.DataFrame, demangle: bool = True, aggregate: bool = True
    ) -> 'MemoryUsageView':
        """Creates a MemoryUsageView from a DataFrame."""
        roots = []
        if demangle:
            with demangler.Demangler() as d:
                roots = _load_df(df, d, aggregate)
        else:
            roots = _load_df(df, None, aggregate)
        return cls(roots)

    @classmethod
    def from_comparison(
        cls, base: 'MemoryUsageView', new: 'MemoryUsageView'
    ) -> 'MemoryUsageView':
        """Duplicates the base view with deltas as compared to the new view."""
        return cls(_compare_node_lists(base.roots, new.roots))

    def toplevel_names(self) -> list[str]:
        return [root.name for root in self.roots]

    def toplevel_pretty_report(self) -> list[str]:
        """Generates pretty-printed memory usage for toplevel frames."""
        result = []
        for root in self.roots:
            pretty_size = _prettify_size(root.value)
            result.append((root.name, root.value, pretty_size))
        for triplet in sorted(result, key=lambda t: t[1], reverse=True):
            yield f'{triplet[0]}: {triplet[2]}'

    def display(self):
        if not _HAS_IPYTHON:
            print('Error: Could not import IPython module')
            return
        saved_roots = self.roots
        unique_root_id = f'metrics-root-{id(self)}'
        env = jinja2.Environment(
            loader=jinja2.FileSystemLoader(_MEMORY_USAGE_DIR)
        )
        env.policies['json.dumps_function'] = MemoryUsageView.to_json
        try:
            # Common heap profiles are a too deep for the default recursion
            # limit.
            previous_recursion_limit = sys.getrecursionlimit()
            sys.setrecursionlimit(10000)
            self.roots = copy.deepcopy(self.roots)
            _prettify_nodes(self.roots)
            template = env.get_template('memory_usage_table.html.j2')
            html_str = template.render(json_data=self, root_id=unique_root_id)
        finally:
            sys.setrecursionlimit(previous_recursion_limit)
            self.roots = saved_roots
        display(HTML(html_str))


class TreeNodeEncoder(json.JSONEncoder):
    """A JSON encoder for TreeNode objects."""

    def default(self, obj):
        if isinstance(obj, TreeNode):
            data = {
                'name': obj.name,
                'value': obj.value,
                'delta': obj.delta,
            }
            if obj.children:
                data['children'] = [
                    self.default(child) for child in obj.children
                ]
            return data
        return super().default(obj)


def _extract_heap_dump(trace_processor: TraceProcessor) -> pd.DataFrame:
    """Extracts a heap dump from a trace file.

    In the returned DataFrame rows with parent callsite_id must precede
    the child rows referencing them.
    """
    # The SQL schema is documented here:
    #   https://perfetto.dev/docs/analysis/sql-tables#stack_profile_callsite
    query = r"""
    WITH AggregatedAllocations AS (
      SELECT
        callsite_id,
        SUM(size) AS total_size_bytes
      FROM
        heap_profile_allocation
      WHERE
        size >= 0
      GROUP BY
        callsite_id
    )
    SELECT
      callsite.id AS callsite_id,
      callsite.parent_id AS parent_callsite_id,
      callsite.depth AS depth,

      frame.name AS frame_name,
      frame.id AS frame_id,

      -- Ensures 0 is returned instead of NULL when the callsite had no recorded
      -- allocations.
      COALESCE(allocations.total_size_bytes, 0) AS total_size_bytes
    FROM
      stack_profile_callsite AS callsite
      JOIN stack_profile_frame AS frame
      ON callsite.frame_id = frame.id

      -- Allocations are only available for leaf callsites. Normally allocations
      -- are callsites with name 'malloc' and 'posix_memalign'.
      LEFT JOIN AggregatedAllocations AS allocations
      ON callsite.id = allocations.callsite_id
    ORDER BY
      depth ASC
    """
    df = trace_processor.query(query).as_pandas_dataframe()
    return df


def _load_df(
    df: pd.DataFrame, d: demangler.Demangler, aggregate: bool
) -> list[TreeNode]:
    """Loads a DataFrame into a MemoryUsageView.roots."""
    # Nodes indexed by callsite_id.
    nodes: dict[int, TreeNode] = {}
    roots: list[TreeNode] = []
    for row in df.itertuples():
        frame_name = row.frame_name
        if pd.isna(frame_name) or not frame_name:
            frame_name = 'unknown'
        elif d:
            frame_name = d.demangle(frame_name)

        callsite_id = int(row.callsite_id)
        assert callsite_id not in nodes, (
            f'Duplicate callsite id: {row.callsite_id}'
        )

        node = TreeNode(name=frame_name, value=int(row.total_size_bytes))
        nodes[callsite_id] = node

        if int(row.depth) == 0:
            assert pd.isna(row.parent_callsite_id), (
                f'Row with zero depth and {row.parent_callsite_id=}'
            )
            roots.append(node)
            continue
        parent_id = int(row.parent_callsite_id)
        assert parent_id in nodes, (
            f'Missing parent callsite for {callsite_id}: {parent_id}'
        )
        parent_node = nodes[parent_id]
        parent_node.children.append(node)
    if aggregate:
        roots = _aggregate_nodes(roots)
    return roots


def _aggregate_nodes(nodes: list[TreeNode]) -> list[TreeNode]:
    """Recursively aggregates values the list of TreeNodes for viewing.

    Replaces same-level nodes with the same name by a single node. Sorts
    all children by value, largest first.

    Note: This function is not idempotent. Applying the aggregation twice will
    lead to wrong results.
    """
    # Group nodes by name, non-recursively.
    grouped_by_name = {}
    for node in nodes:
        name = node.name
        grouped_by_name.setdefault(name, [])
        grouped_by_name[name].append(node)

    # Replace each group of nodes with a single occurrence.
    name_to_grouped_node = {}
    for name, group_list in grouped_by_name.items():
        all_children = []

        # Sum up all individual node values.
        list_total_value = 0
        for node in group_list:
            list_total_value += node.value
            all_children.extend(node.children)

        # Recursively replace children with grouped ones.
        grouped_node = TreeNode(name=name, value=0, delta=0)
        new_children = []
        if all_children:
            new_children = _aggregate_nodes(all_children)
            grouped_node.children = new_children
            for child in new_children:
                grouped_node.value += child.value

        grouped_node.value += list_total_value
        name_to_grouped_node[name] = grouped_node

    # Sort the nodes by their cumulative value.
    return sorted(
        list(name_to_grouped_node.values()),
        key=lambda node: node.value,
        reverse=True,
    )


def _zip_by_name(
    left_nodes: list[TreeNode], right_nodes: list[TreeNode]
) -> list[TreeNode]:
    """Iterates over two lists of nodes, producing pairs with an item from each.

    Nodes are paired by name. This function assumes that each input list has no
    nodes with duplicate names. If for an element in either list there is no
    paired node (i.e. no node with the same name) in the other list, then None
    takes place of the non-existent node.
    """
    name_to_right_node = {n.name: n for n in right_nodes}
    processed_names = set()
    for node1 in left_nodes:
        name = node1.name
        processed_names.add(name)
        node2 = name_to_right_node.get(name)
        yield (node1, node2)
    for node2 in right_nodes:
        if node2.name not in processed_names:
            yield (None, node2)


def _compare_node_lists(
    nodes_base: list[TreeNode], nodes_new: list[TreeNode]
) -> list[TreeNode]:
    """Recursively creates a copy of base nodes with deltas to the new list.

    Puts value difference (from base list to the new list) as node.delta in the
    resulting tree(s). Assumes that each list (or list of children of a node
    does not have duplicated names). This can be achieved by applying
    _aggregate_nodes() above. Toplevel lists and all lists of children are
    individually sorted by delta (biggest to smallest) for simplest
    visualisation to emphasize biggest memory regressions first.
    """
    result = []
    for base, new in _zip_by_name(nodes_base, nodes_new):
        if not base:
            result.append(TreeNode(name=new.name, value=0, delta=new.value))
            continue
        if not new:
            zero_new = TreeNode(
                name=base.name, value=base.value, delta=-base.value
            )
            result.append(zero_new)
            continue
        merged_node = TreeNode(name=base.name, value=base.value)
        merged_node.delta = new.value - base.value
        merged_node.children = _compare_node_lists(base.children, new.children)
        result.append(merged_node)
    return sorted(result, key=lambda node: node.delta, reverse=True)


def _prettify_size(size_bytes: int):
    KIB = 1024
    TIB = KIB**4

    if size_bytes < KIB or size_bytes > TIB:
        # The JS display table replaces zero integers with empty cells. To allow
        # removing such clutter simply do not convert the result to str.
        return size_bytes

    for unit in ['KiB', 'MiB', 'GiB']:
        size_bytes /= KIB
        if size_bytes < KIB or unit == 'GiB':
            result = f'{size_bytes:.2f}'.rstrip('0').rstrip('.')
            return f'{result} {unit}'


def _prettify_nodes(nodes: list[TreeNode]):
    for node in nodes:
        node.value = _prettify_size(node.value)
        node.delta = _prettify_size(node.delta)
        _prettify_nodes(node.children)