############################################################################
# tools/pynuttx/nxgdbmcp/src/gmcp/tools/memory.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 typing import Optional

from mcp.server.fastmcp import Context

from ..utils import _exec_command


def register_memory_tools(gdb_mcp):
    @gdb_mcp.tool()
    async def nxgdb_check_memory_leaks(
        ctx: Context, session_id: str, address: Optional[str] = None
    ) -> str:
        """Check for memory leaks in the target process.

        Analyzes allocated memory blocks to detect potential memory leaks.

        Args:
            ctx (Context): The context object containing runtime/session information.
            session_id (str): Identifier for the GDB debugging session.
            address (Optional[str]): Memory address to analyze. When provided, dumps all
                valid nodes referenced by the memory block containing this address.

        Returns:
            str: Memory leak analysis results showing leaked blocks and their details.
        """
        command = "memleak"
        command += f" --address {address}" if address is not None else ""
        return await _exec_command(ctx, session_id, command)

    @gdb_mcp.tool()
    async def nxgdb_memalloc_inspect(
        ctx: Context,
        session_id: str,
        address: Optional[str] = None,
        log: Optional[str] = None,
        pid: Optional[int] = None,
        min_index: Optional[int] = None,
        max_index: Optional[int] = None,
        biggest: Optional[int] = None,
        top: Optional[int] = None,
        size: Optional[int] = None,
        orphan: Optional[bool] = None,
        no_pool: Optional[bool] = None,
        no_heap: Optional[bool] = None,
        no_group: Optional[bool] = None,
        no_backtrace: Optional[bool] = None,
        no_reverse: Optional[bool] = None,
        no_pid: Optional[int] = None,
        is_free: Optional[bool] = None,
        sort: Optional[str] = None,
    ) -> str:
        """Dump and analyze memory allocations in the target process.

        Executes the `memdump` command to provide detailed inspection of memory
        allocations and heap state. Supports extensive filtering and formatting
        options for customized memory analysis.

        Args:
            ctx (Context): The context object containing runtime/session information.
            session_id (str): Identifier for the GDB debugging session.
            address (Optional[str]): Query which memory block contains this address.
                Example:
                `memdump --address 0x1234` to find the block containing address 0x1234.
            log (Optional[str]): Parse memory blocks from a log file generated by
                running memdump command on the device.
            pid (Optional[int]): Only dump memory allocated or freed by this specific
                task/process ID.
            min_index (Optional[int]): Minimum sequence number to include. Memory blocks
                outside the specified sequence range will be excluded.
            max_index (Optional[int]): Maximum sequence number to include. Memory blocks
                outside the specified sequence range will be excluded.
            biggest (Optional[int]): Sort by memory size (count * block size), equivalent
                to --sort size.
            top (Optional[int]): Display only the top N memory blocks.
            size (Optional[int]): Only dump memory blocks of this specific size.
            orphan (Optional[bool]): Only dump orphaned memory blocks (blocks without
                valid references).
            no_pool (Optional[bool]): Exclude pool memory, only show heap allocations.
            no_heap (Optional[bool]): Exclude heap memory, only show pool allocations.
            no_group (Optional[bool]): Disable memory block grouping. By default, blocks
                with same call stack, size, and PID are grouped together.
            no_backtrace (Optional[bool]): Skip printing memory block backtraces,
                significantly improving dump speed.
            no_reverse (Optional[bool]): Disable reverse sorting of results (default
                sorting is in reverse order).
            no_pid (Optional[int]): Exclude memory blocks with specific PID. Use -1
                to exclude mempool allocations.
            is_free (Optional[bool]): Dump free memory blocks instead of the default
                used memory blocks.
            sort (Optional[str]): Sort the dump results by specified criteria.
                Available options: 'size', 'nodesize', 'count', 'seq', 'address'.

        Returns:
            str: Detailed memory dump analysis showing allocation patterns, sizes,
                 backtraces, and other memory-related information.
        """
        command = "memdump"
        if address is not None:
            command += f" --address {address}"
        if log is not None:
            command += f" --log {log}"
        if pid is not None:
            command += f" --pid {pid}"
        if min_index is not None:
            command += f" --min {min_index}"
        if max_index is not None:
            command += f" --max {max_index}"
        if biggest is not None:
            command += " --biggest"
        if orphan is not None:
            command += " --orphan"
        if top is not None:
            command += f" --top {top}"
        if size is not None:
            command += f" --size {size}"
        if no_pool is not None:
            command += " --no-pool"
        if no_heap is not None:
            command += " --no-heap"
        if no_group is not None:
            command += " --no-group"
        if no_backtrace is not None:
            command += " --no-backtrace"
        if no_reverse is not None:
            command += " --no-reverse"
        if no_pid is not None:
            command += f" --no-pid {no_pid}"
        if is_free is not None:
            command += " --free"
        if sort is not None:
            command += f" --sort {sort}"
        return await _exec_command(ctx, session_id, command)

    @gdb_mcp.tool()
    async def nxgdb_find_memory_pattern(
        ctx: Context,
        session_id: str,
        pattern: str,
        memrange: Optional[str] = None,
        heap_only: Optional[bool] = None,
        globals_only: Optional[bool] = None,
    ) -> str:
        """Find memory addresses containing a specific pattern for memory corruption analysis.

        This tool searches for a specific pattern in memory to help analyze memory corruption
        issues. When memory has been corrupted/trampled with distinctive values, memfind can
        locate all occurrences of that pattern. Use with memdump -a <address> to investigate
        the code that might be causing the corruption.

        Common use case: When you detect memory corruption with a specific pattern/value,
        use memfind to locate all instances of that pattern in memory, then use memdump
        on those addresses to trace back to the source code that wrote those values.

        Args:
            ctx (Context): The context object containing runtime/session information.
            session_id (str): Identifier for the GDB debugging session.
            pattern (str): Pattern to search for in memory. Can be a hex value (e.g., "0xdeadbeef"),
                string pattern, or byte sequence that you suspect is causing memory corruption.
            memrange (Optional[str]): Specific memory range to search within (e.g., "0x1000-0x2000").
                If not specified, searches the entire accessible memory space.
            heap_only (Optional[bool]): If True, restricts search to heap memory only,
                excluding stack and global variables.
            globals_only (Optional[bool]): If True, restricts search to global variables only,
                excluding heap and stack memory.

        Returns:
            str: List of memory addresses where the pattern was found, along with context
                 information that can be used with memdump for further investigation.
        """
        command = f"memfind {pattern}"

        if memrange is not None:
            command += f" --memrange {memrange}"
        if heap_only is not None:
            command += " --heap-only"
        if globals_only is not None:
            command += " --globals-only"

        return await _exec_command(ctx, session_id, command)

    @gdb_mcp.tool()
    async def nxgdb_fetch_memory_layout(
        ctx: Context,
        session_id: str,
        heap_only: Optional[bool] = None,
        globals_only: Optional[bool] = None,
    ) -> str:
        """Show RAM range of heap and sections for memory layout analysis.

        This tool displays comprehensive memory range information including all memory regions
        in the current system. It automatically parses ELF data sections and analyzes memory
        ranges within g_mmheap, providing a quick overview of memory regions for debugging.

        Key features:
        - Displays all system memory ranges
        - Automatically parses ELF data sections
        - Analyzes g_mmheap memory ranges
        - Provides quick memory region overview for debugging

        Use this tool to:
        - Get an overview of the memory layout
        - Understand heap and global variable locations
        - Prepare for targeted memory analysis with other tools

        Args:
            ctx (Context): The context object containing runtime/session information.
            session_id (str): Identifier for the GDB debugging session.
            heap_only (Optional[bool]): If True, show only heap memory ranges,
                excluding other memory sections.
            globals_only (Optional[bool]): If True, show only global variable memory ranges,
                excluding heap and other sections.

        Returns:
            str: Detailed memory range information showing heap ranges, ELF data sections,
                 and other memory regions with their start/end addresses and properties.
        """
        command = "memrange"

        if heap_only is not None:
            command += " --heap-only"
        if globals_only is not None:
            command += " --globals-only"

        return await _exec_command(ctx, session_id, command)
