"""
Utilities for profiling pwndbg.
"""

from __future__ import annotations

import argparse

import pwndbg.commands
import pwndbg.profiling
from pwndbg.commands import CommandCategory

parser = argparse.ArgumentParser(
    description="""
Utilities for profiling Pwndbg.

Check out the `./profiling` folder for other useful utilities.
Use `./profiling/print_stats.py` to generate a report from a `.pstats` file.
"""
)
subparsers = parser.add_subparsers(dest="command")
parser_start = subparsers.add_parser("start", description="Starts the profiler.")
parser_stop = subparsers.add_parser("stop", description="Stop the profiler and save the results.")
parser_stop.add_argument(
    "--file", type=str, default="pwndbg.pstats", help="Output file for profile data."
)


@pwndbg.commands.Command(parser, category=CommandCategory.PWNDBG)
def profiler(command, file="pwndbg.pstats") -> None:
    if command == "start":
        print("Starting profiler.")
        pwndbg.profiling.profiler.start()
    elif command == "stop":
        print("Stopped profiler. Wrote results to", file)
        print(
            "To analyze the results, use the ./profiling/print_stats.py script in the pwndbg repo."
        )
        pwndbg.profiling.profiler.stop(file)
