Lib/profiling/sampling/cli.py

"""Command-line interface for the sampling profiler."""

import argparse
import importlib.util
import locale
import os
import re
import selectors
import socket
import subprocess
import sys
import time
import webbrowser
from contextlib import nullcontext

from .errors import SamplingUnknownProcessError, SamplingModuleNotFoundError, SamplingScriptNotFoundError
from .sample import sample, sample_live, dump_stack, _is_process_running
from .dump import print_stack_dump
from .pstats_collector import PstatsCollector
from .stack_collector import CollapsedStackCollector, FlamegraphCollector, DiffFlamegraphCollector
from .heatmap_collector import HeatmapCollector
from .gecko_collector import GeckoCollector
from .jsonl_collector import JsonlCollector
from .binary_collector import BinaryCollector
from .binary_reader import BinaryReader
from .constants import (
    MICROSECONDS_PER_SECOND,
    PROFILING_MODE_ALL,
    PROFILING_MODE_WALL,
    PROFILING_MODE_CPU,
    PROFILING_MODE_GIL,
    PROFILING_MODE_EXCEPTION,
    SORT_MODE_NSAMPLES,
    SORT_MODE_TOTTIME,
    SORT_MODE_CUMTIME,
    SORT_MODE_SAMPLE_PCT,
    SORT_MODE_CUMUL_PCT,
    SORT_MODE_NSAMPLES_CUMUL,
)

try:
    from ._child_monitor import ChildProcessMonitor
except ImportError:
    # _remote_debugging module not available on this platform (e.g., WASI)
    ChildProcessMonitor = None

try:
    from .live_collector import LiveStatsCollector
except ImportError:
    LiveStatsCollector = None


class CustomFormatter(
    argparse.ArgumentDefaultsHelpFormatter,
    argparse.RawDescriptionHelpFormatter,
):
    """Custom formatter that combines default values display with raw description formatting."""
    pass


class DiffFlamegraphAction(argparse.Action):
    """Custom action for --diff-flamegraph that sets both format and baseline path."""
    def __call__(self, parser, namespace, values, option_string=None):
        namespace.format = 'diff_flamegraph'
        namespace.diff_baseline = values


_HELP_DESCRIPTION = """Sample a process's stack frames and generate profiling data.

Examples:
  # Run and profile a script
  `python -m profiling.sampling run script.py arg1 arg2`

  # Attach to a running process
  `python -m profiling.sampling attach 1234`

  # Dump a running process's current stack
  `python -m profiling.sampling dump 1234`

  # Live interactive mode for a script
  `python -m profiling.sampling run --live script.py`

  # Live interactive mode for a running process
  `python -m profiling.sampling attach --live 1234`

Use `python -m profiling.sampling <command> --help` for command-specific help."""


# Constants for socket synchronization
_SYNC_TIMEOUT_SEC = 15.0
_PROCESS_KILL_TIMEOUT_SEC = 5.0
_READY_MESSAGE = b"ready"
_RECV_BUFFER_SIZE = 1024
_BINARY_PROFILE_HEADER_SIZE = 64
_BINARY_PROFILE_MAGICS = (b"HCAT", b"TACH")

# Format configuration
FORMAT_EXTENSIONS = {
    "pstats": "pstats",
    "collapsed": "txt",
    "flamegraph": "html",
    "diff_flamegraph": "html",
    "gecko": "json",
    "heatmap": "html",
    "jsonl": "jsonl",
    "binary": "bin",
}

COLLECTOR_MAP = {
    "pstats": PstatsCollector,
    "collapsed": CollapsedStackCollector,
    "flamegraph": FlamegraphCollector,
    "diff_flamegraph": DiffFlamegraphCollector,
    "gecko": GeckoCollector,
    "heatmap": HeatmapCollector,
    "jsonl": JsonlCollector,
    "binary": BinaryCollector,
}

BROWSER_COMPATIBLE_FORMATS = ("flamegraph", "diff_flamegraph", "heatmap")


def _setup_child_monitor(args, parent_pid):
    # Build CLI args for child profilers (excluding --subprocesses to avoid recursion)
    child_cli_args = _build_child_profiler_args(args)

    # Build output pattern
    output_pattern = _build_output_pattern(args)

    return ChildProcessMonitor(
        pid=parent_pid,
        cli_args=child_cli_args,
        output_pattern=output_pattern,
    )


def _get_child_monitor_context(args, pid):
    if getattr(args, 'subprocesses', False):
        return _setup_child_monitor(args, pid)
    return nullcontext()


def _build_child_profiler_args(args):
    child_args = []

    # Sampling options
    hz = MICROSECONDS_PER_SECOND // args.sample_interval_usec
    child_args.extend(["-r", str(hz)])
    if args.duration is not None:
        child_args.extend(["-d", str(args.duration)])
    if args.all_threads:
        child_args.append("-a")
    if args.realtime_stats:
        child_args.append("--realtime-stats")
    if args.native:
        child_args.append("--native")
    if not args.gc:
        child_args.append("--no-gc")
    if args.opcodes:
        child_args.append("--opcodes")
    if args.async_aware:
        child_args.append("--async-aware")
        async_mode = getattr(args, 'async_mode', 'running')
        if async_mode != "running":
            child_args.extend(["--async-mode", async_mode])

    # Mode options
    mode = getattr(args, 'mode', 'wall')
    if mode != "wall":
        child_args.extend(["--mode", mode])

    # Format options (skip pstats as it's the default)
    if args.format == "diff_flamegraph":
        child_args.extend(["--diff-flamegraph", args.diff_baseline])
    elif args.format != "pstats":
        child_args.append(f"--{args.format}")

    return child_args


def _build_output_pattern(args):
    """Build output filename pattern for child profilers.

    The pattern uses {pid} as a placeholder which will be replaced with the
    actual child PID using str.replace(), so user filenames with braces are safe.
    """
    if args.outfile:
        # User specified output - add PID to filename
        base, ext = os.path.splitext(args.outfile)
        if ext:
            return f"{base}_{{pid}}{ext}"
        else:
            return f"{args.outfile}_{{pid}}"
    else:
        # Use default pattern based on format (consistent _ separator)
        extension = FORMAT_EXTENSIONS.get(args.format, "txt")
        if args.format == "heatmap":
            return "heatmap_{pid}"
        if args.format == "pstats":
            # pstats defaults to stdout, but for subprocesses we need files
            return "profile_{pid}.pstats"
        return f"{args.format}_{{pid}}.{extension}"


def _parse_mode(mode_string):
    """Convert mode string to mode constant."""
    mode_map = {
        "wall": PROFILING_MODE_WALL,
        "cpu": PROFILING_MODE_CPU,
        "gil": PROFILING_MODE_GIL,
        "exception": PROFILING_MODE_EXCEPTION,
    }
    return mode_map[mode_string]


def _check_process_died(process):
    """Check if process died and raise an error with stderr if available."""
    if process.poll() is None:
        return

    # Process died - try to get stderr for error message
    stderr_msg = ""
    if process.stderr:
        try:
            stderr_msg = process.stderr.read().decode().strip()
        except (OSError, UnicodeDecodeError):
            pass

    if stderr_msg:
        raise RuntimeError(stderr_msg)
    raise RuntimeError(f"Process exited with code {process.returncode}")


def _wait_for_ready_signal(sync_sock, process, timeout):
    """Wait for the ready signal from the subprocess, checking for early death."""
    deadline = time.monotonic() + timeout
    sel = selectors.DefaultSelector()
    sel.register(sync_sock, selectors.EVENT_READ)

    try:
        while True:
            _check_process_died(process)

            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise socket.timeout("timed out")

            if not sel.select(timeout=min(0.1, remaining)):
                continue

            conn, _ = sync_sock.accept()
            try:
                ready_signal = conn.recv(_RECV_BUFFER_SIZE)
            finally:
                conn.close()

            if ready_signal != _READY_MESSAGE:
                raise RuntimeError(f"Invalid ready signal received: {ready_signal!r}")
            return
    finally:
        sel.close()


def _run_with_sync(original_cmd, suppress_output=False):
    """Run a command with socket-based synchronization and return the process."""
    # Create a TCP socket for synchronization with better socket options
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sync_sock:
        # Set SO_REUSEADDR to avoid "Address already in use" errors
        sync_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sync_sock.bind(("127.0.0.1", 0))  # Let OS choose a free port
        sync_port = sync_sock.getsockname()[1]
        sync_sock.listen(1)
        sync_sock.settimeout(_SYNC_TIMEOUT_SEC)

        # Get current working directory to preserve it
        cwd = os.getcwd()

        # Build command using the sync coordinator
        target_args = original_cmd[1:]  # Remove python executable
        cmd = (
            sys.executable,
            "-m",
            "profiling.sampling._sync_coordinator",
            str(sync_port),
            cwd,
        ) + tuple(target_args)

        # Start the process with coordinator
        # When suppress_output=True (live mode), capture stderr so we can
        # report errors if the process dies before signaling ready.
        # When suppress_output=False (normal mode), let stderr inherit so
        # script errors print to the terminal.
        popen_kwargs = {}
        if suppress_output:
            popen_kwargs["stdin"] = subprocess.DEVNULL
            popen_kwargs["stdout"] = subprocess.DEVNULL
            popen_kwargs["stderr"] = subprocess.PIPE

        process = subprocess.Popen(cmd, **popen_kwargs)

        try:
            _wait_for_ready_signal(sync_sock, process, _SYNC_TIMEOUT_SEC)
        except socket.timeout:
            # If we timeout, kill the process and raise an error
            if process.poll() is None:
                process.terminate()
                try:
                    process.wait(timeout=_PROCESS_KILL_TIMEOUT_SEC)
                except subprocess.TimeoutExpired:
                    process.kill()
                    process.wait()
            raise RuntimeError(
                "Process failed to signal readiness within timeout"
            )

        return process


_RATE_PATTERN = re.compile(r'''
      ^                   # Start of string
      (                   # Group 1: The numeric value
          \d+             #   One or more digits (integer part)
          (?:\.\d+)?      #   Optional: decimal point followed by digits
      )                   #   Examples: "10", "0.5", "100.25"
      (                   # Group 2: Optional unit suffix
          hz              #   "hz" - hertz
          | khz           #   "khz" - kilohertz
          | k             #   "k" - shorthand for kilohertz
      )?                  #   Suffix is optional (bare number = Hz)
      $                   # End of string
  ''', re.VERBOSE | re.IGNORECASE)


def _parse_sampling_rate(rate_str: str) -> int:
    """Parse sampling rate string to microseconds."""
    rate_str = rate_str.strip().lower()

    match = _RATE_PATTERN.match(rate_str)
    if not match:
        raise argparse.ArgumentTypeError(
            f"Invalid sampling rate format: {rate_str}. "
            "Expected: number followed by optional suffix (hz, khz, k) with no spaces (e.g., 10khz)"
        )

    number_part = match.group(1)
    suffix = match.group(2) or ''

    # Determine multiplier based on suffix
    suffix_map = {
        'hz': 1,
        'khz': 1000,
        'k': 1000,
    }
    multiplier = suffix_map.get(suffix, 1)
    hz = float(number_part) * multiplier
    if hz <= 0:
        raise argparse.ArgumentTypeError(f"Sampling rate must be positive: {rate_str}")

    interval_usec = int(MICROSECONDS_PER_SECOND / hz)
    if interval_usec < 1:
        raise argparse.ArgumentTypeError(f"Sampling rate too high: {rate_str}")

    return interval_usec


def _add_sampling_options(parser):
    """Add sampling configuration options to a parser."""
    sampling_group = parser.add_argument_group("Sampling configuration")
    sampling_group.add_argument(
        "-r",
        "--sampling-rate",
        type=_parse_sampling_rate,
        default="1khz",
        metavar="RATE",
        dest="sample_interval_usec",
        help="sampling rate (e.g., 10000, 10khz, 10k)",
    )
    sampling_group.add_argument(
        "-d",
        "--duration",
        type=int,
        default=None,
        metavar="SECONDS",
        help="Sampling duration (default: run to completion)",
    )
    sampling_group.add_argument(
        "-a",
        "--all-threads",
        action="store_true",
        help="Sample all threads in the process instead of just the main thread",
    )
    sampling_group.add_argument(
        "--realtime-stats",
        action="store_true",
        help="Print real-time sampling statistics (Hz, mean, min, max) during profiling",
    )
    sampling_group.add_argument(
        "--native",
        action="store_true",
        help='Include artificial `<native>` frames to denote calls to non-Python code',
    )
    sampling_group.add_argument(
        "--no-gc",
        action="store_false",
        dest="gc",
        help='Don\'t include artificial `<GC>` frames to denote active garbage collection',
    )
    sampling_group.add_argument(
        "--opcodes",
        action="store_true",
        help="Gather bytecode opcode information for instruction-level profiling "
        "(shows which bytecode instructions are executing, including specializations).",
    )
    sampling_group.add_argument(
        "--async-aware",
        action="store_true",
        help="Enable async-aware profiling (uses task-based stack reconstruction)",
    )
    sampling_group.add_argument(
        "--subprocesses",
        action="store_true",
        help="Also profile subprocesses. Each subprocess gets its own profiler and output file.",
    )
    sampling_group.add_argument(
        "--blocking",
        action="store_true",
        help="Stop all threads in target process before sampling to get consistent snapshots. "
        "Uses thread_suspend on macOS and ptrace on Linux. Adds overhead but ensures memory "
        "reads are from a frozen state.",
    )


def _add_mode_options(parser):
    """Add mode options to a parser."""
    mode_group = parser.add_argument_group("Mode options")
    mode_group.add_argument(
        "--mode",
        choices=["wall", "cpu", "gil", "exception"],
        default="wall",
        help="Sampling mode: wall (all samples), cpu (only samples when thread is on CPU), "
        "gil (only samples when thread holds the GIL), "
        "exception (only samples when thread has an active exception). "
        "Incompatible with `--async-aware`",
    )
    mode_group.add_argument(
        "--async-mode",
        choices=["running", "all"],
        default="running",
        help='Async profiling mode: "running" (only running task) '
        'or "all" (all tasks including waiting). Requires `--async-aware`',
    )


def _add_format_options(parser, include_compression=True, include_binary=True):
    """Add output format options to a parser."""
    output_group = parser.add_argument_group("Output options")
    format_group = output_group.add_mutually_exclusive_group()
    format_group.add_argument(
        "--pstats",
        action="store_const",
        const="pstats",
        dest="format",
        help="Generate pstats output (default)",
    )
    format_group.add_argument(
        "--collapsed",
        action="store_const",
        const="collapsed",
        dest="format",
        help="Generate collapsed stack traces for flamegraphs",
    )
    format_group.add_argument(
        "--flamegraph",
        action="store_const",
        const="flamegraph",
        dest="format",
        help="Generate interactive HTML flamegraph visualization",
    )
    format_group.add_argument(
        "--gecko",
        action="store_const",
        const="gecko",
        dest="format",
        help="Generate Gecko format for Firefox Profiler",
    )
    format_group.add_argument(
        "--heatmap",
        action="store_const",
        const="heatmap",
        dest="format",
        help="Generate interactive HTML heatmap visualization with line-level sample counts",
    )
    format_group.add_argument(
        "--diff-flamegraph",
        metavar="BASELINE",
        action=DiffFlamegraphAction,
        help="Generate differential flamegraph comparing current profile to `BASELINE` binary file",
    )
    format_group.add_argument(
        "--jsonl",
        action="store_const",
        const="jsonl",
        dest="format",
        help="Generate newline-delimited JSON (JSONL) for programmatic consumers",
    )
    if include_binary:
        format_group.add_argument(
            "--binary",
            action="store_const",
            const="binary",
            dest="format",
            help="Generate high-performance binary format (use `replay` command to convert)",
        )
    parser.set_defaults(format="pstats", diff_baseline=None)

    if include_compression:
        output_group.add_argument(
            "--compression",
            choices=["auto", "zstd", "none"],
            default="auto",
            help="Compression for binary format: auto (use zstd if available), zstd, none",
        )

    output_group.add_argument(
        "-o",
        "--output",
        dest="outfile",
        help="Output path (default: `stdout` for `pstats` text; with `-o`, `pstats` is binary). "
        "Auto-generated for other formats. For heatmap: directory name (default: `heatmap_PID`)",
    )
    output_group.add_argument(
        "--browser",
        action="store_true",
        help=(
            "Automatically open HTML output "
            f"({', '.join('--' + f.replace('_', '-') for f in BROWSER_COMPATIBLE_FORMATS)}) "
            "in browser. "
            "When using `--subprocesses`, only the main process opens the browser"
        ),
    )


def _add_pstats_options(parser):
    """Add pstats-specific display options to a parser."""
    pstats_group = parser.add_argument_group("pstats format options")
    pstats_group.add_argument(
        "--sort",
        choices=[
            "nsamples",
            "tottime",
            "cumtime",
            "sample-pct",
            "cumul-pct",
            "nsamples-cumul",
            "name",
        ],
        default=None,
        help="Sort order for pstats output (default: nsamples)",
    )
    pstats_group.add_argument(
        "-l",
        "--limit",
        type=int,
        default=None,
        help="Limit the number of rows in the output (default: 15)",
    )
    pstats_group.add_argument(
        "--no-summary",
        action="store_true",
        help="Disable the summary section in the pstats output",
    )


def _add_dump_options(parser):
    """Add one-shot stack dump options to a parser."""
    dump_group = parser.add_argument_group("Dump options")
    dump_group.add_argument(
        "-a",
        "--all-threads",
        action="store_true",
        help="Dump all threads in the process instead of just the main thread",
    )
    dump_group.add_argument(
        "--native",
        action="store_true",
        help='Include artificial `<native>` frames to denote calls to non-Python code',
    )
    dump_group.add_argument(
        "--no-gc",
        action="store_false",
        dest="gc",
        help='Don\'t include artificial `<GC>` frames to denote active garbage collection',
    )
    dump_group.add_argument(
        "--opcodes",
        action="store_true",
        help="Show bytecode opcode names when available.",
    )
    dump_group.add_argument(
        "--async-aware",
        action="store_true",
        help="Enable async-aware stack reconstruction",
    )
    dump_group.add_argument(
        "--async-mode",
        choices=["running", "all"],
        default=argparse.SUPPRESS,
        help='Async stack mode: "running" (only running task) '
        'or "all" (all tasks including waiting, default for dump). '
        "Requires `--async-aware`",
    )
    dump_group.add_argument(
        "--blocking",
        action="store_true",
        help="Stop all threads in target process before dumping the stack.",
    )


def _sort_to_mode(sort_choice):
    """Convert sort choice string to SORT_MODE constant."""
    sort_map = {
        "nsamples": SORT_MODE_NSAMPLES,
        "tottime": SORT_MODE_TOTTIME,
        "cumtime": SORT_MODE_CUMTIME,
        "sample-pct": SORT_MODE_SAMPLE_PCT,
        "cumul-pct": SORT_MODE_CUMUL_PCT,
        "nsamples-cumul": SORT_MODE_NSAMPLES_CUMUL,
        "name": -1,
    }
    return sort_map.get(sort_choice, SORT_MODE_NSAMPLES)


def _capture_config_from_args(args):
    return {
        "all_threads": args.all_threads,
        "native": args.native,
        "gc": args.gc,
        "opcodes": args.opcodes,
        "blocking": args.blocking,
    }


def _create_collector(format_type, sample_interval_usec, skip_idle, opcodes=False,
                      mode=None, output_file=None, compression='auto',
                      diff_baseline=None, capture_config=None):
    """Create the appropriate collector based on format type.

    Args:
        format_type: The output format ('pstats', 'collapsed', 'flamegraph',
                    'gecko', 'heatmap', 'jsonl', 'binary', 'diff_flamegraph')
        sample_interval_usec: Sampling interval in microseconds
        skip_idle: Whether to skip idle samples
        opcodes: Whether to collect opcode information (only used by gecko format
                 for creating interval markers in Firefox Profiler)
        mode: Profiling mode for collectors that expose it in metadata
        output_file: Output file path (required for binary format)
        compression: Compression type for binary format ('auto', 'zstd', 'none')
        diff_baseline: Path to baseline binary file for differential flamegraph
        capture_config: Capture feature mapping for binary profiles and diffs

    Returns:
        A collector instance of the appropriate type
    """
    collector_class = COLLECTOR_MAP.get(format_type)
    if collector_class is None:
        raise ValueError(f"Unknown format: {format_type}")

    if format_type == "diff_flamegraph":
        if diff_baseline is None:
            raise ValueError("Differential flamegraph requires a baseline file")
        if not os.path.exists(diff_baseline):
            raise ValueError(f"Baseline file not found: {diff_baseline}")
        return collector_class(
            sample_interval_usec,
            baseline_binary_path=diff_baseline,
            skip_idle=skip_idle,
            mode=mode,
            capture_config=capture_config,
        )

    # Binary format requires output file and compression
    if format_type == "binary":
        if output_file is None:
            raise ValueError("Binary format requires an output file")
        return collector_class(output_file, sample_interval_usec, skip_idle=skip_idle,
                              compression=compression, mode=mode,
                              capture_config=capture_config)

    # Gecko format never skips idle (it needs both GIL and CPU data)
    # and is the only format that uses opcodes for interval markers
    if format_type == "gecko":
        skip_idle = False
        return collector_class(sample_interval_usec, skip_idle=skip_idle, opcodes=opcodes)

    if format_type == "jsonl":
        return collector_class(
            sample_interval_usec, skip_idle=skip_idle, mode=mode
        )

    return collector_class(sample_interval_usec, skip_idle=skip_idle)


def _generate_output_filename(format_type, pid):
    """Generate output filename based on format and PID.

    Args:
        format_type: The output format
        pid: Process ID

    Returns:
        Generated filename
    """
    extension = FORMAT_EXTENSIONS.get(format_type, "txt")
    # For heatmap, use cleaner directory name without extension
    if format_type == "heatmap":
        return f"heatmap_{pid}"
    return f"{format_type}_{pid}.{extension}"


def _open_in_browser(path):
    """Open a file or directory in the default web browser.

    Args:
        path: File path or directory path to open

    For directories (heatmap), opens the index.html file inside.
    """
    abs_path = os.path.abspath(path)

    # For heatmap directories, open the index.html file
    if os.path.isdir(abs_path):
        index_path = os.path.join(abs_path, 'index.html')
        if os.path.exists(index_path):
            abs_path = index_path
        else:
            print(f"Warning: Could not find index.html in {path}", file=sys.stderr)
            return

    file_url = f"file://{abs_path}"
    try:
        webbrowser.open(file_url)
    except Exception as e:
        print(f"Warning: Could not open browser: {e}", file=sys.stderr)


def _validate_replay_input_file(filename):
    """Validate that the replay input looks like a sampling binary profile."""
    try:
        with open(filename, "rb") as file:
            header = file.read(_BINARY_PROFILE_HEADER_SIZE)
    except OSError as exc:
        sys.exit(f"Error: Could not read input file {filename}: {exc}")

    if (
        len(header) < _BINARY_PROFILE_HEADER_SIZE
        or header[:4] not in _BINARY_PROFILE_MAGICS
    ):
        sys.exit(
            "Error: Input file is not a binary sampling profile. "
            "The replay command only accepts files created with --binary"
        )


def _replay_with_reader(args, reader):
    """Replay samples from an open binary reader."""
    info = reader.get_info()
    interval = info['sample_interval_us']

    print(f"Replaying {info['sample_count']} samples from {args.input_file}")
    print(f"  Sample interval: {interval} us")
    print(
        "  Compression: "
        f"{'zstd' if info.get('compression_type', 0) == 1 else 'none'}"
    )

    collector = _create_collector(
        args.format, interval, skip_idle=False,
        mode=info.get("mode"),
        diff_baseline=args.diff_baseline,
        capture_config=info.get("capture_config"),
    )

    def progress_callback(current, total):
        if total > 0:
            pct = current / total
            bar_width = 40
            filled = int(bar_width * pct)
            bar = '█' * filled + '░' * (bar_width - filled)
            print(
                f"\r  [{bar}] {pct*100:5.1f}% ({current:,}/{total:,})",
                end="",
                flush=True,
            )

    count = reader.replay_samples(collector, progress_callback)
    if hasattr(collector, "set_replay_stats"):
        collector.set_replay_stats(info)
    if hasattr(collector, "set_mode"):
        collector.set_mode(info.get("mode"))
    print()

    if args.format == "pstats":
        if args.outfile:
            collector.export(args.outfile)
        else:
            sort_choice = (
                args.sort if args.sort is not None else "nsamples"
            )
            limit = args.limit if args.limit is not None else 15
            sort_mode = _sort_to_mode(sort_choice)
            collector.print_stats(
                sort_mode, limit, not args.no_summary,
                info.get("mode") if info.get("mode") is not None
                else PROFILING_MODE_WALL
            )
    else:
        filename = (
            args.outfile
            or _generate_output_filename(args.format, os.getpid())
        )
        export_ok = collector.export(filename)

        # Auto-open browser for HTML output if --browser flag is set
        if (
            export_ok
            and args.format in BROWSER_COMPATIBLE_FORMATS
            and getattr(args, 'browser', False)
        ):
            _open_in_browser(filename)

    print(f"Replayed {count} samples")


def _handle_output(collector, args, pid, mode):
    """Handle output for the collector based on format and arguments.

    Args:
        collector: The collector instance with profiling data
        args: Parsed command-line arguments
        pid: Process ID (for generating filenames)
        mode: Profiling mode used
    """
    if args.format == "binary":
        # Binary format already wrote to file incrementally, just finalize
        collector.export(None)
        filename = collector.filename
        print(f"Binary profile written to {filename} ({collector.total_samples} samples)")
    elif args.format == "pstats":
        if args.outfile:
            # If outfile is a directory, generate filename inside it
            if os.path.isdir(args.outfile):
                filename = os.path.join(args.outfile, _generate_output_filename(args.format, pid))
                collector.export(filename)
            else:
                collector.export(args.outfile)
        else:
            # Print to stdout with defaults applied
            sort_choice = args.sort if args.sort is not None else "nsamples"
            limit = args.limit if args.limit is not None else 15
            sort_mode = _sort_to_mode(sort_choice)
            collector.print_stats(
                sort_mode, limit, not args.no_summary, mode
            )
    else:
        # Export to file
        if args.outfile and os.path.isdir(args.outfile):
            # If outfile is a directory, generate filename inside it
            filename = os.path.join(args.outfile, _generate_output_filename(args.format, pid))
        else:
            filename = args.outfile or _generate_output_filename(args.format, pid)
        export_ok = collector.export(filename)

        # Auto-open browser for HTML output if --browser flag is set
        if (
            export_ok
            and args.format in BROWSER_COMPATIBLE_FORMATS
            and getattr(args, 'browser', False)
        ):
            _open_in_browser(filename)


def _validate_args(args, parser):
    """Validate format-specific options and live mode requirements.

    Args:
        args: Parsed command-line arguments
        parser: ArgumentParser instance for error reporting
    """
    command = getattr(args, 'command', None)

    if command == "replay":
        return

    # Check if live mode is available
    if hasattr(args, 'live') and args.live and LiveStatsCollector is None:
        parser.error(
            "Live mode requires the curses module, which is not available."
        )

    # --subprocesses is incompatible with --live
    if hasattr(args, 'subprocesses') and args.subprocesses:
        if ChildProcessMonitor is None:
            parser.error(
                "--subprocesses is not available on this platform "
                "(requires _remote_debugging module)."
            )
        if hasattr(args, 'live') and args.live:
            parser.error("--subprocesses is incompatible with --live mode.")

    # Async-aware mode is incompatible with options that need thread data.
    if getattr(args, 'async_aware', False):
        issues = []
        if getattr(args, 'native', False):
            issues.append("--native")
        if not getattr(args, 'gc', True):
            issues.append("--no-gc")
        if getattr(args, 'format', None) == "binary":
            issues.append("--binary")
        if hasattr(args, 'mode') and args.mode != "wall":
            issues.append(f"--mode={args.mode}")
        if hasattr(args, 'all_threads') and args.all_threads:
            issues.append("--all-threads")
        if issues:
            parser.error(
                f"Options {', '.join(issues)} are incompatible with --async-aware. "
                "Async-aware profiling uses task-based stack reconstruction."
            )

    # --async-mode requires --async-aware when explicitly set
    if not getattr(args, 'async_aware', False):
        if command == "dump":
            # dump uses SUPPRESS default, so attr only exists if user passed it
            if hasattr(args, 'async_mode'):
                parser.error("--async-mode requires --async-aware to be enabled.")
        elif hasattr(args, 'async_mode') and args.async_mode != "running":
            parser.error("--async-mode requires --async-aware to be enabled.")

    if command == "dump":
        return

    # Warn about blocking mode with aggressive sampling intervals
    if args.blocking and args.sample_interval_usec < 100:
        print(
            f"Warning: --blocking with a {args.sample_interval_usec} µs interval will stop all threads "
            f"{1_000_000 // args.sample_interval_usec} times per second. "
            "Consider using --sampling-rate 1khz or lower to reduce overhead.",
            file=sys.stderr
        )

    # Live mode is incompatible with format options
    if hasattr(args, 'live') and args.live:
        if args.format != "pstats":
            format_flag = f"--{args.format}"
            parser.error(
                f"--live is incompatible with {format_flag}. Live mode uses a TUI interface."
            )

        # Live mode is also incompatible with pstats-specific options
        issues = []
        if args.sort is not None:
            issues.append("--sort")
        if args.limit is not None:
            issues.append("--limit")
        if args.no_summary:
            issues.append("--no-summary")

        if issues:
            parser.error(
                f"Options {', '.join(issues)} are incompatible with --live. "
                "Live mode uses a TUI interface with its own controls."
            )
        return

    # Validate gecko mode doesn't use non-wall mode
    if args.format == "gecko" and getattr(args, 'mode', 'wall') != "wall":
        parser.error(
            "--mode option is incompatible with --gecko. "
            "Gecko format automatically includes both GIL-holding and CPU status analysis."
        )

    # Validate --opcodes is only used with compatible formats
    opcodes_compatible_formats = ("live", "gecko", "flamegraph", "diff_flamegraph", "heatmap", "binary")
    if getattr(args, 'opcodes', False) and args.format not in opcodes_compatible_formats:
        parser.error(
            f"--opcodes is only compatible with {', '.join('--' + f for f in opcodes_compatible_formats)}."
        )

    # Validate pstats-specific options are only used with pstats format
    if args.format != "pstats":
        issues = []
        if args.sort is not None:
            issues.append("--sort")
        if args.limit is not None:
            issues.append("--limit")
        if args.no_summary:
            issues.append("--no-summary")

        if issues:
            format_flag = f"--{args.format}"
            parser.error(
                f"Options {', '.join(issues)} are only valid with --pstats, not {format_flag}"
            )


def main():
    """Main entry point for the CLI."""
    # Set locale for number formatting, restore on exit
    old_locale = locale.setlocale(locale.LC_ALL, None)
    locale.setlocale(locale.LC_ALL, "")
    try:
        _main()
    finally:
        locale.setlocale(locale.LC_ALL, old_locale)


def _main():
    # Create the main parser
    parser = argparse.ArgumentParser(
        description=_HELP_DESCRIPTION,
        formatter_class=CustomFormatter,
    )

    # Create subparsers for commands
    subparsers = parser.add_subparsers(
        dest="command", required=True, help="Command to run"
    )

    # === RUN COMMAND ===
    run_parser = subparsers.add_parser(
        "run",
        help="Run and profile a script or module",
        formatter_class=CustomFormatter,
        description="""Run and profile a Python script or module

Examples:
  # Run and profile a module
  `python -m profiling.sampling run -m mymodule arg1 arg2`

  # Generate flamegraph from a script
  `python -m profiling.sampling run --flamegraph -o output.html script.py`

  # Profile with custom rate and duration
  `python -m profiling.sampling run -r 5khz -d 30 script.py`

  # Save collapsed stacks to file
  `python -m profiling.sampling run --collapsed -o stacks.txt script.py`

  # Live interactive mode for a script
  `python -m profiling.sampling run --live script.py`""",
    )
    run_parser.add_argument(
        "-m",
        "--module",
        action="store_true",
        help="Run target as a module (like `python -m`)",
    )
    run_parser.add_argument(
        "target",
        help="Script file or module name to profile",
    )
    run_parser.add_argument(
        "args",
        nargs=argparse.REMAINDER,
        help="Arguments to pass to the script or module",
    )
    run_parser.add_argument(
        "--live",
        action="store_true",
        help="Interactive TUI profiler (top-like interface, press 'q' to quit, 's' to cycle sort)",
    )
    _add_sampling_options(run_parser)
    _add_mode_options(run_parser)
    _add_format_options(run_parser)
    _add_pstats_options(run_parser)

    # === ATTACH COMMAND ===
    attach_parser = subparsers.add_parser(
        "attach",
        help="Attach to and profile a running process",
        formatter_class=CustomFormatter,
        description="""Attach to a running process and profile it

Examples:
  # Profile all threads, sort by total time
  `python -m profiling.sampling attach -a --sort tottime 1234`

  # Live interactive mode for a running process
  `python -m profiling.sampling attach --live 1234`""",
    )
    attach_parser.add_argument(
        "pid",
        type=int,
        help="Process ID to attach to",
    )
    attach_parser.add_argument(
        "--live",
        action="store_true",
        help="Interactive TUI profiler (top-like interface, press 'q' to quit, 's' to cycle sort)",
    )
    _add_sampling_options(attach_parser)
    _add_mode_options(attach_parser)
    _add_format_options(attach_parser)
    _add_pstats_options(attach_parser)

    # === DUMP COMMAND ===
    dump_parser = subparsers.add_parser(
        "dump",
        help="Dump a running process's current stack",
        formatter_class=CustomFormatter,
        description="""Dump a running process's current Python stack

Examples:
  # Dump the main thread stack
  `python -m profiling.sampling dump 1234`

  # Dump all thread stacks
  `python -m profiling.sampling dump -a 1234`""",
    )
    dump_parser.add_argument(
        "pid",
        type=int,
        help="Process ID to dump",
    )
    _add_dump_options(dump_parser)

    # === REPLAY COMMAND ===
    replay_parser = subparsers.add_parser(
        "replay",
        help="Replay a binary profile and convert to another format",
        formatter_class=CustomFormatter,
        description="""Replay a binary profile file and convert to another format

Examples:
  # Convert binary to flamegraph
  `python -m profiling.sampling replay --flamegraph -o output.html profile.bin`

  # Convert binary to pstats and print to stdout
  `python -m profiling.sampling replay profile.bin`

  # Convert binary to gecko format
  `python -m profiling.sampling replay --gecko -o profile.json profile.bin`""",
    )
    replay_parser.add_argument(
        "input_file",
        help="Binary profile file to replay",
    )
    _add_format_options(replay_parser, include_compression=False, include_binary=False)
    _add_pstats_options(replay_parser)

    # Parse arguments
    args = parser.parse_args()

    # Validate arguments
    _validate_args(args, parser)

    # Command dispatch table
    command_handlers = {
        "run": _handle_run,
        "attach": _handle_attach,
        "dump": _handle_dump,
        "replay": _handle_replay,
    }

    # Execute the appropriate command
    handler = command_handlers.get(args.command)
    if handler:
        handler(args)
    else:
        parser.error(f"Unknown command: {args.command}")


def _handle_attach(args):
    """Handle the 'attach' command."""
    if not _is_process_running(args.pid):
        raise SamplingUnknownProcessError(args.pid)
    # Check if live mode is requested
    if args.live:
        _handle_live_attach(args, args.pid)
        return

    # Use PROFILING_MODE_ALL for gecko format
    mode = (
        PROFILING_MODE_ALL
        if args.format == "gecko"
        else _parse_mode(args.mode)
    )

    # Determine skip_idle based on mode
    skip_idle = (
        mode != PROFILING_MODE_WALL if mode != PROFILING_MODE_ALL else False
    )

    output_file = None
    if args.format == "binary":
        output_file = args.outfile or _generate_output_filename(args.format, args.pid)

    # Create the appropriate collector
    collector = _create_collector(
        args.format, args.sample_interval_usec, skip_idle, args.opcodes, mode,
        output_file=output_file,
        compression=getattr(args, 'compression', 'auto'),
        diff_baseline=args.diff_baseline,
        capture_config=_capture_config_from_args(args),
    )

    with _get_child_monitor_context(args, args.pid):
        collector = sample(
            args.pid,
            collector,
            duration_sec=args.duration,
            all_threads=args.all_threads,
            realtime_stats=args.realtime_stats,
            mode=mode,
            async_aware=args.async_mode if args.async_aware else None,
            native=args.native,
            gc=args.gc,
            opcodes=args.opcodes,
            blocking=args.blocking,
        )
        _handle_output(collector, args, args.pid, mode)


def _handle_dump(args):
    if not _is_process_running(args.pid):
        raise SamplingUnknownProcessError(args.pid)

    # Async-aware reconstruction requires wall mode so every thread is sampled,
    # not just those holding the GIL or on CPU.
    mode = PROFILING_MODE_WALL if args.async_aware else PROFILING_MODE_ALL
    async_mode = getattr(args, "async_mode", "all") if args.async_aware else None
    try:
        stack_frames = dump_stack(
            args.pid,
            all_threads=args.all_threads,
            mode=mode,
            async_aware=async_mode,
            native=args.native,
            gc=args.gc,
            opcodes=args.opcodes,
            blocking=args.blocking,
        )
    except ProcessLookupError:
        sys.exit(
            f"No stack dump collected - process {args.pid} exited before "
            "its stack could be read."
        )

    print_stack_dump(stack_frames, pid=args.pid)


def _handle_run(args):
    """Handle the 'run' command."""
    # Validate target exists before launching subprocess
    if args.module:
        # Temporarily add cwd to sys.path so we can find modules in the
        # current directory, matching the coordinator's behavior
        cwd = os.getcwd()
        added_cwd = False
        if cwd not in sys.path:
            sys.path.insert(0, cwd)
            added_cwd = True
        try:
            if importlib.util.find_spec(args.target) is None:
                raise SamplingModuleNotFoundError(args.target)
        finally:
            if added_cwd:
                sys.path.remove(cwd)
    else:
        if not os.path.exists(args.target):
            raise SamplingScriptNotFoundError(args.target)

    # Check if live mode is requested
    if args.live:
        _handle_live_run(args)
        return

    # Build the command to run
    if args.module:
        cmd = (sys.executable, "-m", args.target, *args.args)
    else:
        cmd = (sys.executable, args.target, *args.args)

    # Run with synchronization
    try:
        process = _run_with_sync(cmd, suppress_output=False)
    except RuntimeError as e:
        sys.exit(f"Error: {e}")

    # Use PROFILING_MODE_ALL for gecko format
    mode = (
        PROFILING_MODE_ALL
        if args.format == "gecko"
        else _parse_mode(args.mode)
    )

    # Determine skip_idle based on mode
    skip_idle = (
        mode != PROFILING_MODE_WALL if mode != PROFILING_MODE_ALL else False
    )

    output_file = None
    if args.format == "binary":
        output_file = args.outfile or _generate_output_filename(args.format, process.pid)

    # Create the appropriate collector
    collector = _create_collector(
        args.format, args.sample_interval_usec, skip_idle, args.opcodes, mode,
        output_file=output_file,
        compression=getattr(args, 'compression', 'auto'),
        diff_baseline=args.diff_baseline,
        capture_config=_capture_config_from_args(args),
    )

    with _get_child_monitor_context(args, process.pid):
        try:
            collector = sample(
                process.pid,
                collector,
                duration_sec=args.duration,
                all_threads=args.all_threads,
                realtime_stats=args.realtime_stats,
                mode=mode,
                async_aware=args.async_mode if args.async_aware else None,
                native=args.native,
                gc=args.gc,
                opcodes=args.opcodes,
                blocking=args.blocking,
            )
            _handle_output(collector, args, process.pid, mode)
        finally:
            # Terminate the main subprocess - child profilers finish when their
            # target processes exit
            if process.poll() is None:
                process.terminate()
                try:
                    process.wait(timeout=_PROCESS_KILL_TIMEOUT_SEC)
                except subprocess.TimeoutExpired:
                    process.kill()
                    process.wait()


def _handle_live_attach(args, pid):
    """Handle live mode for an existing process."""
    mode = _parse_mode(args.mode)

    # Determine skip_idle based on mode
    skip_idle = mode != PROFILING_MODE_WALL

    # Create live collector with default settings
    collector = LiveStatsCollector(
        args.sample_interval_usec,
        skip_idle=skip_idle,
        sort_by="tottime",  # Default initial sort
        limit=20,  # Default limit
        pid=pid,
        mode=mode,
        opcodes=args.opcodes,
        async_aware=args.async_mode if args.async_aware else None,
    )

    # Sample in live mode
    sample_live(
        pid,
        collector,
        duration_sec=args.duration,
        all_threads=args.all_threads,
        realtime_stats=args.realtime_stats,
        mode=mode,
        async_aware=args.async_mode if args.async_aware else None,
        native=args.native,
        gc=args.gc,
        opcodes=args.opcodes,
        blocking=args.blocking,
    )


def _handle_live_run(args):
    """Handle live mode for running a script/module."""
    # Build the command to run
    if args.module:
        cmd = (sys.executable, "-m", args.target, *args.args)
    else:
        cmd = (sys.executable, args.target, *args.args)

    # Run with synchronization, suppressing output for live mode
    try:
        process = _run_with_sync(cmd, suppress_output=True)
    except RuntimeError as e:
        sys.exit(f"Error: {e}")

    mode = _parse_mode(args.mode)

    # Determine skip_idle based on mode
    skip_idle = mode != PROFILING_MODE_WALL

    # Create live collector with default settings
    collector = LiveStatsCollector(
        args.sample_interval_usec,
        skip_idle=skip_idle,
        sort_by="tottime",  # Default initial sort
        limit=20,  # Default limit
        pid=process.pid,
        mode=mode,
        opcodes=args.opcodes,
        async_aware=args.async_mode if args.async_aware else None,
    )

    # Profile the subprocess in live mode
    try:
        sample_live(
            process.pid,
            collector,
            duration_sec=args.duration,
            all_threads=args.all_threads,
            realtime_stats=args.realtime_stats,
            mode=mode,
            async_aware=args.async_mode if args.async_aware else None,
            native=args.native,
            gc=args.gc,
            opcodes=args.opcodes,
            blocking=args.blocking,
        )
    finally:
        # Clean up the subprocess and get any error output
        returncode = process.poll()
        if returncode is None:
            # Process still running - terminate it
            process.terminate()
            try:
                process.wait(timeout=_PROCESS_KILL_TIMEOUT_SEC)
            except subprocess.TimeoutExpired:
                process.kill()
        # Ensure process is fully terminated
        process.wait()
        # Read any stderr output (tracebacks, errors, etc.)
        if process.stderr:
            with process.stderr:
                try:
                    stderr = process.stderr.read()
                    if stderr:
                        print(stderr.decode(), file=sys.stderr)
                except (OSError, ValueError):
                    # Ignore errors if pipe is already closed
                    pass


def _handle_replay(args):
    """Handle the 'replay' command - convert binary profile to another format."""
    if not os.path.exists(args.input_file):
        sys.exit(f"Error: Input file not found: {args.input_file}")

    _validate_replay_input_file(args.input_file)

    try:
        with BinaryReader(args.input_file) as reader:
            _replay_with_reader(args, reader)
    except (OSError, ValueError) as exc:
        sys.exit(f"Error: {exc}")


if __name__ == "__main__":
    main()