#!/usr/bin/python3
# SPDX-License-Identifier: GPL-2.0-only

"""
RT system profile save/restore functionality.

This module provides functions to save and restore RT system configurations
including cpuset definitions.
"""

import datetime
import sys
from ruamel.yaml import YAML


def save_cpusets(filename, get_cpusets_info_func):
    """Save cpuset configuration to a YAML profile.

    Args:
        filename: Path to output YAML file
        get_cpusets_info_func: Function to get cpuset info
                               (should be get_all_cpusets_info from tuna-cmd.py)

    Returns:
        Number of cpuset entries saved

    Note:
        Only saves user-created cpusets with CPUs assigned.
        System cpusets (.slice, .scope, .mount, .service) are automatically filtered out.
    """
    # Get all user-created cpusets with CPUs assigned
    # Reuses existing filtering logic from get_all_cpusets_info:
    #   skip_empty=True: Don't save cpusets with no CPUs
    #   skip_system=True: Don't save systemd-managed cpusets
    all_cpusets = get_cpusets_info_func(pattern=None, recursive=True,
                                        skip_empty=True, skip_system=True)

    if not all_cpusets:
        print("No user-created cpusets found to save", file=sys.stderr)
        return 0

    # Build list of cpuset entries for YAML
    # Reuses the structure from get_cpuset_info() - just extract what we need
    cpuset_entries = []

    for cs_info in all_cpusets:
        entry = {
            'name': cs_info['name'],
            'cpus': cs_info['cpus'],
            'mems': cs_info['mems'],
            'partition': cs_info['partition']
        }
        cpuset_entries.append(entry)

    # Create profile structure
    profile = {
        'version': '1.0',
        'timestamp': datetime.datetime.now().isoformat(),
        'cpusets': cpuset_entries
    }

    # Write YAML with ruamel.yaml for human readability
    yaml = YAML()
    yaml.default_flow_style = False
    yaml.width = 4096  # Avoid line wrapping

    with open(filename, 'w') as f:
        # Write header comments
        f.write("# Tuna RT System Profile - Cpuset Configuration\n")
        f.write(f"# Generated: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write("#\n")
        f.write("# This file can be applied with: tuna cpuset apply <filename>\n")
        f.write("#\n\n")

        # Write YAML structure
        yaml.dump(profile, f)

    return len(cpuset_entries)


def apply_cpusets(filename, cpuset_module, verbose=False):
    """Apply cpuset configuration from a YAML profile.

    Args:
        filename: Path to input YAML file
        cpuset_module: The cpuset module (passed as parameter for code reuse)
        verbose: If True, print detailed progress messages

    Returns:
        Tuple of (num_applied, num_warnings, num_errors)

    Note:
        Creates cpusets in the order they appear in the profile.
        Sets memory nodes before CPUs (cgroup v2 requirement).
        Existing cpusets with matching names are skipped with a warning.
    """
    # Read YAML file
    yaml = YAML()
    try:
        with open(filename, 'r') as f:
            profile = yaml.load(f)
    except FileNotFoundError:
        print(f"Error: Profile file '{filename}' not found", file=sys.stderr)
        return (0, 0, 1)
    except Exception as e:
        print(f"Error reading profile: {e}", file=sys.stderr)
        return (0, 0, 1)

    # Validate profile structure
    if not profile or 'cpusets' not in profile:
        print("Error: No cpusets found in profile", file=sys.stderr)
        return (0, 0, 1)

    # Check version compatibility
    version = profile.get('version', 'unknown')
    if verbose:
        print(f"Applying profile version {version}")
        timestamp = profile.get('timestamp', 'unknown')
        print(f"Generated: {timestamp}")

    cpuset_entries = profile['cpusets']
    if not cpuset_entries:
        print("No cpusets to apply", file=sys.stderr)
        return (0, 0, 0)

    num_applied = 0
    num_warnings = 0
    num_errors = 0

    # Apply each cpuset
    for entry in cpuset_entries:
        name = entry.get('name')
        cpus = entry.get('cpus', '')
        mems = entry.get('mems', '0')
        partition = entry.get('partition', 'member')

        if not name:
            print("Error: Cpuset entry missing 'name' field", file=sys.stderr)
            num_errors += 1
            continue

        if verbose:
            print(f"Creating cpuset '{name}' (CPUs: {cpus}, mems: {mems}, partition: {partition})")

        try:
            # Check if cpuset already exists
            import os
            cpuset_path = os.path.join('/sys/fs/cgroup', name)
            if os.path.exists(cpuset_path):
                print(f"Warning: Cpuset '{name}' already exists, skipping", file=sys.stderr)
                num_warnings += 1
                continue

            # Create cpuset
            cs = cpuset_module.Cpuset(name)

            # Set memory nodes BEFORE CPUs (cgroup v2 requirement)
            cs.write_memnode(mems)

            # Set CPUs
            if cpus:
                cs.assign_cpus(cpus)

            # Set partition type if isolated
            if partition == 'isolated':
                cs.set_cpu_exclusive()

            num_applied += 1

            if verbose:
                print(f"  ✓ Created cpuset '{name}'")

        except (OSError, ValueError) as e:
            print(f"Error creating cpuset '{name}': {e}", file=sys.stderr)
            num_errors += 1

    return (num_applied, num_warnings, num_errors)
