############################################################################
# tools/pynuttx/nxgdb/mm.py
#
# SPDX-License-Identifier: Apache-2.0
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.  The
# ASF licenses this file to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance with the
# License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
# License for the specific language governing permissions and limitations
# under the License.
#
############################################################################

from __future__ import annotations

import argparse
from typing import Generator, List, Tuple

import gdb

from . import autocompeletion, backtrace, lists, utils
from .backtrace import CONFIG_LIBC_BACKTRACE_DEPTH
from .protocols import mm as p
from .utils import Value

# Note we use mm_freenode_s to check if CONFIG_MM_RECORD_STACK is enabled instead
# of utils.get_symbol_value("CONFIG_MM_RECORD_STACK") because the latter may report
# wrong value on some platforms.

CONFIG_MM_RECORD_STACK = utils.has_field("struct mm_freenode_s", "stack")
CONFIG_MM_RECORD_PID = utils.has_field("struct mm_freenode_s", "pid")
CONFIG_MM_RECORD_SEQNO = utils.has_field("struct mm_freenode_s", "seqno")
CONFIG_MM_RECORD = (
    CONFIG_MM_RECORD_STACK or CONFIG_MM_RECORD_PID or CONFIG_MM_RECORD_SEQNO
)
MM_RECORD_STACK_DEPTH = CONFIG_LIBC_BACKTRACE_DEPTH if CONFIG_MM_RECORD_STACK else 0

mempool_record_s = utils.lookup_type("struct mempool_record_s")
mm_record_size = 0

g_sections = None

PID_MM_INVALID = -100
PID_MM_MEMPOOL = -1


def get_sections():
    global g_sections
    if g_sections is None:
        g_sections = gdb.execute("maintenance info sections", to_string=True)
    return g_sections


def mm_alignup(size: int) -> int:
    size_t = utils.lookup_type("uintptr_t")
    CONFIG_MM_DEFAULT_ALIGNMENT = 2 * size_t.sizeof
    align = CONFIG_MM_DEFAULT_ALIGNMENT
    # @todo: use the actual value of CONFIG_MM_DEFAULT_ALIGNMENT
    # align = utils.get_symbol_value("CONFIG_MM_DEFAULT_ALIGNMENT") or 2 * size_t.sizeof
    size = (size + align - 1) & ~(align - 1)
    return size


if mempool_record_s and CONFIG_MM_RECORD:
    mm_record_size = mm_alignup(mempool_record_s.sizeof)


class MemPoolBlock:
    """
    Memory pool block instance.
    """

    MAGIC_ALLOC = 0xAAAA_AAAA

    def __init__(
        self, addr: int, blocksize: int, overhead: int, pool: MemPool = None
    ) -> None:
        """
        Initialize the memory pool block instance.
        block: must be start address of the block,
        blocksize: block size without backtrace overhead,
        overhead: backtrace overhead size.
        pool: the MemPool instance that this block belongs to.
        """
        self.overhead = overhead
        self.from_pool = True
        self.is_orphan = False
        self.address = addr
        self.blocksize = int(blocksize)
        self.nodesize = int(blocksize) + self.overhead
        self.usersize = self.blocksize
        self.useraddress = self.address
        self.pool = pool
        # Lazy evaluation
        self._backtrace = self._pid = self._seqno = self._magic = self._record = None

    def __repr__(self) -> str:
        return f"block@{hex(self.address)},size:{self.blocksize},seqno:{self.seqno},pid:{self.pid}"

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

    def __hash__(self) -> int:
        return hash((self.pid, self.nodesize, self.backtrace))

    def __eq__(self, value: MemPoolBlock) -> bool:
        return (
            self.pid == value.pid
            and self.nodesize == value.nodesize
            and self.backtrace == value.backtrace
        )

    def contains(self, address: int) -> bool:
        """Check if the address is in block's range, excluding overhead"""
        return self.address <= address < self.address + self.blocksize

    @property
    def record(self) -> p.MemPoolBlock:
        if not self._record:
            addr = int(self.address)
            addr -= mm_record_size
            self._record = (
                gdb.Value(addr).cast(mempool_record_s.pointer()).dereference()
            )
        return self._record

    @property
    def is_free(self) -> bool:
        if not CONFIG_MM_RECORD:
            return self.pool.is_free(self) if self.pool else None

        if not self._magic:
            self._magic = int(self.record["magic"])

        return self._magic != self.MAGIC_ALLOC

    @property
    def seqno(self) -> int:
        if not self._seqno:
            self._seqno = (
                int(self.record["seqno"]) if CONFIG_MM_RECORD_SEQNO else PID_MM_INVALID
            )
        return self._seqno

    @property
    def pid(self) -> int:
        if not self._pid:
            self._pid = (
                int(self.record["pid"]) if CONFIG_MM_RECORD_PID else PID_MM_INVALID
            )
        return self._pid

    @property
    def backtrace(self) -> Tuple[int]:
        if MM_RECORD_STACK_DEPTH <= 0:
            return ()

        if not self._backtrace:
            self._backtrace = tuple(
                backtrace.BacktraceEntry(self.record["stack"]).get()
            )
        return self._backtrace

    @property
    def prevnode(self) -> MemPoolBlock:
        addr = self.address - self.nodesize
        return MemPoolBlock(addr, self.blocksize, self.overhead, pool=self.pool)

    @property
    def nextnode(self) -> MemPoolBlock:
        addr = self.address + self.nodesize
        return MemPoolBlock(addr, self.blocksize, self.overhead, pool=self.pool)

    def read_memory(self) -> memoryview:
        return gdb.selected_inferior().read_memory(self.address, self.blocksize)


class MemPool(Value, p.MemPool):
    """
    Memory pool instance.
    """

    def __init__(self, mpool: Value, name=None) -> None:
        if mpool.type.code == gdb.TYPE_CODE_PTR:
            mpool = mpool.dereference()
        super().__init__(mpool)
        self._blksize = None
        self._nfree = None
        self._nifree = None
        self._overhead = None
        self._free_blks = None

    def __repr__(self) -> str:
        return f"{self.name}@{hex(self.address)},size:{self.size}/{self['blocksize']},nused:{self.nused},nfree:{self.nfree}"

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

    @property
    def name(self) -> str:
        try:
            return self.procfs.name.string()
        except Exception:
            return "<noname>"

    @property
    def memranges(self) -> Generator[Tuple[int, int], None, None]:
        """Memory ranges of the pool"""
        sq_entry_t = utils.lookup_type("sq_entry_t")
        blksize = self.size

        if self.ibase:
            blks = int(self.interruptsize) // blksize
            base = int(self.ibase)
            yield (base, base + blks * blksize)

        if not self.equeue.head:
            return None

        # First queue has size of initialsize
        ninit = int(self.initialsize)
        ninit = ninit and (ninit - sq_entry_t.sizeof) // blksize
        nexpand = (int(self.expandsize) - sq_entry_t.sizeof) // blksize

        for entry in lists.NxSQueue(self.equeue):
            blks = ninit or nexpand
            ninit = 0
            yield (int(entry) - blks * blksize, int(entry))

    @property
    def nqueue(self) -> int:
        return lists.sq_count(self.equeue)

    @property
    def size(self) -> int:
        """Real block size including backtrace overhead"""
        if not self._blksize:
            blksize = self["blocksize"]
            if CONFIG_MM_RECORD:
                blksize = mm_alignup(blksize + mempool_record_s.sizeof)
            self._blksize = int(blksize)
        return self._blksize

    @property
    def overhead(self) -> int:
        if not self._overhead:
            self._overhead = self.size - int(self["blocksize"])
        return self._overhead

    @property
    def nwaiter(self) -> int:
        return (
            -int(self.waitsem.val.semcount) if self.wait and self.expandsize == 0 else 0
        )

    @property
    def nused(self) -> int:
        return int(self.nalloc)

    @property
    def free(self) -> int:
        return (self.nfree + self.nifree) * self.size

    @property
    def nfree(self) -> int:
        if not self._nfree:
            self._nfree = lists.sq_count(self.queue)
        return self._nfree + self.nifree

    @property
    def nifree(self) -> int:
        """Interrupt pool free blocks count"""
        if not self._nifree:
            self._nifree = lists.sq_count(self.iqueue)
        return self._nifree

    @property
    def total(self) -> int:
        nqueue = lists.sq_count(self.equeue)
        sq_entry_t = utils.lookup_type("sq_entry_t")
        blocks = self.nused + self.nfree
        return int(nqueue * sq_entry_t.sizeof + blocks * self.size)

    @property
    def blks(self) -> Generator[MemPoolBlock, None, None]:
        """Iterate over all blocks in the pool"""
        sq_entry_t = utils.lookup_type("sq_entry_t")
        blksize = self.size  # Real block size including backtrace overhead
        blocksize = self["blocksize"]

        def iterate(entry, nblocks):
            base = int(entry) - nblocks * blksize
            while nblocks > 0:
                yield MemPoolBlock(
                    base + mm_record_size, blocksize, self.overhead, pool=self
                )
                base += blksize
                nblocks -= 1

        if self.ibase:
            blks = int(self.interruptsize) // blksize
            yield from iterate(self.ibase + blks * blksize, blks)

        if not self.equeue.head:
            return None

        # First queue has size of initialsize
        ninit = int(self.initialsize)
        ninit = ninit and (ninit - sq_entry_t.sizeof) // blksize
        nexpand = (int(self.expandsize) - sq_entry_t.sizeof) // blksize

        for entry in lists.NxSQueue(self.equeue):
            yield from iterate(entry, ninit or nexpand)
            ninit = 0

    def contains(self, address: int) -> Tuple[bool, Value]:
        ranges = self.memranges
        if not ranges:
            return False, None

        for start, end in ranges:
            if start <= address < end:
                return True, None

    def find(self, address: int) -> Value:
        """Find the block that contains the given address"""
        sq_entry_t = utils.lookup_type("sq_entry_t")
        blksize = self.size
        blocksize = self["blocksize"]

        def get_blk(base):
            blkstart = base + (address - base) // blksize * blksize
            blkstart += mm_record_size
            return MemPoolBlock(blkstart, blocksize, self.overhead, pool=self)

        if self.ibase:
            # Check if it belongs to interrupt pool
            blks = int(self.interruptsize) // blksize
            base = int(self.ibase)
            if base <= address < base + blks * blksize:
                return get_blk(base)

        if not self.equeue.head:
            return None

        # First queue has size of initialsize
        ninit = int(self.initialsize)
        ninit = ninit and (ninit - sq_entry_t.sizeof) // blksize
        nexpand = (int(self.expandsize) - sq_entry_t.sizeof) // blksize

        for entry in lists.NxSQueue(self.equeue):
            blks = ninit or nexpand
            ninit = 0
            base = int(entry) - blks * blksize
            if base <= address < int(entry):
                return get_blk(base)

    def blks_free(self) -> Generator[MemPoolBlock, None, None]:
        """Iterate over all free blocks in the pool"""
        blocksize = self["blocksize"]
        for entry in lists.NxSQueue(self.queue):
            yield MemPoolBlock(int(entry), blocksize, self.overhead, pool=self)

    def blks_used(self) -> Generator[MemPoolBlock, None, None]:
        """Iterate over all used blocks in the pool"""
        return filter(lambda blk: not blk.is_free, self.blks)

    def is_free(self, blk: MemPoolBlock) -> bool:
        """Check if the given block is free in the pool"""
        if self._free_blks is None:
            try:
                self._free_blks = set(free.address for free in self.blks_free())
            except Exception:
                self._free_blks = set()

            # monitor GDB stop event to clear cache
            def clear_free_blks(event):
                self._free_blks = None

            gdb.events.stop.connect(clear_free_blks)

        # If self._free_blks is empty, it's high possible there's problem reading free blocks
        # instead of all memory is used.
        if not self._free_blks:
            return None

        return blk.address in self._free_blks


class MemPoolMultiple(Value, p.MemPoolMultiple):
    """
    Multiple level memory pool instance.
    """

    def __init__(self, mpool: Value, name=None) -> None:
        if mpool.type.code == gdb.TYPE_CODE_PTR:
            mpool = mpool.dereference()
        super().__init__(mpool)

    def __repr__(self) -> str:
        return f"Multiple Level Memory Pool: {self.address}"

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

    @property
    def pools(self) -> Generator[MemPool, None, None]:
        for pool in utils.ArrayIterator(self["pools"], self.npools):
            yield MemPool(pool)

    @property
    def free(self) -> int:
        return sum(pool.free for pool in self.pools)

    @property
    def chunks(self) -> Generator[MemPoolBlock, None, None]:
        for chunk in lists.sq_for_every(self.chunk_queue):
            chunk = chunk.cast(gdb.lookup_type("struct mpool_chunk_s"))
            yield chunk


class MMNode(gdb.Value, p.MMFreeNode):
    """
    One memory node in the memory manager heap, either free or allocated.
    The instance is always dereferenced to the actual node.
    """

    MM_ALLOC_BIT = 0x1
    MM_PREVFREE_BIT = 0x2
    MM_MASK_BIT = MM_ALLOC_BIT | MM_PREVFREE_BIT
    try:
        MM_SIZEOF_ALLOCNODE = utils.sizeof("struct mm_allocnode_s")
        # Although preceding can locates in the previous node, we still count it as overhead
        MM_ALLOCNODE_OVERHEAD = MM_SIZEOF_ALLOCNODE
        MM_MIN_SHIFT = utils.log2ceil(utils.sizeof("struct mm_freenode_s"))
        MM_MIN_CHUNK = 1 << MM_MIN_SHIFT
    except Exception:
        MM_SIZEOF_ALLOCNODE = 0
        MM_ALLOCNODE_OVERHEAD = 0
        MM_MIN_SHIFT = 0
        MM_MIN_CHUNK = 0

    def __init__(self, node: gdb.Value):
        if node.type.code == gdb.TYPE_CODE_PTR:
            node = node.dereference()
        self._backtrace = None
        self._address = None
        self._nodesize = None
        super().__init__(node)

    def __repr__(self):
        return (
            f"{hex(self.address)}({'F' if self.is_free else 'A'}{'F' if self.is_prev_free else 'A'})"
            f" size:{self.nodesize}/{self.prevsize if self.is_prev_free else '-'}"
            f" seq:{self.seqno} pid:{self.pid} "
        )

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

    def __hash__(self) -> int:
        return hash((self.pid, self.nodesize, self.backtrace))

    def __eq__(self, value: MMNode) -> bool:
        return (
            self.pid == value.pid
            and self.nodesize == value.nodesize
            and self.backtrace == value.backtrace
        )

    def contains(self, address):
        """Check if the address is in node's range, excluding overhead"""
        return self.useraddress <= address < self.useraddress + self.usersize

    def read_memory(self):
        return gdb.selected_inferior().read_memory(self.useraddress, self.usersize)

    @property
    def address(self) -> int:
        """Change 'void *' to int"""
        if not self._address:
            self._address = int(super().address)
        return self._address

    @property
    def useraddress(self) -> int:
        """Address of user memory, excluding overhead"""
        return self.address + self.overhead

    @property
    def prevsize(self) -> int:
        """Size of preceding chunk size"""
        return int(self["preceding"]) & ~MMNode.MM_MASK_BIT

    @property
    def nodesize(self) -> int:
        """Size of this chunk, including overhead"""
        if not self._nodesize:
            self._nodesize = int(self["size"]) & ~MMNode.MM_MASK_BIT
        return self._nodesize

    @property
    def usersize(self) -> int:
        """Size of this chunk, excluding overhead"""
        usersize = self.nodesize - MMNode.MM_ALLOCNODE_OVERHEAD

        # This node is allocted, thus the next node's preceding is user memory.
        usersize += 4
        return usersize

    @property
    def flink(self):
        # Only free node has flink and blink
        return MMNode(self["flink"]) if self.is_free and self["flink"] else None

    @property
    def blink(self):
        # Only free node has flink and blink
        return MMNode(self["blink"]) if self.is_free and self["blink"] else None

    @property
    def pid(self) -> int:
        if CONFIG_MM_RECORD_PID:
            return int(self["pid"])
        return PID_MM_INVALID

    @property
    def seqno(self) -> int:
        return int(self["seqno"]) if CONFIG_MM_RECORD_SEQNO else -1

    @property
    def backtrace(self) -> List[Tuple[int, str, str]]:
        if MM_RECORD_STACK_DEPTH <= 0:
            return ()

        # The free mm heap node does not record backtrace, the value may be illegal
        try:
            if not self._backtrace and (
                stack := backtrace.BacktraceEntry(self["stack"]).get()
            ):
                self._backtrace = tuple(stack)
        except gdb.MemoryError:
            self._backtrace = tuple(backtrace.BacktraceEntry(0).get())

        return self._backtrace

    @property
    def prevnode(self) -> MMNode:
        if not self.is_prev_free:
            return None

        addr = int(self.address) - self.prevsize
        type = utils.lookup_type("struct mm_freenode_s").pointer()
        return MMNode(gdb.Value(addr).cast(type))

    @property
    def nextnode(self) -> MMNode:
        if not self.nodesize:
            gdb.write(f"\n\x1b[31;1m Node corrupted: {self} \x1b[m\n")
            return None

        addr = int(self.address) + self.nodesize
        type = utils.lookup_type("struct mm_freenode_s").pointer()
        # Use gdb.Value for better performance
        return MMNode(gdb.Value(addr).cast(type))

    @property
    def is_free(self) -> bool:
        return not self["size"] & MMNode.MM_ALLOC_BIT

    @property
    def is_prev_free(self) -> bool:
        return self["size"] & MMNode.MM_PREVFREE_BIT

    @property
    def is_orphan(self) -> bool:
        # Report orphaned node and node likely to be orphaned(free-used-used-free)
        return self.is_prev_free or self.nextnode.is_free

    @property
    def from_pool(self) -> bool:
        return False

    @property
    def overhead(self) -> int:
        return MMNode.MM_ALLOCNODE_OVERHEAD


class MMHeap(Value, p.MMHeap):
    """
    One memory manager heap. It may contains multiple regions.
    """

    def __init__(self, heap: Value, name=None) -> None:
        mm_heap_s = utils.lookup_type("struct mm_heap_s")
        if isinstance(heap, int) or heap.type.code == gdb.TYPE_CODE_INT:
            heap = gdb.Value(heap).cast(mm_heap_s.pointer()).dereference()
        elif heap.type.code == gdb.TYPE_CODE_PTR:
            heap = heap.dereference()

        if heap.type != mm_heap_s:
            raise ValueError(f"Invalid heap type: {heap.type}")

        super().__init__(heap)

        self.name = name or "<noname>"
        self._regions = None

        # Check if heap node is accessible
        try:
            for start, end in self.regions:
                gdb.selected_inferior().read_memory(start.address, 1)
                gdb.selected_inferior().read_memory(end.address, 1)
        except gdb.MemoryError:
            raise ValueError(f"Heap node not accessible: {heap}")

    def __repr__(self) -> str:
        regions = [
            f"{hex(start.address)}~{hex(end.address)}" for start, end in self.regions
        ]
        return f"{self.name}@{self.address}, {int(self.heapsize) / 1024 :.1f}kB {self.nregions}regions: {','.join(regions)}"

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

    @property
    def curused(self) -> int:
        return int(self.mm_curused)

    @property
    def heapsize(self) -> int:
        return int(self.mm_heapsize)

    @property
    def free(self) -> int:
        return self.heapsize - self.curused

    @property
    def nregions(self) -> int:
        return int(utils.get_field(self, "mm_nregions", default=1))

    @property
    def mm_mpool(self) -> MemPoolMultiple:
        mpool = utils.get_field(self, "mm_mpool", default=None)
        return MemPoolMultiple(mpool) if mpool else None

    @property
    def regions(self) -> List[Tuple[MMNode, MMNode]]:
        if not self._regions:
            regions = self.nregions
            self._regions = []
            for start, end in zip(
                utils.ArrayIterator(self.mm_heapstart, regions),
                utils.ArrayIterator(self.mm_heapend, regions),
            ):
                self._regions.append((MMNode(start), MMNode(end)))
        return self._regions

    @property
    def nodes(self) -> Generator[MMNode, None, None]:
        for start, end in self.regions:
            node = start
            while node and node.address <= end.address:
                yield node
                node = node.nextnode

    def nodes_free(self) -> Generator[MMNode, None, None]:
        return filter(lambda node: node.is_free, self.nodes)

    def nodes_used(self) -> Generator[MMNode, None, None]:
        return filter(lambda node: not node.is_free, self.nodes)

    def contains(self, address: int) -> bool:
        ranges = [[int(start.address), int(end.address)] for start, end in self.regions]
        ranges[0][0] = int(self.address)  # The heap itself is also in the range
        return any(start <= address <= end for start, end in ranges)

    def find(self, address: int) -> MMNode:
        for node in self.nodes:
            if node.address <= address < node.address + node.nodesize:
                return node


def get_heaps() -> List[MMHeap]:
    # parse g_procfs_meminfo to get all heaps
    heaps = []
    meminfo: p.ProcfsMeminfoEntry = utils.gdb_eval_or_none("g_procfs_meminfo")
    if not meminfo and (heap := utils.parse_and_eval("g_mmheap")):
        try:
            heaps.append(MMHeap(heap))
        except Exception:
            pass

    try:
        while meminfo:
            try:
                heap = MMHeap(meminfo.heap, name=meminfo.name.string())
                heaps.append(heap)
            except Exception:
                pass

            meminfo = meminfo.next
    except gdb.MemoryError:
        # procfs not accessible
        pass

    return heaps


def get_pools(heaps: List[Value] = []) -> Generator[MemPool, None, None]:
    for heap in heaps or get_heaps():
        if not (mm_pool := heap.mm_mpool):
            continue

        for pool in mm_pool.pools:
            yield pool


def memory_range(heap=True, globals=True) -> List[Tuple[int, int]]:
    # Execute the GDB command to get section info
    memranges = []
    if globals:
        sections = get_sections()
        # Parse the output to find sections with ALLOC and LOAD
        for line in sections.splitlines():
            if "ALLOC" in line and "READONLY" not in line:
                parts = line.split()
                start = int(parts[1].split("->")[0], 16)
                end = int(parts[1].split("->")[1], 16)
                if start == end:
                    continue

                memranges.append((start, end))

        idle_topstack = int(utils.parse_and_eval("g_idle_topstack").cast("uintptr_t"))
        idle_stacksize = int(utils.parse_and_eval("CONFIG_IDLETHREAD_STACKSIZE"))
        memranges.append((idle_topstack - idle_stacksize, idle_topstack))

    # Get heaps from memdump
    if heap:
        for heap in get_heaps():
            for i in range(heap.nregions):
                start = int(heap["mm_heapstart"][i])
                end = int(heap["mm_heapend"][i]) + MMNode.MM_SIZEOF_ALLOCNODE

                if start == end:
                    continue

                # For the first region in the heap we need to compensate for the heap_s size
                if i == 0:
                    start = int(heap.address)

                for r in memranges:
                    # If the address range is already in memranges, skip
                    if r[0] <= start and r[1] >= end:
                        break
                    # If the new address range includes a range in memranges,
                    # delete the old one and add the new one
                    elif start <= r[0] and end >= r[1]:
                        memranges.remove(r)
                        memranges.append((start, end))
                        break
                else:
                    memranges.append((start, end))
    return sorted(memranges, key=lambda x: x[0])


def get_memrange(
    rangestr: str = "", heap_only=False, globals_only=False
) -> List[Tuple[int, int]]:
    """
    Parse memory range from string or get from heap and globals,
    the string should be in the format of "start1,end1,attr1 start2,end2,attr2".
    """

    memrange = []
    if rangestr:
        values = rangestr.replace('"', "").split(",")
        for i in range(0, len(values), 3):
            start = utils.parse_arg(values[i])
            end = utils.parse_arg(values[i + 1])
            memrange.append((start, end))
    else:
        memrange = memory_range(heap=not globals_only, globals=not heap_only)

    # Merge overlapping ranges
    merged_ranges = []
    for start, end in sorted(memrange):
        if merged_ranges and start <= merged_ranges[-1][1]:
            merged_ranges[-1] = (merged_ranges[-1][0], max(merged_ranges[-1][1], end))
        else:
            merged_ranges.append((start, end))

    return merged_ranges


def get_nodes_dict():
    """
    Return dict of all memory nodes, including memory pool.
    """

    # If there is a cached result, return it directly
    if hasattr(get_nodes_dict, "_cached_nodes"):
        return get_nodes_dict._cached_nodes

    nodes = []
    for heap in get_heaps():
        nodes.extend(
            {
                "name": heap.name,
                "address": node.address,
                "size": node.nodesize,
                "seqno": node.seqno,
                "pid": node.pid,
                "free": node.is_free,
                "from_pool": False,
                "backtrace": node.backtrace,
            }
            for node in heap.nodes
        )

        for pool in get_pools([heap]):
            nodes.extend(
                {
                    "name": f"{heap.name}@{blk.nodesize}pool",
                    "address": blk.address,
                    "size": blk.nodesize,
                    "seqno": blk.seqno,
                    "pid": blk.pid,
                    "free": blk.is_free,
                    "from_pool": True,
                    "backtrace": blk.backtrace,
                }
                for blk in pool.blks
            )

    # Cache the result
    get_nodes_dict._cached_nodes = nodes
    return nodes


class MMHeapInfo(gdb.Command):
    """Show basic heap information"""

    def __init__(self):
        super().__init__("mm heap", gdb.COMMAND_USER)

    @utils.dont_repeat_decorator
    def invoke(self, arg: str, from_tty: bool) -> None:
        for heap in get_heaps():
            regions = [(start.address, end.address) for start, end in heap.regions]
            gdb.write(f"{heap} - has {len(list(heap.nodes))} nodes, regions:")
            gdb.write(" ".join(f"{hex(start)}~{hex(end)}" for start, end in regions))
            gdb.write("\n")


@autocompeletion.complete
class MMPoolInfo(gdb.Command):
    """Show basic heap information"""

    def get_argparser(self):
        parser = argparse.ArgumentParser(description="Dump memory pool information.")
        parser.add_argument(
            "--heap",
            type=str,
            metavar="file",
            help="Which heap's pool to show",
            default=None,
        )
        return parser

    def __init__(self):
        super().__init__("mm pool", gdb.COMMAND_USER)
        utils.alias("mempool", "mm pool")
        self.parser = self.get_argparser()

    @utils.dont_repeat_decorator
    def invoke(self, arg: str, from_tty: bool) -> None:
        try:
            args = self.parser.parse_args(gdb.string_to_argv(arg))
        except SystemExit:
            return

        heaps = [utils.parse_and_eval(args.heap)] if args.heap else get_heaps()
        if not (pools := list(get_pools(heaps))):
            gdb.write("No pools found.\n")
            return

        count = len(pools)
        gdb.write(f"Total {count} pools\n")

        name_max = max(len(pool.name) for pool in pools) + 11  # 11: "@0x12345678"
        formatter = (
            "{:>%d} {:>11} {:>9} {:>9} {:>9} {:>9} {:>9} {:>9} {:>9} {:>9}\n" % name_max
        )
        head = (
            "",
            "total",
            "blocksize",
            "bsize",
            "overhead",
            "nused",
            "nfree",
            "nifree",
            "nwaiter",
            "nqueue",
        )

        gdb.write(formatter.format(*head))
        for pool in pools:
            gdb.write(
                formatter.format(
                    f"{pool.name}@{pool.address:#x}",
                    pool.total,
                    pool.blocksize,
                    pool.size,
                    pool.overhead,
                    pool.nused,
                    pool.nfree,
                    pool.nifree,
                    pool.nwaiter,
                    pool.nqueue,
                )
            )