# Copyright 2022 The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Unit tests for the Task Mixins."""
import math
import os
import shutil
import signal
import subprocess
import tempfile
import textwrap
import time
from pathlib import Path
from typing import Union
from unittest import TestCase

import psutil

from debusine.tasks import Task
from debusine.tasks._task_mixins import TaskRunCommandMixin


class TaskWithRunCommandMixin(Task, TaskRunCommandMixin):
    """Used in TaskRunCommandMixinTests."""

    def __init__(
        self,
        abort_after_aborted_calls: Union[int, float],
        wait_popen_sigkill=False,
        pids_directory=None,
        wait_for_sigusr1=True,
        send_signal_to_cmd=None,
    ):
        """
        Initialise object used in the tests.

        :param abort_after_aborted_calls: abort() method will return False
          until this number. Then will return True on each subsequent call
        :param wait_popen_sigkill: if True the first call to _wait_popen()
          will do super()._wait_popen() AND also kill the process group.
        :param pids_directory: the directory where the PID.log files are
          expected to appear.
        :param wait_for_sigusr1: if True on the first _wait_popen will
          block until SIGUSR1 is received
        :param send_signal_to_cmd: if it's a signal (SIGUSR1, SIGUSR2) will
          run the command normally. When self.aborted is accessed twice it
          will send the signal to the command
        """
        super().__init__()

        self._abort_after_aborted_calls = abort_after_aborted_calls
        self._wait_popen_sigkill = wait_popen_sigkill
        self._pids_directory = pids_directory
        self._wait_for_sigusr1 = wait_for_sigusr1
        self._send_signal_to_cmd = send_signal_to_cmd

        self._children_initialized = False
        self._popen = None

        self.aborted_call_count = 0

    @property
    def aborted(self):
        """
        Return False (not aborted) or True.

        True when has been called more times than abort_after_aborted_calls
        parameter used in the __init__().

        If send_signal_to_cmd it will send the signal to the PID.
        """
        self.aborted_call_count += 1

        if self.aborted_call_count == 2 and self._send_signal_to_cmd:
            os.kill(self._popen.pid, self._send_signal_to_cmd)

        return self.aborted_call_count > self._abort_after_aborted_calls

    def _do_wait_for_sigusr1(self):
        """
        Will wait up to 5 seconds to receive SIGUSR1.

        SIGUSR1 is sent by the signal-logger.sh when the last process
        is spawned.
        """
        if self._children_initialized:
            return

        result = signal.sigtimedwait({signal.SIGUSR1}, 5)

        if result is None:  # pragma: no cover
            # No signal received after the timeout. This should never
            # happen: signal-logger.sh should have sent on time.
            # (the timeout is to fail faster in case that signal-logger.sh
            # is not working)
            raise RuntimeError("SIGUSR1 not received")

        self._children_initialized = True

    def wait_cmd_zombie_process(self):
        """
        Wait and get the returncode for cmd.

        TaskMixins SIGKILLed cmd. cmd probably was in a non-interruptable
        call in the Kernel and TaskMixins did not wait and cmd is not
        left zombie. Here it collects the returncode, so it is
        not zombie and avoids a Popen.__del__ ResourceWarning
        """
        self._popen.wait(timeout=5)

    def _wait_popen(self, popen, timeout):  # noqa: U100
        # Calls TaskMixins._wait_popen() with a short timeout

        # Depending on self._wait_popen_sigkill will send a SIGKILL
        # to the process group

        self._popen = popen  # Used by wait_cmd_zombie_process()

        if self._wait_for_sigusr1:
            self._do_wait_for_sigusr1()

        try:
            result = super()._wait_popen(popen, 0.001)
        except subprocess.TimeoutExpired as exc:
            result = exc

        if self._wait_popen_sigkill:
            pid = int(next(self._pids_directory.glob("*.log")).stem)

            process_group = os.getpgid(pid)

            os.killpg(process_group, signal.SIGKILL)

            # If _wait_popen is called again: no killing again
            self._wait_popen_sigkill = False

        if isinstance(result, Exception):
            raise result
        else:
            return result


class TaskRunCommandMixinTests(TestCase):
    """Test the TaskRunCommandMixin class."""

    def setUp(self):
        """Set up test object."""
        self.temp_directory = Path(
            tempfile.mkdtemp(prefix="debusine-test-task-mixins-")
        )

        self.addCleanup(lambda: shutil.rmtree(self.temp_directory))

    def run_signal_logger(self, task, options):
        """Run signal logger in task with options."""
        # Block SIGUSR1 signal in case that it is delivered before
        # _wait_for_children_initialized collects it

        signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGUSR1})

        self.addCleanup(
            signal.pthread_sigmask, signal.SIG_UNBLOCK, {signal.SIGUSR1}
        )

        script_path = self.write_signal_logger()

        return task.run_cmd(
            [
                script_path,
                self.temp_directory,
                str(os.getpid()),
                *options,
            ]
        )

    def run_success_failure(self, task):
        """Run success/failure script in task."""
        script_path = self.write_success_failure()

        return task.run_cmd([script_path])

    @staticmethod
    def pid_exist_not_zombie(pid):
        """Return True if the pid exist and is not a zombie process."""
        try:
            status = psutil.Process(pid).status()
            return status != "zombie"
        except psutil.NoSuchProcess:
            return False

    def assert_processes_are_terminated(self, expected_number_files_to_check):
        """
        Assert that the PIDs in self.temp_directory don't exist or a zombie.

        :param expected_number_files_to_check: number of expected processes
           (PID files) to check.
        """
        start_time = time.time()
        time_out = 2

        files_to_check = list(self.temp_directory.glob("*.log"))
        number_of_files_to_check = len(files_to_check)

        self.assertEqual(
            expected_number_files_to_check,
            number_of_files_to_check,
            f"Expected {expected_number_files_to_check} "
            f"processes found {number_of_files_to_check}",
        )

        while True:
            alive_pids = False
            elapsed = time.time() - start_time

            if elapsed > time_out:  # pragma: no cover
                # "no cover" because if the test pass this branch
                # is not used. Not having it the Test would be waiting
                # endless for the processes to disappear
                break

            for file in files_to_check:
                pid = int(file.stem)

                if self.pid_exist_not_zombie(pid):  # pragma: no cover
                    # Usually not triggered: the processes are killed very fast
                    # hence the no cover statement.
                    alive_pids = True
                else:  # pragma: no cover
                    # Usually triggered, "no cover" needed because of the issue
                    # explained in
                    # https://github.com/nedbat/coveragepy/issues/1025
                    pass

            if alive_pids:  # pragma: no cover
                # Usually not triggered because processes are killed very fast
                # hence the no cover statement.
                pass
            else:  # pragma: no cover
                # Usually triggered, "no cover" needed because of the issue
                # explained in https://github.com/nedbat/coveragepy/issues/1025
                break  # All processes are gone

        self.assertLess(
            elapsed, time_out, "Timeout waiting for processes to disappear"
        )

    def test_execute_cmd_cancelled_killed_with_sigterm(self):
        """Execute script, kills with SIGTERM. No processes running after."""
        task = TaskWithRunCommandMixin(abort_after_aborted_calls=2)

        # Create three processes
        success = self.run_signal_logger(task, ["3"])

        # The processes were killed in Task.run_cmd (the task
        # was aborted after two aborted calls)
        self.assert_processes_are_terminated(3)

        self.assertFalse(success)

    def test_execute_cmd_cancelled_killed_with_sigkill(self):
        """Execute script, two children, killed by SIGKILL (not by SIGTERM)."""
        task = TaskWithRunCommandMixin(abort_after_aborted_calls=1)

        success = self.run_signal_logger(task, ["2", "TERM"])

        self.assert_processes_are_terminated(2)

        self.assertFalse(success)

        task.wait_cmd_zombie_process()

    def test_execute_cmd_finished_before_killpg_sigterm(self):
        """Execute script. SIGTERM is called but processes were already dead."""
        task = TaskWithRunCommandMixin(
            abort_after_aborted_calls=1,
            wait_popen_sigkill=True,
            pids_directory=self.temp_directory,
        )

        success = self.run_signal_logger(task, ["2"])

        self.assert_processes_are_terminated(2)

        self.assertFalse(success)

    def test_execute_cmd_finished_before_killpg_sigkill(self):
        """Execute script. SIGKILL is called but processes were already dead."""
        task = TaskWithRunCommandMixin(
            abort_after_aborted_calls=2,
            wait_popen_sigkill=True,
            pids_directory=self.temp_directory,
        )

        success = self.run_signal_logger(task, ["2", "TERM"])

        self.assert_processes_are_terminated(2)

        self.assertFalse(success)

    def test_command_finish_with_success(self):
        """Execute script. No abort, cmd return code is 0 (success)."""
        task = TaskWithRunCommandMixin(
            abort_after_aborted_calls=math.inf,
            wait_for_sigusr1=False,
            send_signal_to_cmd=signal.SIGUSR1,
        )

        success = self.run_success_failure(task)

        self.assertTrue(success)

    def test_command_finish_with_failure(self):
        """Execute script. No abort, cmd return code is 0 (success)."""
        task = TaskWithRunCommandMixin(
            abort_after_aborted_calls=math.inf,
            wait_for_sigusr1=False,
            send_signal_to_cmd=signal.SIGUSR2,
        )

        success = self.run_success_failure(task)

        self.assertFalse(success)

    def write_success_failure(self):
        """
        Write to self.temp_directory a script that traps SIGUSR1 and SIGUSR2.

        For SIGUSR1 exits with exitcode == 0, SIGUSR2 exits with exitcode == 1.

        :return: script path
        """
        script_file = Path(self.temp_directory) / "success-failure.sh"

        with script_file.open("w") as f:
            f.write(
                textwrap.dedent(
                    """\
                #!/bin/sh

                trap "exit 0" USR1
                trap "exit 1" USR2

                while true
                do
                    true
                    # with a sleep there is a delay until the script
                    # process the signals
                done
                """
                )
            )

        script_file.chmod(0o700)

        return script_file

    def write_signal_logger(self):
        """
        Write to self.temp_directory a script to test finishing of processes.

        :return: script path
        """
        script_file = Path(self.temp_directory) / "signal-logger.sh"

        with script_file.open("w") as f:
            f.write(
                textwrap.dedent(
                    """\
            #!/bin/sh

            # Arguments:
            # $1: output directory for the log files. Must exist
            # $2: PID of the process to send a SIGUSR1 when all children
            #     has been created
            # $3: number of generation of processes to launch
            # $4: (optional) If "TERM": traps TERM signal (logs it) and do not
            #                die

            log() {
                echo "$1" >> "$output_directory/$$.log"
            }

            output_directory="$1"
            pid_to_notify="$2"
            to_spawn="$3"
            term="$4"

            if [ "$term" = "TERM" ]
            then
                trap "log SIGTERM" TERM
            fi

            log "Started $$"

            if [ "$to_spawn" -gt 1 ]
            then
                to_spawn=$((to_spawn-1))
                cmd="$0 $output_directory $pid_to_notify $to_spawn $term"
                $cmd &
                log "$$ launched $!"
            else
                kill -s 10 $pid_to_notify
            fi

            while true
            do
                sleep 1
            done
            """
                )
            )

        script_file.chmod(0o700)

        return script_file
