#!/usr/bin/python3
# -*- coding: utf-8 -*-
# SPDX-License-Identifier: GPL-2.0-only
"""
Test suite for tuna cpuset save/apply commands

Tests the cpuset save and apply functionality:
- save: Saving cpuset configuration to YAML
- apply: Restoring cpuset configuration from YAML
- Round-trip: save → destroy → apply → verify
- Edge cases: existing cpusets, missing files, empty profiles, etc.

All tests clean up created files and cpusets afterwards.
"""

import unittest
import os
import sys
import tempfile
from pathlib import Path

# Add tuna to path if running from tests directory
if Path(__file__).parent.name == 'tests':
    sys.path.insert(0, str(Path(__file__).parent.parent))

from tuna import cpuset, profile


# Import the CLI functions from tuna-cmd.py
tuna_cmd_path = Path(__file__).parent.parent / 'tuna-cmd.py'
spec = __import__('importlib.util').util.spec_from_file_location("tuna_cmd", tuna_cmd_path)
tuna_cmd = __import__('importlib.util').util.module_from_spec(spec)
spec.loader.exec_module(tuna_cmd)

# Import the functions we want to test
save_cpusets_cmd = tuna_cmd.save_cpusets_cmd
apply_cpusets_cmd = tuna_cmd.apply_cpusets_cmd
get_all_cpusets_info = tuna_cmd.get_all_cpusets_info


@unittest.skipUnless(os.geteuid() == 0, "Requires root permissions")
@unittest.skipUnless(cpuset.CpusetsInit().supported, "Requires cgroup v2 support")
class TestSaveApply(unittest.TestCase):
    """Test save and apply commands for cpuset configuration"""

    def setUp(self):
        """Create temporary file and test cpusets before each test"""
        self.ci = cpuset.CpusetsInit()

        # Create temporary YAML file
        self.temp_fd, self.temp_file = tempfile.mkstemp(suffix='.yaml', prefix='tuna-test-')
        os.close(self.temp_fd)  # Close fd, we'll use the filename

        # Clean up any existing test cpusets (including from manual tests)
        for pattern in ['test_save_*', 'test_rt', 'test_housekeeping']:
            try:
                cpuset.cleanup_cpusets(pattern, force=True, recursive=False)
            except:
                pass

    def tearDown(self):
        """Clean up test cpusets and temporary file after each test"""
        # Clean up test cpusets (including from manual tests)
        for pattern in ['test_save_*', 'test_rt', 'test_housekeeping']:
            try:
                cpuset.cleanup_cpusets(pattern, force=True, recursive=False)
            except:
                pass

        # Remove temporary YAML file
        if os.path.exists(self.temp_file):
            os.unlink(self.temp_file)

    def create_test_cpusets(self):
        """Helper: Create two test cpusets for testing"""
        # Create test_save_rt (isolated)
        cs_rt = cpuset.Cpuset('test_save_rt')
        cs_rt.write_memnode('0')
        cs_rt.assign_cpus('2-3')
        cs_rt.set_cpu_exclusive()

        # Create test_save_hk (member)
        cs_hk = cpuset.Cpuset('test_save_hk')
        cs_hk.write_memnode('0')
        cs_hk.assign_cpus('0-1')

        return cs_rt, cs_hk

    def test_save_creates_yaml_file(self):
        """Test that save command creates a YAML file"""
        self.create_test_cpusets()

        # Save cpusets
        save_cpusets_cmd(self.temp_file)

        # Verify file exists and has content
        self.assertTrue(os.path.exists(self.temp_file))
        self.assertGreater(os.path.getsize(self.temp_file), 0)

    def test_save_captures_all_cpusets(self):
        """Test that save captures all user cpusets"""
        self.create_test_cpusets()

        # Save cpusets
        save_cpusets_cmd(self.temp_file)

        # Read and verify YAML content
        from ruamel.yaml import YAML
        yaml = YAML()
        with open(self.temp_file, 'r') as f:
            data = yaml.load(f)

        self.assertIn('cpusets', data)
        self.assertEqual(len(data['cpusets']), 2)

        # Check cpuset names
        names = {cs['name'] for cs in data['cpusets']}
        self.assertEqual(names, {'test_save_rt', 'test_save_hk'})

    def test_save_captures_partition_types(self):
        """Test that save preserves partition types (isolated vs member)"""
        self.create_test_cpusets()

        # Save cpusets
        save_cpusets_cmd(self.temp_file)

        # Read YAML
        from ruamel.yaml import YAML
        yaml = YAML()
        with open(self.temp_file, 'r') as f:
            data = yaml.load(f)

        # Find each cpuset in saved data
        cpusets_by_name = {cs['name']: cs for cs in data['cpusets']}

        self.assertEqual(cpusets_by_name['test_save_rt']['partition'], 'isolated')
        self.assertEqual(cpusets_by_name['test_save_hk']['partition'], 'member')

    def test_save_captures_cpu_ranges(self):
        """Test that save captures CPU assignments correctly"""
        self.create_test_cpusets()

        # Save cpusets
        save_cpusets_cmd(self.temp_file)

        # Read YAML
        from ruamel.yaml import YAML
        yaml = YAML()
        with open(self.temp_file, 'r') as f:
            data = yaml.load(f)

        # Find each cpuset in saved data
        cpusets_by_name = {cs['name']: cs for cs in data['cpusets']}

        self.assertEqual(cpusets_by_name['test_save_rt']['cpus'], '2-3')
        self.assertEqual(cpusets_by_name['test_save_hk']['cpus'], '0-1')

    def test_apply_recreates_cpusets(self):
        """Test that apply recreates cpusets from YAML"""
        # Create and save cpusets
        self.create_test_cpusets()
        save_cpusets_cmd(self.temp_file)

        # Destroy cpusets
        cpuset.cleanup_cpusets('test_save_*', force=True, recursive=False)

        # Verify they're gone
        remaining = cpuset.list_cpusets(pattern='test_save_*', recursive=False)
        self.assertEqual(len(remaining), 0)

        # Apply profile to recreate
        apply_cpusets_cmd(self.temp_file, verbose=False)

        # Verify cpusets were recreated
        recreated = cpuset.list_cpusets(pattern='test_save_*', recursive=False)
        self.assertEqual(len(recreated), 2)
        self.assertIn('test_save_rt', recreated)
        self.assertIn('test_save_hk', recreated)

    def test_apply_preserves_partition_types(self):
        """Test that apply restores partition types correctly"""
        # Create, save, destroy, apply
        self.create_test_cpusets()
        save_cpusets_cmd(self.temp_file)
        cpuset.cleanup_cpusets('test_save_*', force=True, recursive=False)
        apply_cpusets_cmd(self.temp_file, verbose=False)

        # Verify partition types
        cs_rt = cpuset.Cpuset('test_save_rt', existing_ok=True)
        cs_hk = cpuset.Cpuset('test_save_hk', existing_ok=True)

        self.assertEqual(cs_rt.get_partition_type(), 'isolated')
        self.assertEqual(cs_hk.get_partition_type(), 'member')

    def test_apply_preserves_cpu_assignments(self):
        """Test that apply restores CPU assignments correctly"""
        # Create, save, destroy, apply
        self.create_test_cpusets()
        save_cpusets_cmd(self.temp_file)
        cpuset.cleanup_cpusets('test_save_*', force=True, recursive=False)
        apply_cpusets_cmd(self.temp_file, verbose=False)

        # Verify CPU assignments
        cs_rt = cpuset.Cpuset('test_save_rt', existing_ok=True)
        cs_hk = cpuset.Cpuset('test_save_hk', existing_ok=True)

        self.assertEqual(cs_rt.get_cpus(), '2-3')
        self.assertEqual(cs_hk.get_cpus(), '0-1')

    def test_apply_with_existing_cpusets_warns(self):
        """Test that apply warns when cpusets already exist"""
        # Create and save cpusets
        self.create_test_cpusets()
        save_cpusets_cmd(self.temp_file)

        # Apply without destroying (cpusets already exist)
        # This should return warnings
        num_applied, num_warnings, num_errors = profile.apply_cpusets(
            self.temp_file, cpuset, verbose=False
        )

        self.assertEqual(num_applied, 0)  # Nothing new created
        self.assertEqual(num_warnings, 2)  # Both cpusets already exist
        self.assertEqual(num_errors, 0)

    def test_apply_missing_file_returns_error(self):
        """Test that apply returns error for missing file"""
        missing_file = '/tmp/nonexistent-tuna-test.yaml'

        num_applied, num_warnings, num_errors = profile.apply_cpusets(
            missing_file, cpuset, verbose=False
        )

        self.assertEqual(num_applied, 0)
        self.assertEqual(num_warnings, 0)
        self.assertEqual(num_errors, 1)

    def test_roundtrip_preserves_configuration(self):
        """Test full round-trip: save → destroy → apply → verify"""
        # Step 1: Create test cpusets
        cs_rt, cs_hk = self.create_test_cpusets()
        original_rt_cpus = cs_rt.get_cpus()
        original_hk_cpus = cs_hk.get_cpus()
        original_rt_partition = cs_rt.get_partition_type()
        original_hk_partition = cs_hk.get_partition_type()

        # Step 2: Save configuration
        save_cpusets_cmd(self.temp_file)

        # Step 3: Destroy cpusets
        cpuset.cleanup_cpusets('test_save_*', force=True, recursive=False)

        # Step 4: Apply configuration
        apply_cpusets_cmd(self.temp_file, verbose=False)

        # Step 5: Verify configuration matches original
        cs_rt_restored = cpuset.Cpuset('test_save_rt', existing_ok=True)
        cs_hk_restored = cpuset.Cpuset('test_save_hk', existing_ok=True)

        self.assertEqual(cs_rt_restored.get_cpus(), original_rt_cpus)
        self.assertEqual(cs_hk_restored.get_cpus(), original_hk_cpus)
        self.assertEqual(cs_rt_restored.get_partition_type(), original_rt_partition)
        self.assertEqual(cs_hk_restored.get_partition_type(), original_hk_partition)

    def test_save_with_no_cpusets(self):
        """Test that save handles case with no cpusets gracefully"""
        # Don't create any cpusets
        # save_cpusets_cmd will print to stderr but shouldn't crash

        # This should complete without error but save 0 cpusets
        num_saved = profile.save_cpusets(self.temp_file, get_all_cpusets_info)
        self.assertEqual(num_saved, 0)

    def test_save_skips_system_cpusets(self):
        """Test that save automatically filters out system cpusets"""
        # Create one test cpuset
        cs = cpuset.Cpuset('test_save_single')
        cs.write_memnode('0')
        cs.assign_cpus('0')

        # Save (should only save test_save_single, not any .slice/.scope/.service cpusets)
        save_cpusets_cmd(self.temp_file)

        # Read YAML
        from ruamel.yaml import YAML
        yaml = YAML()
        with open(self.temp_file, 'r') as f:
            data = yaml.load(f)

        # Verify only our test cpuset was saved (no system cpusets)
        saved_names = [cs['name'] for cs in data['cpusets']]
        for name in saved_names:
            self.assertFalse(name.endswith(('.slice', '.scope', '.mount', '.service')))

        # Cleanup
        cpuset.destroy_cpuset('test_save_single', force=True)


if __name__ == '__main__':
    unittest.main()
