#!/usr/bin/python3
# -*- coding: utf-8 -*-
# SPDX-License-Identifier: GPL-2.0-only
#
# Copyright 2026 John Kacur <jkacur@redhat.com>
#
# Originally prototyped in tuna, implemented in rteval,
# and adapted back to tuna with tuna-specific enhancements.
#
""" Module for creating and using cpusets (cgroup v2) """

import errno
import os
import logging
from glob import glob

logger = logging.getLogger(__name__)

# Blocklist of critical processes that should never be moved to custom cpusets
# Moving these processes can break system functionality and prevent clean shutdown
PROCESS_BLOCKLIST = {
    'systemd',           # PID 1 - init system
    'systemd-logind',    # Session/power management - critical for shutdown
    'systemd-journald',  # Logging daemon
    'dbus-daemon',       # D-Bus system message bus
    'dbus-broker',       # Alternative D-Bus implementation
}


class Cpuset:
    """ Class for manipulating cpusets """

    mpath = '/sys/fs/cgroup'

    def __init__(self, name=None, auto_destroy=False, existing_ok=False):
        """
        Initialize a cpuset

        Args:
            name: Name of the cpuset to create
            auto_destroy: If True, cpuset will be automatically destroyed when
                         exiting a context manager (rteval style).
                         If False (default), cpuset persists after program exits
                         (tuna style).
            existing_ok: If True, suppress warning when cpuset already exists.
                        Use this when accessing an existing cpuset rather than
                        creating a new one (default: False).
        """
        self.cpuset_name = None
        self._cpuset_path = None
        self.auto_destroy = auto_destroy
        self.create_cpuset(name, existing_ok=existing_ok)

    @property
    def cpuset_path(self):
        """ Return the cpuset path, for example /sys/fs/cgroup/tuna_0 """
        return self._cpuset_path

    def __str__(self):
        return f'cpuset_name={self.cpuset_name}, cpuset_path={self.cpuset_path}'

    def create_cpuset(self, name=None, existing_ok=False):
        """
        Create a cpuset (cgroup) below /sys/fs/cgroup

        If the cpuset already exists, prints a warning (unless existing_ok=True)
        and returns without creating a new one.

        Args:
            name: Name of the cpuset to create
            existing_ok: If True, don't warn when cpuset already exists
        """
        if name is None:
            raise ValueError("cpuset name cannot be None")
        self.cpuset_name = name
        path = os.path.join(Cpuset.mpath, self.cpuset_name)

        # Check if cpuset already exists
        if os.path.exists(path):
            if not existing_ok:
                logger.warning(f"Cpuset '{name}' already exists at {path}")
                print(f"Warning: cpuset '{name}' already exists at {path}")
            self._cpuset_path = path
            return

        os.mkdir(path)
        self._cpuset_path = path
        # Enable cpuset controller in parent if needed
        self._enable_controllers()

    def _enable_controllers(self):
        """ Enable cpuset controller in parent cgroup """
        # Get parent path
        parent_path = os.path.dirname(self._cpuset_path)
        if parent_path == '/sys/fs':
            parent_path = '/sys/fs/cgroup'

        subtree_control = os.path.join(parent_path, 'cgroup.subtree_control')
        try:
            # Read current controllers
            with open(subtree_control, 'r', encoding='utf-8') as f:
                current = f.read().strip()

            # Enable cpuset if not already enabled
            # Controllers are space-separated, so split and check
            if 'cpuset' not in current.split():
                with open(subtree_control, 'w', encoding='utf-8') as f:
                    f.write('+cpuset')
        except (OSError, PermissionError):
            # May fail if already enabled or no permissions
            pass

    def assign_cpus(self, cpu_str):
        """
        Assign cpus in list-format to a cpuset
        Note: In cgroup v2, you must call write_memnode() before assign_cpus()
        """
        path = os.path.join(self._cpuset_path, "cpuset.cpus")
        with open(path, 'w', encoding='utf-8') as f:
            f.write(cpu_str)

    @staticmethod
    def _get_process_name(pid):
        """
        Get the process name (comm) for a given PID
        Returns None if the process doesn't exist or can't be read
        """
        try:
            with open(f'/proc/{pid}/comm', 'r', encoding='utf-8') as f:
                return f.read().strip()
        except (OSError, FileNotFoundError):
            return None

    @staticmethod
    def _is_process_blocklisted(pid):
        """
        Check if a process is on the blocklist of critical system processes
        Returns True if the process should not be moved to a custom cpuset
        """
        comm = Cpuset._get_process_name(pid)
        if comm is None:
            return False
        return comm in PROCESS_BLOCKLIST

    def write_pid(self, pid):
        """
        Place a pid in a cpuset
        Returns True if successful, False if the process cannot be moved
        Raises OSError for unexpected errors

        Note: Critical system processes (systemd, systemd-logind, etc.) are
        automatically skipped to prevent system instability and shutdown issues.
        """
        # Check if process is on the blocklist
        if self._is_process_blocklisted(pid):
            comm = self._get_process_name(pid)
            logger.info(f"Skipping blocklisted process {pid} ({comm})")
            return False

        path = os.path.join(self._cpuset_path, "cgroup.procs")
        pid_str = str(pid)
        try:
            with open(path, 'w', encoding='utf-8') as f:
                f.write(pid_str)
                f.flush()
            return True
        except OSError as err:
            # EINVAL: Invalid argument (process can't be moved, or invalid PID format)
            # ESRCH: No such process (process died between read and write)
            # EBUSY: Device or resource busy (process can't be moved)
            # EACCES: Permission denied (some kernel threads can't be moved)
            if err.errno in (errno.EINVAL, errno.ESRCH, errno.EBUSY, errno.EACCES):
                return False
            else:
                raise

    def write_memnode(self, memnode):
        """ Place memnode in a cpuset """
        path = os.path.join(self._cpuset_path, "cpuset.mems")
        memnode_str = str(memnode)
        with open(path, 'w', encoding='utf-8') as f:
            f.write(memnode_str)

    def write_cpu_exclusive(self, cpu_exclusive):
        """
        Set CPU partition type (cgroup v2)
        In v2, cpu_exclusive is replaced by cpuset.cpus.partition

        Args:
            cpu_exclusive: True/1/'1' for 'isolated', False/0/'0' for 'member'
        """
        path = os.path.join(self._cpuset_path, "cpuset.cpus.partition")
        if not os.path.exists(path):
            # cpuset.cpus.partition may not exist in all kernels
            return

        # Convert to boolean: handles bool, int, or string input
        if isinstance(cpu_exclusive, str):
            is_exclusive = cpu_exclusive not in ('0', '', 'false', 'False')
        else:
            is_exclusive = bool(cpu_exclusive)

        partition_type = 'isolated' if is_exclusive else 'member'
        try:
            with open(path, 'w', encoding='utf-8') as f:
                f.write(partition_type)
        except OSError:
            # Partition may not be supported or may fail
            pass

    def set_cpu_exclusive(self):
        """ Turn cpu_exclusive on (set partition to isolated) """
        self.write_cpu_exclusive(True)

    def unset_cpu_exclusive(self):
        """ Turn cpu_exclusive off (set partition to member) """
        self.write_cpu_exclusive(False)

    def write_memory_migrate(self, memory_migrate):
        """
        Memory migrate is not available in cgroup v2
        This is kept for API compatibility but does nothing
        """
        pass

    def set_memory_migrate(self):
        """ Memory_migrate not available in cgroup v2 (no-op for compatibility) """
        pass

    def unset_memory_migrate(self):
        """ Memory migrate not available in cgroup v2 (no-op for compatibility) """
        pass

    def destroy(self):
        """
        Remove this cpuset (cgroup directory)
        The cpuset must be empty (no processes) before it can be removed.
        Raises OSError if the directory is not empty or doesn't exist.
        """
        if self._cpuset_path and os.path.exists(self._cpuset_path):
            os.rmdir(self._cpuset_path)

    def __enter__(self):
        """
        Context manager entry point
        Allows usage: with Cpuset('name') as cpuset: ...
        """
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Context manager exit point
        Only destroys the cpuset if auto_destroy=True was set during initialization.
        This allows rteval-style automatic cleanup when desired, but preserves
        tuna's default behavior of persistent cpusets.
        """
        if self.auto_destroy:
            self.destroy()
        return False  # Don't suppress exceptions

    def get_tasks(self):
        """ Get the tasks (PIDs) currently in a cpuset """
        path = os.path.join(self.cpuset_path, "cgroup.procs")
        with open(path, 'r', encoding='utf-8') as f:
            return [line.strip() for line in f if line.strip()]

    def get_cpus(self):
        """
        Get the current CPU assignment as a string
        Returns string like "0-3" or "0,2,4-7" or empty string if none assigned
        """
        path = os.path.join(self._cpuset_path, "cpuset.cpus")
        with open(path, 'r', encoding='utf-8') as f:
            return f.read().strip()

    def get_memnode(self):
        """
        Get the current memory node assignment as a string
        Returns string like "0" or "0-1" or empty string if none assigned
        """
        path = os.path.join(self._cpuset_path, "cpuset.mems")
        with open(path, 'r', encoding='utf-8') as f:
            return f.read().strip()

    def get_partition_type(self):
        """
        Get the current partition type
        Returns 'isolated', 'member', or 'root' depending on cpuset.cpus.partition
        Returns None if partition file doesn't exist
        """
        path = os.path.join(self._cpuset_path, "cpuset.cpus.partition")
        if not os.path.exists(path):
            return None
        with open(path, 'r', encoding='utf-8') as f:
            return f.read().strip()

    def modify_cpus(self, add_cpus=None, remove_cpus=None):
        """
        Modify CPU assignment by adding and/or removing CPUs

        Args:
            add_cpus: List of CPU numbers to add (e.g., [4, 5, 6, 7])
            remove_cpus: List of CPU numbers to remove (e.g., [0, 1])

        Note: CPUs are added before removal. This is important when replacing CPUs
        to avoid having zero CPUs assigned (which may cause errors).
        """
        from tuna import tuna

        # Get current CPUs
        current_cpus_str = self.get_cpus()
        if current_cpus_str:
            current_cpus = set(tuna.cpustring_to_list(current_cpus_str))
        else:
            current_cpus = set()

        # Add new CPUs
        if add_cpus:
            current_cpus.update(add_cpus)

        # Remove CPUs
        if remove_cpus:
            current_cpus.difference_update(remove_cpus)

        # Convert back to string and assign
        if current_cpus:
            new_cpus_str = tuna.collapse_cpulist(sorted(current_cpus))
            self.assign_cpus(new_cpus_str)
        else:
            # If no CPUs left, assign empty string
            self.assign_cpus("")


class CpusetsInit(Cpuset):
    """
    CpusetsInit initializes cpusets (cgroup v2)
    It verifies that cgroup v2 is mounted and cpuset controller is available
    This class represents the root cgroup /sys/fs/cgroup
    """

    mpath = '/sys/fs/cgroup'

    def __init__(self):
        self._supported = self._cpuset_supported()
        self._numa_nodes = self._get_numa_nodes() if self._supported else 0
        if self._supported:
            # Set attributes directly for the root cgroup
            self.cpuset_name = None
            self._cpuset_path = Cpuset.mpath
            # Root cgroup is never auto-destroyed
            self.auto_destroy = False

    @property
    def supported(self):
        """ Return True if cpuset is supported """
        return self._supported

    @property
    def numa_nodes(self):
        """ Return the number of numa nodes """
        return self._numa_nodes

    @staticmethod
    def _cpuset_supported():
        """
        Return True if cpusets are supported
        For cgroup v2, check if:
        1. cgroup2 filesystem is available
        2. /sys/fs/cgroup is mounted
        3. cpuset controller is available
        """
        # Check if cgroup2 is in /proc/filesystems
        cgroup2_found = False
        dpath = '/proc/filesystems'
        try:
            with open(dpath, encoding='utf-8') as fp:
                for line in fp:
                    elems = line.strip().split()
                    if len(elems) == 2 and elems[1] == "cgroup2":
                        cgroup2_found = True
                        break
        except OSError:
            return False

        if not cgroup2_found:
            return False

        # Check if /sys/fs/cgroup is mounted and has cpuset controller
        cgroup_path = '/sys/fs/cgroup'
        controllers_file = os.path.join(cgroup_path, 'cgroup.controllers')

        if not os.path.exists(controllers_file):
            return False

        try:
            with open(controllers_file, encoding='utf-8') as f:
                controllers = f.read().strip().split()
                return 'cpuset' in controllers
        except OSError:
            return False

    @staticmethod
    def _get_numa_nodes():
        """
        return the number of numa nodes
        used by init to set _numa_nodes
        Users of this class should use the "numa_nodes" method
        """
        return len(glob('/sys/devices/system/node/node*'))


class TaskMigrate:
    """ A class for migrating a set of pids from one cpuset to another """

    def __init__(self, cpuset_from, cpuset_to):
        self.cpuset_from = cpuset_from
        self.cpuset_to = cpuset_to
        # memory_migrate is not available in cgroup v2, but keep for API compat
        self.cpuset_to.set_memory_migrate()
        self.migrated = 0
        self.failed = 0

    def migrate(self):
        """ Do the migration, only attempt on pids where this is allowed """
        self.migrated = 0
        self.failed = 0
        tasks = self.cpuset_from.get_tasks()
        for task in tasks:
            if self.cpuset_to.write_pid(task):
                self.migrated += 1
            else:
                self.failed += 1
        return self.migrated, self.failed


#
# NUMA topology detection (adapted from rteval/systopology.py)
#

class NumaNode:
    """Represents a NUMA node with its CPUs.

    Simplified from rteval.systopology.NumaNode for cpuset memory node detection.
    """

    def __init__(self, nodeid, cpus):
        """Initialize NUMA node.

        Args:
            nodeid: NUMA node ID (integer)
            cpus: Set of CPU numbers in this node
        """
        self.nodeid = nodeid
        self.cpus = cpus

    def __contains__(self, cpu):
        """Check if CPU is in this NUMA node."""
        return cpu in self.cpus

    def __int__(self):
        """Return node ID as integer."""
        return self.nodeid


def discover_numa_nodes():
    """Discover NUMA topology.

    Returns dict mapping node IDs to NumaNode objects.
    Handles both multi-node systems and single-node systems.

    Adapted from rteval.systopology.SysTopology.
    """
    nodes = {}
    node_paths = glob('/sys/devices/system/node/node[0-9]*')

    if node_paths:
        # Multi-node NUMA system
        for path in sorted(node_paths):
            nodeid = int(os.path.basename(path)[4:])
            cpulist_file = os.path.join(path, 'cpulist')
            try:
                with open(cpulist_file, 'r', encoding='utf-8') as f:
                    cpulist_str = f.read().strip()
                    # Parse CPU list string (e.g., "0-7" or "0,2,4-7")
                    from tuna import tuna
                    cpu_list = tuna.cpustring_to_list(cpulist_str)
                    nodes[nodeid] = NumaNode(nodeid, set(cpu_list))
            except (OSError, ValueError):
                # Skip nodes we can't read
                continue
    else:
        # Single-node system (no /sys/devices/system/node/node*)
        # Fallback like rteval.systopology.SimNumaNode
        try:
            with open('/sys/devices/system/cpu/possible', 'r', encoding='utf-8') as f:
                cpulist_str = f.read().strip()
                from tuna import tuna
                cpu_list = tuna.cpustring_to_list(cpulist_str)
                nodes[0] = NumaNode(0, set(cpu_list))
        except (OSError, ValueError):
            # Last resort: assume node 0 exists
            nodes[0] = NumaNode(0, set())

    return nodes


def get_numa_nodes_for_cpus(cpu_list):
    """Determine which NUMA nodes contain the given CPUs.

    Args:
        cpu_list: List of CPU numbers (e.g., [0, 1, 8, 9])

    Returns:
        String suitable for cpuset.mems (e.g., "0", "0-1", "0,2")

    Example:
        On a 2-node system with CPUs 0-7 on node0, 8-15 on node1:
        - get_numa_nodes_for_cpus([0,1,2]) returns "0"
        - get_numa_nodes_for_cpus([8,9,10]) returns "1"
        - get_numa_nodes_for_cpus([0,1,8,9]) returns "0-1"
    """
    if not cpu_list:
        # No CPUs specified, default to node 0
        return "0"

    nodes = discover_numa_nodes()

    # Find which nodes contain our CPUs
    required_nodes = set()
    for cpu in cpu_list:
        for nodeid, node in nodes.items():
            if cpu in node:
                required_nodes.add(nodeid)
                break

    if not required_nodes:
        # Fallback if no matches found (shouldn't happen)
        return "0"

    # Convert node IDs to string format
    from tuna import tuna
    return tuna.collapse_cpulist(sorted(required_nodes))


#
# Tuna-specific utility functions for cpuset discovery and management
#

def list_cpusets(pattern=None, recursive=True):
    """
    List all cpusets under /sys/fs/cgroup

    Args:
        pattern: Optional glob pattern to filter (e.g., "tuna_*", "rteval_*")
                If None, returns all cpusets.
        recursive: If True, find nested cpusets; if False, only top-level

    Returns:
        List of cpuset paths relative to /sys/fs/cgroup

    Note: Real-world scenario - systemd, containers, and users may create
    nested cgroup hierarchies, so recursive search is important
    """
    cpusets = []
    base_path = '/sys/fs/cgroup'

    if not recursive:
        # Only search top-level directories
        import fnmatch
        for entry in os.listdir(base_path):
            entry_path = os.path.join(base_path, entry)
            if os.path.isdir(entry_path):
                # Check if it's a cpuset by looking for cpuset.cpus file
                if os.path.exists(os.path.join(entry_path, 'cpuset.cpus')):
                    if pattern is None or fnmatch.fnmatch(entry, pattern):
                        cpusets.append(entry)
    else:
        # Recursive search using os.walk
        for root, dirs, files in os.walk(base_path):
            # Check if this directory has cpuset.cpus (indicates it's a cpuset)
            if 'cpuset.cpus' in files:
                # Get relative path from base_path
                rel_path = os.path.relpath(root, base_path)
                if rel_path == '.':
                    # Skip the root cgroup itself
                    continue
                if pattern is None:
                    cpusets.append(rel_path)
                else:
                    # Match pattern against the basename only
                    import fnmatch
                    if fnmatch.fnmatch(os.path.basename(rel_path), pattern):
                        cpusets.append(rel_path)

    return sorted(cpusets)


def destroy_cpuset(name_or_path, force=False):
    """
    Destroy a cpuset by name or path

    Args:
        name_or_path: Cpuset name (e.g., "tuna_0") or full path
                     (e.g., "/sys/fs/cgroup/tuna_0")
        force: If True, migrate tasks out to root cgroup before destroying

    Raises:
        OSError: If the cpuset doesn't exist or can't be removed
        ValueError: If trying to destroy the root cgroup

    Note: When cleaning nested hierarchies, must destroy children before parents
    (cgroups can't be removed if they have children)
    """
    base_path = '/sys/fs/cgroup'

    # Normalize the path
    if name_or_path.startswith('/'):
        cpuset_path = name_or_path
    else:
        cpuset_path = os.path.join(base_path, name_or_path)

    # Prevent destroying root cgroup
    if cpuset_path == base_path or cpuset_path == '/sys/fs/cgroup':
        raise ValueError("Cannot destroy the root cgroup")

    if not os.path.exists(cpuset_path):
        raise OSError(f"Cpuset does not exist: {cpuset_path}")

    # If force=True, migrate tasks to root cgroup first
    if force:
        root_cgroup = CpusetsInit()
        if not root_cgroup.supported:
            raise RuntimeError("Cpusets not supported on this system")

        # Create temporary Cpuset object for this path
        cpuset = Cpuset.__new__(Cpuset)
        cpuset._cpuset_path = cpuset_path
        cpuset.cpuset_name = os.path.basename(cpuset_path)

        # Migrate tasks to root
        tm = TaskMigrate(cpuset, root_cgroup)
        migrated, failed = tm.migrate()
        logger.info(f"Migrated {migrated} tasks to root cgroup, {failed} failed")

    # Destroy the cpuset
    os.rmdir(cpuset_path)
    logger.info(f"Destroyed cpuset: {cpuset_path}")


def destroy_cpuset_recursive(name_or_path, force=False):
    """
    Recursively destroy a cpuset and all its children

    Args:
        name_or_path: Cpuset name or full path
        force: If True, migrate tasks out before destroying

    This is useful when a cpuset has nested children that need to be
    destroyed before the parent can be removed.
    """
    base_path = '/sys/fs/cgroup'

    # Normalize the path
    if name_or_path.startswith('/'):
        cpuset_path = name_or_path
    else:
        cpuset_path = os.path.join(base_path, name_or_path)

    if not os.path.exists(cpuset_path):
        return  # Already destroyed or doesn't exist

    # Find all children
    children = []
    for entry in os.listdir(cpuset_path):
        child_path = os.path.join(cpuset_path, entry)
        if os.path.isdir(child_path) and os.path.exists(os.path.join(child_path, 'cgroup.procs')):
            children.append(child_path)

    # Recursively destroy children first
    for child_path in children:
        destroy_cpuset_recursive(child_path, force=force)

    # Now destroy this cpuset
    try:
        destroy_cpuset(cpuset_path, force=force)
    except (OSError, ValueError) as e:
        logger.warning(f"Failed to destroy {cpuset_path}: {e}")


def cleanup_cpusets(pattern, force=True, recursive=True):
    """
    Clean up multiple cpusets matching a pattern

    Args:
        pattern: Glob pattern (e.g., "tuna_*" or "rteval_*")
        force: If True, migrate tasks out before destroying
        recursive: If True, find and clean nested cpusets too

    Note: When cleaning nested hierarchies, must destroy children before parents
    (cgroups can't be removed if they have children). This function will
    automatically handle nested children even if they don't match the pattern.
    """
    # Find all matching cpusets
    cpusets = list_cpusets(pattern=pattern, recursive=recursive)

    if not cpusets:
        logger.info(f"No cpusets found matching pattern: {pattern}")
        return

    # Sort in reverse order by depth to ensure children are destroyed before parents
    # Depth is determined by the number of path separators
    cpusets_by_depth = sorted(cpusets, key=lambda x: x.count(os.sep), reverse=True)

    destroyed_count = 0
    failed_count = 0
    # Keep track of what we've already destroyed to avoid double-counting
    destroyed_paths = set()

    for cpuset_name in cpusets_by_depth:
        cpuset_path = os.path.join('/sys/fs/cgroup', cpuset_name)

        # Skip if already destroyed
        if cpuset_path in destroyed_paths or not os.path.exists(cpuset_path):
            continue

        try:
            # Use recursive destroy to handle any nested children
            # (even those that don't match the pattern)
            destroy_cpuset_recursive(cpuset_name, force=force)
            destroyed_count += 1
            destroyed_paths.add(cpuset_path)
        except (OSError, ValueError) as e:
            logger.error(f"Failed to destroy {cpuset_name}: {e}")
            failed_count += 1

    logger.info(f"Cleanup complete: {destroyed_count} destroyed, {failed_count} failed")
    if destroyed_count > 0 or failed_count > 0:
        print(f"Cleanup complete: {destroyed_count} cpusets destroyed, {failed_count} failed")


if __name__ == '__main__':

    cpusets_init = CpusetsInit()
    print(f'cpusets_init.supported = {cpusets_init.supported}')
    if cpusets_init.supported:
        print(f'cpusets_init.numa_nodes = {cpusets_init.numa_nodes}')
        print(f'cpusets_init.cpuset_path = {cpusets_init.cpuset_path}')

        # Creating and manipulating cgroups requires root permissions
        print("\nTo test cpuset creation, run as root:")
        print("  sudo python3 -m tuna.cpuset")

        try:
            # Example 1: Using context manager with auto_destroy (rteval style)
            print("\n" + "="*60)
            print("Example 1: Context Manager with auto_destroy=True")
            print("="*60)
            with Cpuset('tuna_demo', auto_destroy=True) as cpuset:
                print(f"Created cpuset: {cpuset.cpuset_path}")
                cpuset.write_memnode('0')
                cpuset.assign_cpus('0-3')
                print(f"Configured cpuset with CPUs 0-3")
            print("Context manager exited - cpuset automatically destroyed!")

            # Example 2: Persistent cpuset (tuna style, default)
            print("\n" + "="*60)
            print("Example 2: Persistent Cpuset (tuna style, default)")
            print("="*60)
            print("Creating tuna_0...")
            cpuset0 = Cpuset('tuna_0')  # auto_destroy=False by default
            print("Creating tuna_1...")
            cpuset1 = Cpuset('tuna_1')

            # In cgroup v2, must set mems before cpus
            print("Setting memory nodes...")
            cpuset0.write_memnode('0')
            cpuset1.write_memnode('0')

            print("Assigning CPUs to tuna_0 (0-4)...")
            cpuset0.assign_cpus('0-4')
            print("Assigning CPUs to tuna_1 (5-7)...")
            cpuset1.assign_cpus('5-7')

            print("Setting cpu_exclusive on tuna_1...")
            cpuset1.set_cpu_exclusive()

            print("\nCpusets created and will persist after program exits.")
            print("Use cleanup_cpusets('tuna_*') to remove them later.")

            # Example 3: Tuna-specific discovery and cleanup functions
            print("\n" + "="*60)
            print("Example 3: Tuna-specific Discovery and Cleanup")
            print("="*60)
            print("\nListing all cpusets matching 'tuna_*':")
            tuna_cpusets = list_cpusets(pattern='tuna_*')
            for cs in tuna_cpusets:
                print(f"  - {cs}")

            print("\nCleaning up with cleanup_cpusets('tuna_*', force=True)...")
            cleanup_cpusets('tuna_*', force=True)

            print("\nAfter cleanup, listing cpusets again:")
            remaining = list_cpusets(pattern='tuna_*')
            if remaining:
                print(f"  Still found: {remaining}")
            else:
                print("  No tuna_* cpusets remaining (cleanup successful)")

            # Success message
            print("\n" + "="*60)
            print("✅ All examples completed successfully!")
            print("="*60)

        except PermissionError as e:
            print(f"\nPermission denied: {e}")
            print("Run as root to create and manipulate cgroups")
        except Exception as e:
            print(f"\nError: {e}")
            import traceback
            traceback.print_exc()
