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

Tests the --cpuset filter option added to show_threads command.
"""

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


class TestShowThreadsCpusetFilterNonRoot(unittest.TestCase):
    """Test show_threads --cpuset filter without root (basic tests)"""

    def test_help_shows_cpuset_option(self):
        """Test that --cpuset option appears in help"""
        result = subprocess.run(
            [sys.executable, 'tuna-cmd.py', 'show_threads', '--help'],
            capture_output=True,
            text=True
        )
        self.assertEqual(result.returncode, 0)
        self.assertIn('--cpuset', result.stdout)
        self.assertIn('Show only threads in the specified cpuset', result.stdout)


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

    @classmethod
    def setUpClass(cls):
        """Set up test environment once for all tests"""
        cls.ci = cpuset.CpusetsInit()
        cls.test_cpuset_name = 'test_show_filter'
        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_filter_shows_only_cpuset_processes(self):
        """Test that --cpuset filter shows only processes in specified cpuset"""
        # Create test cpuset
        cs = cpuset.Cpuset(self.test_cpuset_name)
        cs.assign_cpus('0-1')
        cs.write_memnode('0')

        # Start a sleep process in the background
        self.sleep_proc = subprocess.Popen(['sleep', '60'])
        sleep_pid = self.sleep_proc.pid
        time.sleep(0.1)  # Give process time to start

        # Move sleep process to test cpuset
        cs.write_pid(sleep_pid)
        time.sleep(0.1)  # Give cgroup time to update

        # Verify process is in the cpuset
        with open(f'/proc/{sleep_pid}/cgroup', 'r') as f:
            cgroup = f.read().strip()
        self.assertIn(self.test_cpuset_name, cgroup)

        # Run show_threads with cpuset filter
        result = subprocess.run(
            [sys.executable, 'tuna-cmd.py', 'show_threads', '--cpuset', self.test_cpuset_name],
            capture_output=True,
            text=True
        )

        self.assertEqual(result.returncode, 0)
        # Check that sleep process is in output
        self.assertIn(str(sleep_pid), result.stdout)
        self.assertIn('sleep', result.stdout)

    def test_filter_with_nonexistent_cpuset(self):
        """Test that filtering with non-existent cpuset shows no processes"""
        # Run show_threads with non-existent cpuset
        result = subprocess.run(
            [sys.executable, 'tuna-cmd.py', 'show_threads', '--cpuset', 'nonexistent_cpuset'],
            capture_output=True,
            text=True
        )

        self.assertEqual(result.returncode, 0)
        # Output should be empty or just header
        lines = [line for line in result.stdout.split('\n') if line.strip()]
        # Should have at most 2 lines (header rows) and no process entries
        self.assertLessEqual(len(lines), 2, "Should show no processes for non-existent cpuset")

    def test_filter_excludes_other_cpuset_processes(self):
        """Test that filter excludes processes from other cpusets"""
        # Create first test cpuset
        cs1 = cpuset.Cpuset(self.test_cpuset_name)
        cs1.assign_cpus('0-1')
        cs1.write_memnode('0')

        # Create second cpuset
        cs2_name = f'{self.test_cpuset_name}_2'
        try:
            cs2 = cpuset.Cpuset(cs2_name)
            cs2.assign_cpus('0-1')
            cs2.write_memnode('0')

            # Start two sleep processes
            self.sleep_proc = subprocess.Popen(['sleep', '60'])
            sleep_pid1 = self.sleep_proc.pid
            sleep_proc2 = subprocess.Popen(['sleep', '60'])
            sleep_pid2 = sleep_proc2.pid
            time.sleep(0.1)

            # Move first to test_cpuset_name, second to test_cpuset_name_2
            cs1.write_pid(sleep_pid1)
            cs2.write_pid(sleep_pid2)
            time.sleep(0.1)

            # Filter for first cpuset
            result = subprocess.run(
                [sys.executable, 'tuna-cmd.py', 'show_threads', '--cpuset', self.test_cpuset_name],
                capture_output=True,
                text=True
            )

            self.assertEqual(result.returncode, 0)
            # Should show first process but not second
            self.assertIn(str(sleep_pid1), result.stdout)
            self.assertNotIn(str(sleep_pid2), result.stdout)

            # Clean up second process and cpuset
            sleep_proc2.kill()
            sleep_proc2.wait()

        finally:
            # Clean up second cpuset
            try:
                cpuset.destroy_cpuset(cs2_name, force=True)
            except:
                pass

    def test_filter_with_cgroups_option(self):
        """Test that --cpuset works together with -G (cgroups display)"""
        # 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 with both --cpuset and -G
        result = subprocess.run(
            [sys.executable, 'tuna-cmd.py', 'show_threads', '--cpuset', self.test_cpuset_name, '-G'],
            capture_output=True,
            text=True
        )

        self.assertEqual(result.returncode, 0)
        self.assertIn(str(sleep_pid), result.stdout)
        # Should show cgroup path in output
        self.assertIn(self.test_cpuset_name, result.stdout)

    def test_filter_with_thread_list(self):
        """Test that --cpuset can be combined with -t (thread list)"""
        # 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 with both --cpuset and -t
        result = subprocess.run(
            [sys.executable, 'tuna-cmd.py', 'show_threads', '--cpuset', self.test_cpuset_name, '-t', str(sleep_pid)],
            capture_output=True,
            text=True
        )

        self.assertEqual(result.returncode, 0)
        self.assertIn(str(sleep_pid), result.stdout)

    def test_filter_empty_cpuset(self):
        """Test filtering on an empty cpuset (no processes)"""
        # Create empty cpuset
        cs = cpuset.Cpuset(self.test_cpuset_name)
        cs.assign_cpus('0-1')
        cs.write_memnode('0')

        # Filter for empty cpuset
        result = subprocess.run(
            [sys.executable, 'tuna-cmd.py', 'show_threads', '--cpuset', self.test_cpuset_name],
            capture_output=True,
            text=True
        )

        self.assertEqual(result.returncode, 0)
        # Should show no processes
        lines = [line for line in result.stdout.split('\n') if line.strip()]
        self.assertLessEqual(len(lines), 2, "Empty cpuset should show no processes")


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