#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
pymem - Python memory profiler

A simple wrapper that traces Python process memory allocations
and outputs analysis results in JSON format to stdout.

Usage:
    pymem trace <pid> [options]

Features:
    - Streaming mode: zero disk usage (uses FIFO)
    - Real-time processing: attach and analyze simultaneously
    - JSON output to stdout with enriched metadata
    - Process memory timeline collection
    - Python 2/3 compatible
"""

from __future__ import print_function
import os
import sys
import subprocess
import argparse
import shutil
import tempfile
import time
import threading
import signal
import json
import re
from datetime import datetime

# Minimum GNU GDB for attach (inferior calls / xstate); see Debian #898048, upstream 51547df6 (fixed in 8.0).
MIN_GDB_VERSION = (8, 0)

# `gdb --version` first-line patterns for vendor rebuilds (GNU GDB + distro tag). Most specific first.
# See also: trailing-version and (GDB) N.M heuristics in parse_gdb_version_from_version_text.
_GDB_VERSION_VENDOR_RES = (
    re.compile(r"Red\s+Hat\s+Enterprise\s+Linux\s+(\d+)\.(\d+)\.\d+"),
    re.compile(r"Alibaba\s+Cloud\s+Linux\s+(\d+)\.(\d+)[.\-]"),
    re.compile(r"Rocky\s+Linux\s+(\d+)\.(\d+)[.\-]"),
    re.compile(r"AlmaLinux(?:\s+Linux)?\s+(\d+)\.(\d+)[.\-]"),
    re.compile(r"CentOS(?:\s+Stream)?(?:\s+Linux)?\s+(\d+)\.(\d+)[.\-]"),
    re.compile(r"Oracle\s+Linux(?:\s+Server)?\s+(\d+)\.(\d+)[.\-]"),
)

# Python 2/3 compatibility for TimeoutExpired
try:
    TimeoutExpired = subprocess.TimeoutExpired
except AttributeError:  # pragma: no cover
    # Python 2 doesn't have TimeoutExpired, define a compatible version
    class TimeoutExpired(Exception):
        """Compatibility class for Python 2 - TimeoutExpired equivalent"""

        def __init__(self, cmd=None, timeout=None, output=None, stderr=None):
            self.cmd = cmd
            self.timeout = timeout
            self.output = output
            self.stderr = stderr
            super(TimeoutExpired, self).__init__(
                "Command {0} timed out after {1} seconds".format(cmd, timeout)
            )


def wait_with_timeout(proc, timeout):
    """
    Wait for process to terminate with timeout support for Python 2/3.

    In Python 3.3+, subprocess.Popen.wait() supports timeout parameter.
    In Python 2.7, we need to manually implement timeout logic.

    Args:
        proc: subprocess.Popen instance
        timeout: Timeout in seconds (int or float)

    Returns:
        int: Process exit code

    Raises:
        TimeoutExpired: If process doesn't terminate within timeout
    """
    if sys.version_info[0] >= 3:
        # Python 3: use native timeout support
        return proc.wait(timeout=timeout)
    else:
        # Python 2: manual timeout implementation (pragma: no cover for py2)
        start_time = time.time()
        while proc.poll() is None:
            elapsed = time.time() - start_time
            if elapsed >= timeout:
                raise TimeoutExpired(cmd=getattr(proc, "args", None), timeout=timeout)
            # Check every 0.1 seconds
            time.sleep(0.1)
        return proc.returncode


def collect_process_meminfo(pid):
    """
    Collect memory information for a process from /proc/{pid}/status.

    Args:
        pid: Process ID

    Returns:
        list: Memory information array [timestamp, VmRSS, RssAnon, RssFile, RssShmem]
              - timestamp: Unix timestamp (seconds since epoch, float)
              - VmRSS: Total RSS in KB
              - RssAnon: Anonymous RSS in KB (user memory)
              - RssFile: File-backed RSS in KB (page cache)
              - RssShmem: Shared memory RSS in KB

    Returns None if process doesn't exist or data unavailable.
    """
    try:
        status_file = "/proc/{0}/status".format(pid)
        if not os.path.exists(status_file):
            return None

        mem_data = {}

        with open(status_file, "r") as f:
            for line in f:
                parts = line.strip().split()
                if len(parts) < 2:
                    continue

                key = parts[0].rstrip(":")
                value = parts[1]

                # Collect relevant memory fields (in KB)
                if key in ("VmRSS", "RssAnon", "RssFile", "RssShmem"):
                    try:
                        mem_data[key] = int(value)
                    except ValueError:
                        pass

        # Ensure we have at least VmRSS
        if "VmRSS" not in mem_data:
            return None

        # Return as array: [timestamp, VmRSS, RssAnon, RssFile, RssShmem]
        return [
            time.time(),
            mem_data.get("VmRSS", 0),
            mem_data.get("RssAnon", 0),
            mem_data.get("RssFile", 0),
            mem_data.get("RssShmem", 0),
        ]

    except (IOError, OSError):
        return None


def start_meminfo_collector(pid, duration, meminfo_list, stop_event, verbose=False):
    """
    Background thread that periodically collects process memory information.

    Collects up to 50 samples over the specified duration, with minimum 1s interval.

    Args:
        pid: Process ID to monitor
        duration: Total collection duration in seconds
        meminfo_list: List to append collected data to
        stop_event: Threading event to signal stop
        verbose: Whether to output debug info
    """
    try:
        # Calculate collection parameters
        max_samples = 50
        min_interval = 1.0  # seconds

        if duration <= 0:
            # For unlimited duration, use default interval
            interval = 1.0
        else:
            # Calculate interval: max(min_interval, duration / max_samples)
            interval = max(min_interval, float(duration) / max_samples)

        if verbose:
            sys.stderr.write(
                "[meminfo] Starting collection: pid={0}, duration={1}s, interval={2:.1f}s\n".format(
                    pid, duration if duration > 0 else "unlimited", interval
                )
            )
            sys.stderr.flush()

        sample_count = 0
        while not stop_event.is_set():
            # Collect memory info
            mem_data = collect_process_meminfo(pid)
            if mem_data:
                meminfo_list.append(mem_data)
                sample_count += 1

                if verbose and sample_count % 10 == 0:
                    sys.stderr.write(
                        "[meminfo] Collected {0} samples\n".format(sample_count)
                    )
                    sys.stderr.flush()
            else:
                # Process no longer exists or can't read data
                if verbose:
                    sys.stderr.write(
                        "[meminfo] Process {0} no longer available, stopping collection\n".format(
                            pid
                        )
                    )
                    sys.stderr.flush()
                break

            # Stop if we've reached max samples (only when duration is set)
            if duration > 0 and sample_count >= max_samples:
                break

            # Wait for next interval
            stop_event.wait(interval)

        if verbose:
            sys.stderr.write(
                "[meminfo] Collection stopped: {0} samples collected\n".format(
                    sample_count
                )
            )
            sys.stderr.flush()

    except Exception as e:
        if verbose:
            sys.stderr.write("[meminfo] Error in collector: {0}\n".format(e))
            sys.stderr.flush()


def format_response(
    code, message, tracing_data=None, meminfo_data=None, start_time=None, end_time=None
):
    """
    Format output as standardized JSON response.

    Args:
        code: Response code ('Success' or error code)
        message: Error message (empty for Success)
        tracing_data: Tracing data from pymemtrace (dict or None)
        meminfo_data: Process memory timeline data (list or None)
        start_time: Trace start timestamp (Unix timestamp, float)
        end_time: Trace end timestamp (Unix timestamp, float)

    Returns:
        dict: Formatted response
    """
    response = {
        "code": code,
        "message": message,
        "data": {
            "tracing": {
                "start_time": start_time if start_time is not None else 0,
                "end_time": end_time if end_time is not None else 0,
                "data": tracing_data if tracing_data else {},
            },
            "process_meminfo": meminfo_data if meminfo_data else [],
        },
    }
    return response


def find_pymemtrace_binary():
    """Find the pymemtrace binary in the sysak tools directory."""
    # Get script directory - compatible with Python 2 and 3
    script_dir = os.path.dirname(os.path.abspath(__file__))

    # Possible locations (in order of preference)
    possible_paths = [
        # Same directory as script (most common case)
        os.path.join(script_dir, "pymemtrace"),
    ]

    # Add kernel-specific path if available
    try:
        import platform

        kernel_release = platform.release()
        possible_paths.append(os.path.join(script_dir, kernel_release, "pymemtrace"))
    except (AttributeError, OSError):
        pass

    # Add other possible paths
    possible_paths.extend(
        [
            # In current directory (for development)
            os.path.join(script_dir, "pymem-trace/release/pymemtrace"),
            # In system path
            shutil.which("pymemtrace") if hasattr(shutil, "which") else None,
        ]
    )

    for path in possible_paths:
        if path and os.path.isfile(path) and os.access(path, os.X_OK):
            return path

    # Try to find in installed sysak location
    sysak_tools_path = "/usr/local/sysak/.sysak_components/tools"
    if os.path.isdir(sysak_tools_path):
        # First try same directory as script
        try:
            for kernel_version in os.listdir(sysak_tools_path):
                kernel_path = os.path.join(
                    sysak_tools_path, kernel_version, "pymemtrace"
                )
                if os.path.isfile(kernel_path):
                    return kernel_path
        except OSError:
            pass
        # Also try root tools directory
        root_tools_path = os.path.join(sysak_tools_path, "pymemtrace")
        if os.path.isfile(root_tools_path):
            return root_tools_path

    return None


def find_helper_file(filename):
    """Find helper files (_attach.gdb, libpymemtrace_inject.so) in the same directory as pymemtrace."""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(script_dir, filename)


def check_process_exists(pid):
    """
    Check if a process exists.

    Args:
        pid: Process ID

    Returns:
        tuple: (success: bool, error_message: str)
    """
    proc_path = "/proc/{0}".format(pid)
    if not os.path.exists(proc_path):
        return False, "Process {0} does not exist".format(pid)
    return True, ""


def get_process_executable(pid):
    """
    Get the executable path of a process.

    Args:
        pid: Process ID

    Returns:
        str: Executable path, or None if unavailable
    """
    try:
        exe_path = "/proc/{0}/exe".format(pid)
        return os.readlink(exe_path)
    except (OSError, IOError):
        return None


def check_python_process(pid):
    """
    Check if a process is a Python process.

    Args:
        pid: Process ID

    Returns:
        tuple: (success: bool, error_message: str, executable: str)
    """
    exe = get_process_executable(pid)
    if not exe:
        return (
            False,
            "Cannot determine process executable (permission denied or process exited)",
            None,
        )

    # Check if it's a Python executable
    exe_name = os.path.basename(exe)
    exe_lower = exe_name.lower()

    # Support various Python executable names:
    # - python, python2, python3, python3.6, python3.11, etc.
    # - platform-python, platform-python3.6 (RHEL 8/AlmaLinux/CentOS 8 system Python)
    # - pypy, pypy3 (will be caught by later CPython check)
    is_python = (
        exe_name.startswith("python")
        or exe_name.startswith("platform-python")
        or "python" in exe_lower
        or exe_lower.startswith(
            "pypy"
        )  # PyPy support (will fail later in check_cpython)
    )

    if not is_python:
        return (
            False,
            "Process {0} is not a Python process (executable: {1})".format(
                pid, exe_name
            ),
            exe,
        )

    return True, "", exe


def get_python_version(pid):
    """
    Get Python version of a running process by reading /proc/{pid}/environ.

    Args:
        pid: Process ID

    Returns:
        tuple: (version_tuple, version_string) or (None, None) if unavailable
               version_tuple: (major, minor, micro) as integers
               version_string: e.g., "3.8.10"
    """
    try:
        # Try to get version from process cmdline
        cmdline_path = "/proc/{0}/cmdline".format(pid)
        with open(cmdline_path, "r") as f:
            cmdline = f.read()

        # Parse cmdline for python executable
        parts = cmdline.split("\x00")
        if parts and "python" in parts[0]:
            exe = parts[0]

            # Try to execute python --version
            try:
                import subprocess

                # Python 2/3 compatible version check without timeout parameter
                proc = subprocess.Popen(
                    [exe, "--version"], stdout=subprocess.PIPE, stderr=subprocess.STDOUT
                )
                try:
                    wait_with_timeout(proc, 2)
                    result = proc.stdout.read()
                except TimeoutExpired:
                    proc.kill()
                    proc.wait()
                    raise

                version_str = result.decode("utf-8", errors="ignore").strip()

                # Parse version string like "Python 3.8.10"
                if "Python" in version_str:
                    version_part = version_str.split()[1]
                    version_nums = version_part.split(".")
                    if len(version_nums) >= 2:
                        major = int(version_nums[0])
                        minor = int(version_nums[1])
                        micro = int(version_nums[2]) if len(version_nums) > 2 else 0
                        return (major, minor, micro), version_part
            except:
                pass
    except:
        pass

    return None, None


def check_python_version(pid):
    """
    Check if Python version is >= 3.6.

    Args:
        pid: Process ID

    Returns:
        tuple: (success: bool, error_message: str, version_string: str)
    """
    version_tuple, version_str = get_python_version(pid)

    if not version_tuple:
        # Cannot determine version, allow to proceed (will be caught by pymemtrace)
        return True, "", "unknown"

    major, minor, micro = version_tuple

    if major < 3 or (major == 3 and minor < 6):
        return (
            False,
            "Python version {0} is too old (requires >= 3.6)".format(version_str),
            version_str,
        )

    return True, "", version_str


def check_cpython(pid):
    """
    Check if the process is running CPython (not PyPy, Jython, etc.).

    Args:
        pid: Process ID

    Returns:
        tuple: (success: bool, error_message: str, implementation: str)
    """
    exe = get_process_executable(pid)
    if not exe:
        return True, "", "unknown"  # Cannot determine, allow to proceed

    exe_lower = exe.lower()

    # Check for known non-CPython implementations
    if "pypy" in exe_lower:
        return False, "PyPy is not supported (only CPython)", "PyPy"
    elif "jython" in exe_lower:
        return False, "Jython is not supported (only CPython)", "Jython"
    elif "ironpython" in exe_lower:
        return False, "IronPython is not supported (only CPython)", "IronPython"

    # Assume CPython if it's just "python*"
    return True, "", "CPython"


def get_process_start_time(pid):
    """
    Get process start time to uniquely identify a process.

    This reads the start time from /proc/{pid}/stat (field 22), which is
    the time the process started after system boot (in clock ticks).
    This value uniquely identifies a process and doesn't change even if
    the PID is reused by a different process.

    Args:
        pid: Process ID

    Returns:
        str: Process start time (clock ticks since boot), or None if failed
    """
    try:
        with open("/proc/{0}/stat".format(pid), "r") as f:
            stat = f.read()
            # Field 22 (index 21) is starttime
            # Note: process name can contain spaces and parens, so we need to
            # find the last ')' and split from there
            paren_end = stat.rfind(")")
            if paren_end != -1:
                fields = stat[paren_end + 2 :].split()
                if len(fields) >= 20:
                    return fields[19]  # starttime is 20th field after ')'
    except (FileNotFoundError, ProcessLookupError, IOError, PermissionError):
        pass
    return None


def is_process_alive(pid, expected_start_time=None):
    """
    Check if a process is still alive, optionally verifying it's the same process.

    This function handles several edge cases:
    - Process death: returns False
    - PID reuse: if expected_start_time provided, returns False for different process
    - Zombie processes: returns True (PID still exists, will be caught by timeout)
    - Permission errors: returns True (conservative, assume alive if can't verify)

    Args:
        pid: Process ID
        expected_start_time: Optional. Expected process start time from get_process_start_time().
                           If provided, verifies this is the same process (not reused PID).

    Returns:
        bool: True if process exists (and is same if start_time provided), False otherwise
    """
    # If we have expected start time, use it to verify same process
    if expected_start_time is not None:
        current_start_time = get_process_start_time(pid)
        if current_start_time is None:
            # Can't read process info - likely dead
            return False
        # Compare start times - protects against PID reuse
        return current_start_time == expected_start_time

    # Fallback to simple check using kill(0)
    try:
        os.kill(pid, 0)
        return True
    except PermissionError:
        # No permission to signal, but process exists
        # This is conservative - if we can't verify, assume alive
        return True
    except (OSError, ProcessLookupError):
        # Process not found
        return False


def pre_check_target_process(pid, verbose=False):
    """
    Perform pre-flight checks on the target process.

    Args:
        pid: Process ID
        verbose: Whether to print verbose messages

    Returns:
        tuple: (success: bool, error_code: str, error_message: str)
    """
    if verbose:
        sys.stderr.write(
            "[pymem] Performing pre-flight checks on process {0}...\n".format(pid)
        )
        sys.stderr.flush()

    # 1. Check if process exists
    success, msg = check_process_exists(pid)
    if not success:
        return False, "ProcessNotFound", msg

    if verbose:
        sys.stderr.write("[pymem] ✓ Process exists\n")
        sys.stderr.flush()

    # 2. Check if it's a Python process
    success, msg, exe = check_python_process(pid)
    if not success:
        return False, "NotPythonProcess", msg

    if verbose:
        sys.stderr.write("[pymem] ✓ Python process: {0}\n".format(exe))
        sys.stderr.flush()

    # 3. Check Python version
    success, msg, version = check_python_version(pid)
    if not success:
        return False, "PythonVersionTooOld", msg

    if verbose and version != "unknown":
        sys.stderr.write("[pymem] ✓ Python version: {0}\n".format(version))
        sys.stderr.flush()

    # 4. Check if it's CPython
    success, msg, impl = check_cpython(pid)
    if not success:
        return False, "NotCPython", msg

    if verbose and impl != "unknown":
        sys.stderr.write("[pymem] ✓ Implementation: {0}\n".format(impl))
        sys.stderr.flush()

    if verbose:
        sys.stderr.write("[pymem] ✓ All pre-flight checks passed\n")
        sys.stderr.flush()

    return True, "", ""


def stream_stderr(proc, name, verbose=False):
    """Stream stderr from a subprocess (only in verbose mode)."""
    if not verbose:
        # Silently consume stderr
        try:
            for line in iter(proc.stderr.readline, b""):
                pass
        except:
            pass
        return

    # Verbose mode: print stderr
    try:
        for line in iter(proc.stderr.readline, b""):
            if line:
                sys.stderr.write(
                    "[{0}] {1}".format(name, line.decode("utf-8", errors="replace"))
                )
                sys.stderr.flush()
    except:
        pass


def read_json_output(proc):
    """Read JSON output from processor stdout."""
    output = []
    try:
        for line in iter(proc.stdout.readline, b""):
            if line:
                output.append(line)
    except:
        pass
    return b"".join(output)


def trace_command(
    pymemtrace_bin,
    pid,
    duration=0,
    gotplt_hook=True,
    pymem_hook=False,
    pymem_sample_rate=0,
    native=True,
    verbose=False,
):
    """
    Trace a Python process memory allocations and output enriched JSON to stdout.

    Workflow:
    1. Start process memory collector (background thread)
    2. Create a FIFO for streaming data
    3. Start processor (reads from FIFO, outputs JSON to stdout)
    4. Start attach (writes trace data to FIFO)
    5. Wait for both to complete
    6. Stop memory collector
    7. Format and output enriched JSON
    8. Clean up FIFO

    Args:
        pymemtrace_bin: Path to pymemtrace binary
        pid: Target process PID
        duration: Trace duration in seconds (0 = unlimited)
        gotplt_hook: Enable GOT/PLT hook (malloc/calloc/etc)
        pymem_hook: Enable PyMem API hook (PyObject_Malloc/etc)
        pymem_sample_rate: When pymem_hook is on, record 1 in N allocations (0 or 1 = all; >1 reduces overhead)
        native: Enable native stack traces
        verbose: Verbose output to stderr

    Returns:
        int: Always returns 0; check JSON output's 'code' field for actual status
    """
    fifo_path = None
    processor_proc = None
    attach_proc = None
    meminfo_collector_thread = None
    meminfo_stop_event = threading.Event()
    meminfo_list = []
    start_time = time.time()
    end_time = None

    try:
        # Pre-flight checks
        check_ok, error_code, error_msg = pre_check_target_process(pid, verbose)
        if not check_ok:
            # Pre-check failed, return error immediately
            response = format_response(
                code=error_code,
                message=error_msg,
                start_time=start_time,
                end_time=time.time(),
            )
            output_json = json.dumps(response, indent=2, ensure_ascii=False)
            if hasattr(sys.stdout, "buffer"):
                sys.stdout.buffer.write(output_json.encode("utf-8"))
            else:
                sys.stdout.write(output_json)
            sys.stdout.write("\n")
            sys.stdout.flush()
            return 0

        # Get process start time to protect against PID reuse
        target_start_time = get_process_start_time(pid)
        if verbose and target_start_time:
            sys.stderr.write(
                "[pymem] Target process start time: {0}\n".format(target_start_time)
            )
            sys.stderr.flush()  # Always return 0, error is in JSON

        # Start memory info collector
        meminfo_collector_thread = threading.Thread(
            target=start_meminfo_collector,
            args=(pid, duration, meminfo_list, meminfo_stop_event, verbose),
        )
        meminfo_collector_thread.daemon = True
        meminfo_collector_thread.start()

        # Create FIFO
        fifo_path = "/tmp/pymem-trace-{0}-{1}.fifo".format(pid, os.getpid())

        # Remove if exists
        if os.path.exists(fifo_path):
            try:
                os.remove(fifo_path)
            except OSError:
                pass

        # Create FIFO
        os.mkfifo(fifo_path, 0o600)
        if verbose:
            sys.stderr.write("[pymem] Created FIFO: {0}\n".format(fifo_path))
            sys.stderr.flush()

        # Start processor: outputs JSON to stdout
        processor_cmd = [
            pymemtrace_bin,
            "report",
            fifo_path,
            "/dev/stdout",  # Output to stdout
            "--streaming",
            "--format",
            "json",
        ]

        if verbose:
            sys.stderr.write(
                "[pymem] Starting processor: {0}\n".format(" ".join(processor_cmd))
            )
            sys.stderr.flush()

        processor_proc = subprocess.Popen(
            processor_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )

        # Start stderr streaming thread for processor
        processor_stderr_thread = threading.Thread(
            target=stream_stderr, args=(processor_proc, "processor", verbose)
        )
        processor_stderr_thread.daemon = True
        processor_stderr_thread.start()

        # Start stdout reading thread for processor
        json_output = []

        def read_processor_output():
            try:
                for line in iter(processor_proc.stdout.readline, b""):
                    if line:
                        json_output.append(line)
            except:
                pass

        processor_stdout_thread = threading.Thread(target=read_processor_output)
        processor_stdout_thread.daemon = True
        processor_stdout_thread.start()

        # Give processor time to open FIFO
        time.sleep(0.5)

        # Start attach
        attach_cmd = [pymemtrace_bin, "attach", str(pid), fifo_path, "--use-fifo"]

        if duration > 0:
            attach_cmd.extend(["--duration", str(duration)])
        if not native:
            attach_cmd.append("--no-native")
        if not gotplt_hook:
            attach_cmd.append("--no-gotplt-hook")
        if pymem_hook:
            attach_cmd.append("--pymem-hook")
        if pymem_sample_rate > 1:
            attach_cmd.extend(["--pymem-sample-rate", str(pymem_sample_rate)])
        # Note: don't pass --verbose to pymemtrace attach, we handle output ourselves

        if verbose:
            sys.stderr.write(
                "[pymem] Starting attach: {0}\n".format(" ".join(attach_cmd))
            )
            sys.stderr.flush()

        attach_proc = subprocess.Popen(
            attach_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )

        # Start stderr streaming thread for attach
        attach_stderr_thread = threading.Thread(
            target=stream_stderr, args=(attach_proc, "attach", verbose)
        )
        attach_stderr_thread.daemon = True
        attach_stderr_thread.start()

        # Wait for attach to complete with target process monitoring
        # Poll both attach process and target process
        if verbose:
            sys.stderr.write("[pymem] Waiting for attach to complete...\n")
            sys.stderr.flush()

        attach_exit = None
        poll_interval = 0.5  # Check every 0.5 seconds
        max_wait_time = (
            (duration + 30) if duration > 0 else 300
        )  # duration + 30s buffer, or 5 min default
        # Cap wait for attach phase so we fail fast on attach hang/failure (do not wait duration+30)
        attach_phase_timeout = 25
        waited_time = 0

        while attach_exit is None and waited_time < max_wait_time and waited_time < attach_phase_timeout:
            attach_exit = attach_proc.poll()

            if attach_exit is None:
                # Check if target process is still alive and is the same process
                if not is_process_alive(pid, target_start_time):
                    if verbose:
                        sys.stderr.write(
                            "[pymem] WARNING: Target process {0} is no longer alive or changed\n".format(
                                pid
                            )
                        )
                        sys.stderr.flush()
                        sys.stderr.flush()

                    # Target crashed, terminate both attach and processor
                    if verbose:
                        sys.stderr.write(
                            "[pymem] Terminating attach and processor due to target crash\n"
                        )
                        sys.stderr.flush()

                    attach_proc.terminate()
                    processor_proc.terminate()  # Also terminate processor

                    try:
                        wait_with_timeout(attach_proc, 5)
                    except TimeoutExpired:
                        attach_proc.kill()
                        attach_proc.wait()

                    attach_exit = -1  # Mark as error
                    break

                time.sleep(poll_interval)
                waited_time += poll_interval

        # If still not done, force termination (attach phase timeout or full max_wait_time)
        if attach_exit is None:
            if verbose:
                sys.stderr.write(
                    "[pymem] WARNING: Attach phase timeout after {0}s, terminating\n".format(
                        min(attach_phase_timeout, max_wait_time)
                    )
                )
                sys.stderr.flush()
            attach_proc.terminate()
            try:
                attach_exit = wait_with_timeout(attach_proc, 5)
            except TimeoutExpired:
                attach_proc.kill()
                attach_exit = attach_proc.wait()

        if verbose:
            sys.stderr.write(
                "[pymem] Attach completed with exit code: {0}\n".format(attach_exit)
            )
            sys.stderr.flush()

        # If attach failed, stop meminfo and exit immediately (do not wait for processor)
        if attach_exit != 0:
            meminfo_stop_event.set()
            if meminfo_collector_thread:
                meminfo_collector_thread.join(timeout=2)
            if processor_proc and processor_proc.poll() is None:
                processor_proc.terminate()
                try:
                    wait_with_timeout(processor_proc, 3)
                except TimeoutExpired:
                    processor_proc.kill()
            if attach_exit == -1:
                # Set in attach wait loop when target PID is gone or reused before attach finished
                error_code, error_msg = (
                    "TargetProcessCrashed",
                    "Target process {0} crashed or exited during tracing".format(pid),
                )
            elif attach_exit == 1:
                error_code, error_msg = "AttachFailed", "Failed to attach to process (exit code: {0})".format(attach_exit)
            elif attach_exit == 2:
                error_code, error_msg = "PermissionDenied", "Permission denied when attaching to process"
            else:
                error_code, error_msg = "AttachError", "Attach failed with exit code {0}".format(attach_exit)
            end_time = time.time()
            response = format_response(
                code=error_code,
                message=error_msg,
                meminfo_data=meminfo_list,
                start_time=start_time,
                end_time=end_time,
            )
            output_json = json.dumps(response, indent=2, ensure_ascii=False)
            if hasattr(sys.stdout, "buffer"):
                sys.stdout.buffer.write(output_json.encode("utf-8"))
            else:
                sys.stdout.write(output_json)
            sys.stdout.write("\n")
            sys.stdout.flush()
            return 0

        # Wait for processor to complete with timeout
        if verbose:
            sys.stderr.write("[pymem] Waiting for processor to complete...\n")
            sys.stderr.flush()

        processor_exit = None
        waited_time = 0

        while processor_exit is None and waited_time < max_wait_time:
            processor_exit = processor_proc.poll()

            if processor_exit is None:
                time.sleep(poll_interval)
                waited_time += poll_interval

        # If still not done, force termination
        if processor_exit is None:
            if verbose:
                sys.stderr.write(
                    "[pymem] WARNING: Processor timeout after {0}s, terminating\n".format(
                        max_wait_time
                    )
                )
                sys.stderr.flush()
            processor_proc.terminate()
            try:
                processor_exit = wait_with_timeout(processor_proc, 5)
            except TimeoutExpired:
                processor_proc.kill()
                processor_exit = processor_proc.wait()

        if verbose:
            sys.stderr.write(
                "[pymem] Processor completed with exit code: {0}\n".format(
                    processor_exit
                )
            )
            sys.stderr.flush()

        # Stop memory info collector
        end_time = time.time()
        meminfo_stop_event.set()
        if meminfo_collector_thread:
            meminfo_collector_thread.join(timeout=2)

        # Wait for stdout thread to finish reading
        processor_stdout_thread.join(timeout=5)

        # Parse tracing JSON output
        tracing_data = {}
        if json_output:
            try:
                output_str = b"".join(json_output).decode("utf-8", errors="replace")
                tracing_data = json.loads(output_str)
            except (ValueError, json.JSONDecodeError) as e:
                if verbose:
                    sys.stderr.write(
                        "[pymem] Warning: failed to parse tracing JSON: {0}\n".format(e)
                    )
                    sys.stderr.flush()
                # Keep as empty dict if parsing fails
                tracing_data = {}

        # Determine error code and message based on exit codes
        if attach_exit == 0 and processor_exit == 0:
            error_code = "Success"
            error_msg = ""
        elif processor_exit != 0:
            # attach_exit is always 0 here (non-zero attach returns earlier). Prefer
            # ProcessorError over TargetProcessCrashed when the target PID is gone: the
            # workload often exits normally while trace duration is still running.
            error_code = "ProcessorError"
            error_msg = "Processor failed with exit code {0}".format(processor_exit)
            if not is_process_alive(pid, target_start_time):
                error_msg += (
                    " (target process {0} is no longer present; may have exited, crashed, "
                    "or PID was reused)"
                ).format(pid)
        else:
            error_code = "TraceError"
            error_msg = "Trace failed: attach_exit={0}, processor_exit={1}".format(
                attach_exit, processor_exit
            )

        # Format final response
        response = format_response(
            code=error_code,
            message=error_msg,
            tracing_data=tracing_data,
            meminfo_data=meminfo_list,
            start_time=start_time,
            end_time=end_time,
        )

        # Output enriched JSON to stdout
        output_json = json.dumps(response, indent=2, ensure_ascii=False)
        if hasattr(sys.stdout, "buffer"):
            sys.stdout.buffer.write(output_json.encode("utf-8"))
        else:
            sys.stdout.write(output_json)
        sys.stdout.write("\n")
        sys.stdout.flush()

        # Always return 0; actual status is in JSON 'code' field
        return 0

    except KeyboardInterrupt:
        if verbose:
            sys.stderr.write("\n[pymem] Interrupted by user, cleaning up...\n")
            sys.stderr.flush()

        # Stop memory collector
        end_time = time.time()
        meminfo_stop_event.set()
        if meminfo_collector_thread:
            meminfo_collector_thread.join(timeout=1)

        # Terminate processes
        if attach_proc and attach_proc.poll() is None:
            attach_proc.terminate()
            try:
                wait_with_timeout(attach_proc, 5)
            except:
                attach_proc.kill()

        if processor_proc and processor_proc.poll() is None:
            processor_proc.terminate()
            try:
                wait_with_timeout(processor_proc, 5)
            except:
                processor_proc.kill()

        # Output error response
        response = format_response(
            code="Interrupted",
            message="Trace interrupted by user",
            meminfo_data=meminfo_list,
            start_time=start_time,
            end_time=end_time,
        )
        output_json = json.dumps(response, indent=2, ensure_ascii=False)
        if hasattr(sys.stdout, "buffer"):
            sys.stdout.buffer.write(output_json.encode("utf-8"))
        else:
            sys.stdout.write(output_json)
        sys.stdout.write("\n")
        sys.stdout.flush()

        return 0  # Always return 0, interruption is in JSON

    except Exception as e:
        if verbose:
            sys.stderr.write("[pymem] Error: {0}\n".format(e))
            import traceback

            traceback.print_exc()
            sys.stderr.flush()

        # Stop memory collector
        end_time = time.time()
        meminfo_stop_event.set()
        if meminfo_collector_thread:
            meminfo_collector_thread.join(timeout=1)

        # Output error response
        error_msg = str(e)
        response = format_response(
            code="SystemError",
            message=error_msg,
            meminfo_data=meminfo_list,
            start_time=start_time,
            end_time=end_time,
        )
        output_json = json.dumps(response, indent=2, ensure_ascii=False)
        if hasattr(sys.stdout, "buffer"):
            sys.stdout.buffer.write(output_json.encode("utf-8"))
        else:
            sys.stdout.write(output_json)
        sys.stdout.write("\n")
        sys.stdout.flush()

        return 0  # Always return 0, error is in JSON

    finally:
        # Ensure memory collector is stopped
        meminfo_stop_event.set()
        if meminfo_collector_thread and meminfo_collector_thread.is_alive():
            meminfo_collector_thread.join(timeout=1)

        # Clean up FIFO
        if fifo_path and os.path.exists(fifo_path):
            try:
                os.remove(fifo_path)
                if verbose:
                    sys.stderr.write("[pymem] Cleaned up FIFO: {0}\n".format(fifo_path))
                    sys.stderr.flush()
            except OSError as e:
                if verbose:
                    sys.stderr.write(
                        "[pymem] Warning: failed to remove FIFO: {0}\n".format(e)
                    )
                    sys.stderr.flush()


def check_system_gdb():
    """Check that system GDB is available (required for attach). Return (True, path) or (False, None)."""
    if hasattr(shutil, "which"):
        path = shutil.which("gdb")
        if path:
            return True, path
    # Fallback: try running gdb --version
    try:
        with open(os.devnull, "w") as devnull:
            ret = subprocess.call(
                ["gdb", "--version"],
                stdout=devnull,
                stderr=devnull,
            )
        if ret == 0:
            return True, "gdb"
    except (OSError, IOError):
        pass
    return False, None


def _gdb_version_tuple_ok(major, minor):
    """True if (major, minor) looks like a real GDB release (not a year or kernel-style number)."""
    return 4 <= major <= 40 and 0 <= minor <= 99


def parse_gdb_version_from_version_text(version_text):
    """
    Parse (major, minor) from `gdb --version` output (first line is enough).

    Compatible with common ``gdb --version`` shapes (Python 2 and 3):

    * Upstream / minimal: ``GNU gdb (GDB) 12.1``, ``GNU gdb (GDB) 8.3.1``
    * Debian/Ubuntu: ``GNU gdb (Debian 12.1-3) 12.1`` — trailing ``N.M`` is the GDB version
    * Fedora: ``GNU gdb (GDB) Fedora 32 9.2`` — trailing ``9.2``
    * RHEL rebuild tag: ``GNU gdb (GDB) Red Hat Enterprise Linux 7.6.1-64.el7`` — embedded ``7.6``
    * Alibaba / similar: ``GNU gdb (GDB) Alibaba Cloud Linux 9.2-7.2.0.5.al8`` — gdb RPM ``9.2``,
      not the ``7.2`` inside the release field (handled by vendor regex or max fallback)
    * Rocky/Alma/CentOS/Oracle: vendor prefix + ``X.Y-`` style
    * Fallback: scan all ``N.M`` tokens in a plausible GDB range and take ``max`` when vendor
      strings embed multiple versions (e.g. ``9.2-7.2.0.5``).

    Returns None if no plausible version is found.
    """
    if not version_text:
        return None
    # Python 3: subprocess may return bytes; Python 2: str
    if sys.version_info[0] >= 3 and isinstance(version_text, bytes):
        version_text = version_text.decode("utf-8", "replace")
    line = version_text.strip().split("\n")[0]

    for rx in _GDB_VERSION_VENDOR_RES:
        m = rx.search(line)
        if m:
            major, minor = int(m.group(1)), int(m.group(2))
            if _gdb_version_tuple_ok(major, minor):
                return major, minor

    # ``GNU gdb (GDB) 10.2`` — version immediately after ``(GDB)``
    m = re.search(r"\(GDB\)\s*(\d+)\.(\d+)(?:\.\d+)?\b", line)
    if m:
        major, minor = int(m.group(1)), int(m.group(2))
        if _gdb_version_tuple_ok(major, minor):
            return major, minor

    # Line ends with canonical gdb version: ``... ) 12.1``, ``... Fedora 32 9.2``
    m = re.search(r"(\d+)\.(\d+)(?:\.\d+)?\s*$", line)
    if m:
        major, minor = int(m.group(1)), int(m.group(2))
        if _gdb_version_tuple_ok(major, minor):
            return major, minor

    # Last resort: every ``N.M`` token; vendor release strings may list several (e.g. ``9.2-7.2``).
    candidates = []
    for m in re.finditer(r"\b(\d+)\.(\d+)(?:\.\d+)?\b", line):
        major, minor = int(m.group(1)), int(m.group(2))
        if _gdb_version_tuple_ok(major, minor):
            candidates.append((major, minor))
    if candidates:
        return max(candidates)
    return None


def resolve_gdb_executable_for_version_check(gdb_ok, gdb_path, bundled_gdb):
    """Which gdb binary pymemtrace will effectively use from PATH, or bundled if no system gdb."""
    if gdb_ok:
        if gdb_path and gdb_path != "gdb" and os.path.isfile(gdb_path):
            return gdb_path
        if hasattr(shutil, "which"):
            w = shutil.which("gdb")
            if w:
                return w
        return "gdb"
    if os.path.isfile(bundled_gdb) and os.access(bundled_gdb, os.X_OK):
        return bundled_gdb
    return None


def check_gdb_meets_minimum(gdb_executable, minimum=MIN_GDB_VERSION, timeout=8):
    """
    Run gdb --version and ensure version >= minimum (major, minor).

    Returns:
        (True, None, version_str) if OK
        (False, error_message, None or version_str) if too old or unparsable
    """
    if not gdb_executable:
        return False, "GDB executable path is empty.", None
    try:
        # timeout= exists only on Python 3.3+; Python 2 has no timeout on check_output.
        if sys.version_info[:2] >= (3, 3):
            out = subprocess.check_output(
                [gdb_executable, "--version"],
                stderr=subprocess.STDOUT,
                timeout=timeout,
            )
        else:
            out = subprocess.check_output(
                [gdb_executable, "--version"],
                stderr=subprocess.STDOUT,
            )
        if sys.version_info[0] >= 3 and isinstance(out, bytes):
            out = out.decode("utf-8", "replace")
    except TimeoutExpired as e:
        return False, "Timed out running `{0} --version`: {1}".format(gdb_executable, e), None
    except (OSError, subprocess.CalledProcessError) as e:
        return False, "Could not run `{0} --version`: {1}".format(gdb_executable, e), None

    ver = parse_gdb_version_from_version_text(out)
    if ver is None:
        return (
            False,
            (
                "Could not parse GDB version from `gdb --version`. "
                "Require GNU GDB {0}.{1} or later for attach."
            ).format(minimum[0], minimum[1]),
            None,
        )
    ver_str = "{0}.{1}".format(ver[0], ver[1])
    if ver < minimum:
        return (
            False,
            (
                "GDB version is {0}; pymem trace requires GNU GDB {1}.{2} or later. "
                "Older GDB often fails attach with \"Couldn't write extended state status: Bad address\" "
                "when calling dlopen from the target (fixed in GDB 8.0+). "
                "Install a newer GDB and ensure it is first in PATH (e.g. /usr/local/bin/gdb)."
            ).format(ver_str, minimum[0], minimum[1]),
            ver_str,
        )
    return True, None, ver_str


def main():
    """Main entry point."""
    # Find pymemtrace first
    pymemtrace_bin = find_pymemtrace_binary()

    if not pymemtrace_bin:
        sys.stderr.write("Error: pymemtrace binary not found.\n")
        sys.stderr.write("Please ensure pymem-trace is built and available.\n")
        sys.exit(0)  # Always return 0

    # Parse arguments
    parser = argparse.ArgumentParser(
        prog="pymem",
        description="Python memory profiler - trace and analyze memory allocations",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Trace for 30 seconds and output JSON
  %(prog)s trace 12345 --duration 30

  # Trace with both hook modes enabled
  %(prog)s trace 12345 --duration 60 --gotplt-hook --pymem-hook

  # Trace without native stacks (faster)
  %(prog)s trace 12345 --duration 10 --no-native

  # Verbose mode
  %(prog)s trace 12345 --duration 30 -v

Output:
  JSON data is written to stdout
  Progress messages are written to stderr
        """,
    )

    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # trace command
    trace_parser = subparsers.add_parser(
        "trace",
        help="Trace a Python process and output JSON analysis",
        description="Attach to a running Python process, trace memory allocations, and output analysis in JSON format to stdout.",
    )
    trace_parser.add_argument("pid", type=int, help="Target process PID")
    trace_parser.add_argument(
        "--duration",
        type=int,
        default=30,
        help="Trace duration in seconds (default: 30, 0 = unlimited)",
    )
    trace_parser.add_argument(
        "--gotplt-hook",
        action="store_true",
        default=True,
        help="Enable GOT/PLT hook (malloc/calloc/etc) [default: enabled]",
    )
    trace_parser.add_argument(
        "--no-gotplt-hook", action="store_true", help="Disable GOT/PLT hook"
    )
    trace_parser.add_argument(
        "--pymem-hook",
        action="store_true",
        help="Enable PyMem API hook (PyObject_Malloc/etc) [default: disabled]. Can add significant overhead for allocation-heavy workloads.",
    )
    trace_parser.add_argument(
        "--pymem-sample-rate",
        type=int,
        default=0,
        metavar="N",
        help="When --pymem-hook is on, record only 1 in N allocations (default: 1 = all). Use N>1 (e.g. 10 or 100) to reduce target process impact.",
    )
    trace_parser.add_argument(
        "--native",
        action="store_true",
        default=True,
        help="Enable native stack traces [default: enabled]",
    )
    trace_parser.add_argument(
        "--no-native", action="store_true", help="Disable native stack traces"
    )
    trace_parser.add_argument(
        "--verbose", "-v", action="store_true", help="Verbose output to stderr"
    )

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(0)  # Always return 0

    if args.command == "trace":
        # Require GDB for attach: either system GDB or bundled gdb next to pymemtrace
        tool_dir = os.path.dirname(os.path.abspath(pymemtrace_bin))
        bundled_gdb = os.path.join(tool_dir, "gdb")
        gdb_ok, gdb_path = check_system_gdb()
        if not gdb_ok and not (os.path.isfile(bundled_gdb) and os.access(bundled_gdb, os.X_OK)):
            err_msg = (
                "GDB not found. pymem trace requires GDB for attach. "
                "Please install GDB (e.g. yum install gdb) or use a build that includes bundled gdb in the tools directory."
            )
            response = format_response("GDBNotFound", err_msg)
            output_json = json.dumps(response, ensure_ascii=False)
            if hasattr(sys.stdout, "buffer"):
                sys.stdout.buffer.write(output_json.encode("utf-8"))
            else:
                sys.stdout.write(output_json)
            sys.stdout.write("\n")
            sys.stdout.flush()
            sys.exit(0)

        gdb_exe = resolve_gdb_executable_for_version_check(
            gdb_ok, gdb_path, bundled_gdb
        )
        ver_ok, ver_err, ver_str = check_gdb_meets_minimum(gdb_exe)
        if not ver_ok:
            response = format_response("GDBVersionTooOld", ver_err)
            output_json = json.dumps(response, ensure_ascii=False)
            if hasattr(sys.stdout, "buffer"):
                sys.stdout.buffer.write(output_json.encode("utf-8"))
            else:
                sys.stdout.write(output_json)
            sys.stdout.write("\n")
            sys.stdout.flush()
            sys.exit(0)
        if args.verbose and ver_str:
            sys.stderr.write(
                "[pymem] GDB version {0} (minimum {1}.{2})\n".format(
                    ver_str, MIN_GDB_VERSION[0], MIN_GDB_VERSION[1]
                )
            )
            sys.stderr.flush()

        # Handle flag logic
        gotplt_hook = args.gotplt_hook and not args.no_gotplt_hook
        native = args.native and not args.no_native

        # Execute trace command
        try:
            pymem_sample_rate = args.pymem_sample_rate if args.pymem_sample_rate >= 1 else 1
            result = trace_command(
                pymemtrace_bin,
                args.pid,
                duration=args.duration,
                gotplt_hook=gotplt_hook,
                pymem_hook=args.pymem_hook,
                pymem_sample_rate=pymem_sample_rate,
                native=native,
                verbose=args.verbose,
            )
            sys.exit(result)  # result is always 0 now
        except KeyboardInterrupt:
            if args.verbose:
                sys.stderr.write("\n[pymem] Interrupted by user\n")
            sys.exit(0)  # Always return 0, interruption is in JSON
        except Exception as e:
            if args.verbose:
                sys.stderr.write("[pymem] Error: {0}\n".format(e))
                import traceback

                traceback.print_exc()
            sys.exit(0)  # Always return 0, error is in JSON
    else:
        parser.print_help()
        sys.exit(0)  # Always return 0


if __name__ == "__main__":
    main()
