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

Tests the cpuset name extraction and display in show_threads -G output.
"""

import unittest
import os
import sys
import subprocess
import time
import importlib.util
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
import importlib.util

# Load tuna-cmd.py as a module
spec = importlib.util.spec_from_file_location("tuna_cmd", "tuna-cmd.py")
tuna_cmd = importlib.util.module_from_spec(spec)


class TestExtractCpusetName(unittest.TestCase):
    """Test extract_cpuset_name() function (non-root tests)"""

    @classmethod
    def setUpClass(cls):
        """Load the extract_cpuset_name function"""
        spec.loader.exec_module(tuna_cmd)
        cls.extract_cpuset_name = staticmethod(tuna_cmd.extract_cpuset_name)

    def test_extract_cpuset_name_basic(self):
        """Test basic cpuset name extraction"""
        result = self.extract_cpuset_name("0::/tuna0")
        self.assertEqual(result, "tuna0")

    def test_extract_cpuset_name_root(self):
        """Test root cgroup extraction"""
        result = self.extract_cpuset_name("0::/")
        self.assertEqual(result, "")

    def test_extract_cpuset_name_system_slice(self):
        """Test system.slice extraction"""
        result = self.extract_cpuset_name("0::/system.slice/ssh.service")
        self.assertEqual(result, "system.slice")

    def test_extract_cpuset_name_user_slice(self):
        """Test user.slice extraction"""
        result = self.extract_cpuset_name("0::/user.slice/user-1000.slice/session-2.scope")
        self.assertEqual(result, "user.slice")

    def test_extract_cpuset_name_nested(self):
        """Test nested cpuset extraction"""
        result = self.extract_cpuset_name("0::/parent/child")
        self.assertEqual(result, "parent")

    def test_extract_cpuset_name_empty_string(self):
        """Test empty string input"""
        result = self.extract_cpuset_name("")
        self.assertEqual(result, "")

    def test_extract_cpuset_name_none(self):
        """Test None input"""
        result = self.extract_cpuset_name(None)
        self.assertEqual(result, "")


@unittest.skipUnless(os.geteuid() == 0, "Requires root permissions")
@unittest.skipUnless(cpuset.CpusetsInit().supported, "Requires cgroup v2 support")
class TestShowThreadsCgroupsDisplay(unittest.TestCase):
    """Test show_threads -G cpuset display with root permissions"""

    @classmethod
    def setUpClass(cls):
        """Set up test environment once for all tests"""
        cls.ci = cpuset.CpusetsInit()
        cls.test_cpuset_name = 'test_cgroups_display'
        cls.sleep_proc = None

    def setUp(self):
        """Clean up before each test"""
        # Clean up any leftover test cpusets
        try:
            cpuset.destroy_cpuset(self.test_cpuset_name, force=True)
        except:
            pass

    def tearDown(self):
        """Clean up after each test"""
        # Kill sleep process if running
        if self.sleep_proc and self.sleep_proc.poll() is None:
            self.sleep_proc.kill()
            self.sleep_proc.wait()
            self.sleep_proc = None

        # Destroy test cpuset
        try:
            cpuset.destroy_cpuset(self.test_cpuset_name, force=True)
        except:
            pass

    def test_cgroups_displays_cpuset_name(self):
        """Test that -G displays cpuset name (not full path)"""
        # Create test cpuset
        cs = cpuset.Cpuset(self.test_cpuset_name)
        cs.assign_cpus('0-1')
        cs.write_memnode('0')

        # Start a sleep process
        self.sleep_proc = subprocess.Popen(['sleep', '60'])
        sleep_pid = self.sleep_proc.pid
        time.sleep(0.1)

        # Move to cpuset
        cs.write_pid(sleep_pid)
        time.sleep(0.1)

        # Run show_threads with -G
        result = subprocess.run(
            [sys.executable, 'tuna-cmd.py', 'show_threads', '-G', '-t', str(sleep_pid)],
            capture_output=True,
            text=True
        )

        self.assertEqual(result.returncode, 0)
        # Should show cpuset name, not full path
        self.assertIn(self.test_cpuset_name, result.stdout)
        # Should NOT show "0::/" prefix
        self.assertNotIn("0::", result.stdout)

    def test_cgroups_header_shows_cpuset(self):
        """Test that -G header shows 'cpuset' not 'cgroup'"""
        # Import necessary modules to test the header function directly
        import io
        import sys as test_sys
        from importlib import reload

        # Reload tuna_cmd to get fresh import
        spec = importlib.util.spec_from_file_location("tuna_cmd", "tuna-cmd.py")
        tuna_cmd = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(tuna_cmd)

        # Initialize i18n to set up the _() function
        tuna_cmd.i18n_init()

        # Capture output from ps_show_header
        old_stdout = test_sys.stdout
        test_sys.stdout = io.StringIO()

        try:
            # Call ps_show_header with cgroups=True
            tuna_cmd.ps_show_header(has_ctxt_switch_info=False, cgroups=True)
            output = test_sys.stdout.getvalue()

            # Header should say "cpuset" not "cgroup"
            self.assertIn("cpuset", output)
        finally:
            test_sys.stdout = old_stdout

    def test_cgroups_displays_root_as_empty(self):
        """Test that processes in root cgroup show empty cpuset"""
        # Get a process that's likely in root cgroup (PID 1 - systemd)
        result = subprocess.run(
            [sys.executable, 'tuna-cmd.py', 'show_threads', '-G', '-t', '1'],
            capture_output=True,
            text=True
        )

        self.assertEqual(result.returncode, 0)
        # Should NOT show "0::/" for root cgroup
        self.assertNotIn("0::/", result.stdout)

    def test_cgroups_displays_system_slice_shortened(self):
        """Test that system.slice processes show shortened name"""
        # Find a process in system.slice (many system daemons are there)
        # We'll look for any process and check the display format
        result = subprocess.run(
            [sys.executable, 'tuna-cmd.py', 'show_threads', '-G'],
            capture_output=True,
            text=True
        )

        self.assertEqual(result.returncode, 0)
        # If any system.slice processes exist, they should show "system.slice"
        # not the full path like "0::/system.slice/service.service"
        if "system.slice" in result.stdout:
            # Should NOT show "0::/" prefix
            self.assertNotIn("0::/system.slice/", result.stdout)


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