#!/usr/bin/env python3
import psutil
import time
import csv
import os
import argparse
from datetime import datetime

"""
CARMA Platform CPU Monitor Script
Requirements:
    - Python 3.6 or higher
    - psutil package (install with: pip3 install psutil)
Usage:
    1. Open a new terminal before starting CARMA Platform
    2. Navigate to the directory containing this script:
       cd /path/to/script/
    3. Run this script:
       python3 monitor-ros-cpu.py
    4. In a different terminal, start CARMA Platform:
       carma start all
    5. The script will automatically monitor and log CPU/memory usage
       of all ROS2 nodes and related processes during CARMA operation
    6. To stop monitoring:
       - Press Ctrl+C ONCE in the monitoring terminal
       - The CSV output file will be saved in the logs directory following ROS bag naming convention:
         'logs/cpu_usage_ros2_nodes_YYYY_MM_DD-HH_MM_SS.csv'
Output:
    - CSV file containing timestamp, process info, CPU and memory usage
    - Data can be used to analyze CARMA Platform resource utilization
"""

# Define ROS-related keywords to filter processes
ROS_KEYWORDS = {
    "ros",
    "node",
    "rviz",
    "rqt",
    "/opt/ros/",  # ROS installation path
    "/opt/carma/",  # CARMA ROS installation path
    "roscore",
    "rosmaster",
    "roslaunch",
    "rostopic",
    "rosnode",
    "rosbag",
    "ros2",
    "ros1_bridge",
    "rmw",  # ROS middleware
    "fastrtps",
    "cyclonedds",
    "rclcpp",
    "rclpy",
    "noetic",
    "foxy",
    "humble",
}

# Define processes to exclude (to avoid false positives)
# NOTE: Detection of these keywords overwrites the ROS_KEYWORDS
EXCLUDE_KEYWORDS = {"code", "chrome", "firefox", "vscode", "gnome"}


def parse_args():
    parser = argparse.ArgumentParser(description="Monitor CPU usage of ROS2 nodes")
    parser.add_argument(
        "--output-dir",
        "-o",
        default="carma-cpu-usage-logs",
        help="Directory to store output files (default: carma-cpu-usage-logs)",
    )
    parser.add_argument(
        "--include-pattern",
        "-i",
        help="Additional comma-separated patterns to include in process filtering",
    )
    return parser.parse_args()


def setup_logging_directory(output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    timestamp = datetime.now().strftime("%Y_%m_%d-%H_%M_%S")
    filename = f"cpu_usage_ros2_nodes_{timestamp}.csv"
    return os.path.join(output_dir, filename)


def is_ros_related_process(proc_info, cmdline):
    """
    Check if a process is ROS-related based on name and command line
    """
    # Convert process information to lowercase for case-insensitive matching
    name_lower = proc_info["name"].lower()
    cmdline_lower = cmdline.lower()

    # Check exclusions first
    if any(excl in name_lower or excl in cmdline_lower for excl in EXCLUDE_KEYWORDS):
        return False

    # Check for ROS-related keywords in process name and command line
    return any(
        keyword in name_lower or keyword in cmdline_lower for keyword in ROS_KEYWORDS
    )


def get_process_environment(pid):
    """
    Try to get ROS-related environment variables for a process
    """
    try:
        proc = psutil.Process(pid)
        env = proc.environ()
        ros_env = {k: v for k, v in env.items() if "ROS" in k}
        return bool(ros_env)
    except (psutil.NoSuchProcess, psutil.AccessDenied):
        return False


def main():
    args = parse_args()

    # Add any additional include patterns from command line
    if args.include_pattern:
        additional_patterns = set(args.include_pattern.split(","))
        ROS_KEYWORDS.update(additional_patterns)

    output_file = setup_logging_directory(args.output_dir)

    with open(output_file, mode="a") as file:
        writer = csv.writer(file)
        writer.writerow(
            [
                "Timestamp",
                "PID",
                "Process Name",
                "CPU (%)",
                "Memory (%)",
                "Command Line",
                "Total CPU (%)",
                "Total CPU Num",
                "Total Memory (%)",
                "Total Memory (GB)",
            ]
        )

        print(f"Starting to monitor the CPU usage data and saving to: {output_file}")

        # Total CPU and memory size in GB doesn't change
        total_memory_gb = psutil.virtual_memory().total / (1024**3)
        total_cpus = os.cpu_count()

        while True:
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
            total_cpu_percent = psutil.cpu_percent(interval=None)
            total_memory_percent = (
                psutil.virtual_memory().percent
            )  # Get total memory usage

            for proc in psutil.process_iter(
                ["pid", "name", "cpu_percent", "memory_percent", "cmdline"]
            ):
                try:
                    cmdline = (
                        " ".join(proc.info["cmdline"]) if proc.info["cmdline"] else ""
                    )

                    if is_ros_related_process(proc.info, cmdline):
                        pid = proc.info["pid"]
                        name = proc.info["name"]
                        cpu_percent = proc.info["cpu_percent"]
                        memory_percent = proc.info["memory_percent"]

                        writer.writerow(
                            [
                                timestamp,
                                pid,
                                name,
                                cpu_percent,
                                memory_percent,
                                cmdline,
                                total_cpu_percent,
                                total_cpus,
                                total_memory_percent,
                                total_memory_gb,
                            ]
                        )

                except (
                    psutil.NoSuchProcess,
                    psutil.AccessDenied,
                    psutil.ZombieProcess,
                ):
                    continue

            time.sleep(1)


if __name__ == "__main__":
    main()
