#!/usr/bin/env python3
# Copyright 2024 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from contextlib import ExitStack
from textwrap import dedent
import os
import platform
import subprocess
import tempfile
import typing as T
import shlex
import sys
import unittest

_SCRIPT_DIR = os.path.realpath(os.path.dirname(__file__))
_JOBSERVER_POOL_SCRIPT = os.path.join(_SCRIPT_DIR, "jobserver_pool.py")
_JOBSERVER_TEST_HELPER_SCRIPT = os.path.join(_SCRIPT_DIR, "jobserver_test_helper.py")

_PLATFORM_IS_WINDOWS = platform.system() == "Windows"

# Set this to True to debug command invocations.
_DEBUG = False

default_env = dict(os.environ)
default_env.pop("NINJA_STATUS", None)
default_env.pop("MAKEFLAGS", None)
default_env["TERM"] = "dumb"
NINJA_PATH = os.path.abspath("./ninja")


class BuildDir:
    def __init__(self, build_ninja: str):
        self.build_ninja = dedent(build_ninja)
        self.d: T.Optional[tempfile.TemporaryDirectory] = None

    def __enter__(self):
        self.d = tempfile.TemporaryDirectory()
        with open(os.path.join(self.d.name, "build.ninja"), "w") as f:
            f.write(self.build_ninja)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.d.cleanup()

    @property
    def path(self) -> str:
        assert self.d
        return self.d.name

    def run(
        self,
        cmd_flags: T.Sequence[str] = [],
        env: T.Dict[str, str] = default_env,
    ) -> None:
        """Run a command, raise exception on error. Do not capture outputs."""
        ret = subprocess.run(cmd_flags, env=env)
        ret.check_returncode()

    def ninja_run(
        self,
        ninja_args: T.List[str],
        prefix_args: T.List[str] = [],
        extra_env: T.Dict[str, str] = {},
    ) -> "subprocess.CompletedProcess[str]":
        ret = self.ninja_spawn(
            ninja_args,
            prefix_args=prefix_args,
            extra_env=extra_env,
            capture_output=False,
        )
        ret.check_returncode()
        return ret

    def ninja_clean(self) -> None:
        self.ninja_run(["-t", "clean"])

    def ninja_spawn(
        self,
        ninja_args: T.List[str],
        prefix_args: T.List[str] = [],
        extra_env: T.Dict[str, str] = {},
        capture_output: bool = True,
    ) -> "subprocess.CompletedProcess[str]":
        """Run Ninja command and capture outputs."""
        cmd_args = prefix_args + [NINJA_PATH, "-C", self.path] + ninja_args
        if _DEBUG:
            cmd_str = " ".join(shlex.quote(c) for c in cmd_args)
            print(f"CMD [{cmd_str}]", file=sys.stderr)
        return subprocess.run(
            cmd_args,
            text=True,
            stdout=subprocess.PIPE if capture_output else None,
            stderr=subprocess.PIPE if capture_output else None,
            env={**default_env, **extra_env},
        )

    def ninja_popen(
        self,
        ninja_args: T.List[str],
        prefix_args: T.List[str] = [],
        extra_env: T.Dict[str, str] = {},
        capture_output: bool = True,
    ) -> subprocess.Popen[str]:
        """Start Ninja command and return the Popen object (i.e. without blocking)."""
        cmd_args = prefix_args + [NINJA_PATH, "-C", self.path] + ninja_args
        if _DEBUG:
            cmd_str = " ".join(shlex.quote(c) for c in cmd_args)
            print(f"CMD [{cmd_str}]", file=sys.stderr)
        return subprocess.Popen(
            cmd_args,
            text=True,
            stdout=subprocess.PIPE if capture_output else None,
            stderr=subprocess.PIPE if capture_output else None,
            env={**default_env, **extra_env},
        )


def span_output_file(span_n: int) -> str:
    return "out%02d" % span_n


def generate_build_plan(command_count: int) -> str:
    """Generate a Ninja build plan for |command_count| parallel tasks.

    Each task calls the test helper script which waits for 50ms
    then writes its own start and end time to its output file.
    """
    result = f"""
rule span
    command = {sys.executable} -S {_JOBSERVER_TEST_HELPER_SCRIPT} --duration-ms=50 $out

"""

    for n in range(command_count):
        result += "build %s: span\n" % span_output_file(n)

    result += "build all: phony %s\n" % " ".join(
        [span_output_file(n) for n in range(command_count)]
    )
    return result

def generate_double_build_plan(
        wait_1_ms: int = 500,
        wait_2_ms: int = 500,
        offset: int = 0,
    ) -> str:

    """Generate a Ninja build plan that invokes the test helper script in two parallel commands, each one with its own duration argument

    - `wait_1_ms` and `wait_2_ms` control task durations.
    - `offset`: start index of the output files generated by this plan, allows multiple plans without conflicting filenames.
    """
    result = f"""
rule short
    command = {sys.executable} -S {_JOBSERVER_TEST_HELPER_SCRIPT} --duration-ms={wait_1_ms} $out

rule long
    command = {sys.executable} -S {_JOBSERVER_TEST_HELPER_SCRIPT} --duration-ms={wait_2_ms} $out
"""

    result += "build %s: long\n" % span_output_file(offset)
    result += "build %s: short\n" % span_output_file(offset + 1)

    result += "build all: phony %s\n" % " ".join(
        [span_output_file(n + offset) for n in range(2)]
    )

    return result

def get_spans(build_dir: str, command_count: int) -> T.List[T.Tuple[int, int]]:
    """
    Get span outputs for a single build directory.

    - `build_dir`: Directory containing span output files.
    - `command_count`: Number of commands executed in the build.

    Returns a list of (start, end) tuples parsed from the span output files, sorted by start time.
    """
    return get_spans_for_multiple_builds([build_dir], [command_count])

def get_spans_for_multiple_builds(build_dirs: T.List[str], command_counts: T.List[int]) -> T.List[T.Tuple[int, int]]:
    """
    Get span outputs across multiple build directories. Assumes that the output files
    in each build directory are named sequentially according to the `span_output_file` function,
    and that the total number of output files across all build directories is equal to the sum of `command_counts`.

    - `build_dirs`: Build directories containing span output files.
    - `command_counts`: Number of commands executed in each build.

    Returns a list of (start, end) span tuples in increasing order by start time.
    """
    spans: T.List[T.Tuple[int, int]] = []
    offset:int = 0
    for b_path, command_count in zip(build_dirs, command_counts):
        for n in range(command_count):
            out_file = os.path.join(b_path, span_output_file(n + offset))
            with open(out_file, "rb") as f:
                content = f.read().decode("utf-8")
            lines = content.splitlines()
            assert len(lines) == 2, f"Unexpected output file content: [{content}]"
            spans.append((int(lines[0]), int(lines[1])))

        # Update the offset for the next build directory to ensure unique output file names.
        offset += command_count
    return sorted(spans, key=lambda x: x[0])

def compute_max_overlapped_spans(build_dir: str, command_count: int) -> int:
    """Compute the maximum number of overlapped spanned tasks.

    This reads the output files from |build_dir| and look at their start and end times
    to compute the maximum number of tasks that were run in parallel.
    """
    # Read the output files.
    if command_count < 2:
        return 0

    spans: T.List[T.Tuple[int, int]] = get_spans(build_dir, command_count)
    # Stupid but simple, for each span, count the number of other spans that overlap it.
    max_overlaps = 1
    for n in range(command_count):
        cur_start, cur_end = spans[n]
        cur_overlaps = 1
        for m in range(command_count):
            other_start, other_end = spans[m]
            if n != m and other_end > cur_start and other_start < cur_end:
                cur_overlaps += 1

        if cur_overlaps > max_overlaps:
            max_overlaps = cur_overlaps

    return max_overlaps

class JobserverTest(unittest.TestCase):

    def test_no_jobserver_client(self):
        task_count = 4
        build_plan = generate_build_plan(task_count)
        with BuildDir(build_plan) as b:
            output = b.run([NINJA_PATH, "-C", b.path, f"-j{task_count}", "all"])

            max_overlaps = compute_max_overlapped_spans(b.path, task_count)
            self.assertEqual(max_overlaps, task_count)

            b.ninja_clean()
            output = b.run([NINJA_PATH, "-C", b.path, "-j1", "all"])

            max_overlaps = compute_max_overlapped_spans(b.path, task_count)
            self.assertEqual(max_overlaps, 1)

    def _run_client_test(self, jobserver_args: T.List[str]) -> None:
        task_count = 4
        build_plan = generate_build_plan(task_count)
        with BuildDir(build_plan) as b:
            # First, run the full tasks with with {task_count} tokens, this should allow all
            # tasks to run in parallel.
            ret = b.ninja_run(
                ninja_args=["all"],
                prefix_args=jobserver_args + [f"--jobs={task_count}"],
            )
            max_overlaps = compute_max_overlapped_spans(b.path, task_count)
            self.assertEqual(max_overlaps, task_count)

            # Second, use 2 tokens only, and verify that this was enforced by Ninja.
            b.ninja_clean()
            b.ninja_run(
                ["all"],
                prefix_args=jobserver_args + ["--jobs=2"],
            )
            max_overlaps = compute_max_overlapped_spans(b.path, task_count)
            self.assertEqual(max_overlaps, 2)

            # Third, verify that --jobs=1 serializes all tasks.
            b.ninja_clean()
            b.ninja_run(
                ["all"],
                prefix_args=jobserver_args + ["--jobs=1"],
            )
            max_overlaps = compute_max_overlapped_spans(b.path, task_count)
            self.assertEqual(max_overlaps, 1)

            # Finally, verify that -j1 overrides the pool.
            b.ninja_clean()
            b.ninja_run(
                ["-j1", "all"],
                prefix_args=jobserver_args + [f"--jobs={task_count}"],
            )
            max_overlaps = compute_max_overlapped_spans(b.path, task_count)
            self.assertEqual(max_overlaps, 1)

            # On Linux, use taskset to limit the number of available cores to 1
            # and verify that the jobserver overrides the default Ninja parallelism
            # and that {task_count} tasks are still spawned in parallel.
            if platform.system() == "Linux":
                # First, run without a jobserver, with a single CPU, Ninja will
                # use a parallelism of 2 in this case (GuessParallelism() in ninja.cc)
                b.ninja_clean()
                b.ninja_run(
                    ["all"],
                    prefix_args=["taskset", "-c", "0"],
                )
                max_overlaps = compute_max_overlapped_spans(b.path, task_count)
                self.assertEqual(max_overlaps, 2)

                # Now with a jobserver with {task_count} tasks.
                b.ninja_clean()
                b.ninja_run(
                    ["all"],
                    prefix_args=jobserver_args
                    + [f"--jobs={task_count}"]
                    + ["taskset", "-c", "0"],
                )
                max_overlaps = compute_max_overlapped_spans(b.path, task_count)
                self.assertEqual(max_overlaps, task_count)

    @unittest.skipIf(_PLATFORM_IS_WINDOWS, "These test methods do not work on Windows")
    def test_jobserver_client_with_posix_fifo(self):
        self._run_client_test([sys.executable, "-S", _JOBSERVER_POOL_SCRIPT])

    @unittest.skipIf(_PLATFORM_IS_WINDOWS, "These test methods do not work on Windows")
    def test_jobserver_client_with_posix_pipe(self):
        # Verify that setting up a --pipe server does not make Ninja exit with an error.
        # Instead, a warning is printed.
        task_count = 4
        build_plan = generate_build_plan(task_count)
        with BuildDir(build_plan) as b:

            prefix_args = [
                sys.executable,
                "-S",
                _JOBSERVER_POOL_SCRIPT,
                "--pipe",
                f"--jobs={task_count}",
            ]

            def run_ninja_with_jobserver_pipe(args):
                ret = b.ninja_spawn(args, prefix_args=prefix_args)
                ret.check_returncode()
                return ret.stdout, ret.stderr

            output, error = run_ninja_with_jobserver_pipe(["all"])
            if _DEBUG:
                print(f"OUTPUT [{output}]\nERROR [{error}]\n", file=sys.stderr)
            self.assertTrue(error.find("Pipe-based protocol is not supported!") >= 0)

            max_overlaps = compute_max_overlapped_spans(b.path, task_count)
            self.assertEqual(max_overlaps, task_count)

            # Using an explicit -j<N> ignores the jobserver pool.
            b.ninja_clean()
            output, error = run_ninja_with_jobserver_pipe(["-j1", "all"])
            if _DEBUG:
                print(f"OUTPUT [{output}]\nERROR [{error}]\n", file=sys.stderr)
            self.assertFalse(error.find("Pipe-based protocol is not supported!") >= 0)

            max_overlaps = compute_max_overlapped_spans(b.path, task_count)
            self.assertEqual(max_overlaps, 1)

    @unittest.skipIf(_PLATFORM_IS_WINDOWS, "These test methods do not work on Windows")
    def test_jobserver_client_with_posix_fifo_token_efficiency(self):
        # Due to implicit slots, the max number of parallel tasks is task_count + num_clients - 1
        task_count = 2
        client_tasks = 2

        with ExitStack() as stack:
            tmp_dir = tempfile.mkdtemp()
            stack.callback(os.rmdir, tmp_dir)

            fifo_path = os.path.join(tmp_dir, "jobserver_fifo_test")

            # Create FIFO for Jobserver - see notes about implicit slots above
            os.mkfifo(fifo_path)  # type: ignore
            stack.callback(lambda: os.path.exists(fifo_path) and os.remove(fifo_path))

            read_fd = os.open(fifo_path, os.O_RDONLY | os.O_NONBLOCK)
            stack.callback(os.close, read_fd)

            write_fd = os.open(fifo_path, os.O_WRONLY | os.O_NONBLOCK)
            stack.callback(os.close, write_fd)

            os.write(write_fd, (task_count - 1) * b"x")

            env = dict(os.environ)
            env["MAKEFLAGS"] = f" -j{task_count} --jobserver-auth=fifo:" + fifo_path

            # First build plan has 2 short tasks, second build plan has 2 long tasks
            build_plan_1 = generate_double_build_plan(wait_1_ms=1000, wait_2_ms=1000, offset=0)
            build_plan_2 = generate_double_build_plan(wait_1_ms=2000, wait_2_ms=2000, offset=client_tasks)

            with BuildDir(build_plan_1) as b1, BuildDir(build_plan_2) as b2:
                # Start all builds in parallel, with the same jobserver environment.
                p1 = b1.ninja_popen(["all"], extra_env=env)
                p2 = b2.ninja_popen(["all"], extra_env=env)

                # Wait for all builds to finish and verify they all succeeded.
                p1.communicate()
                p2.communicate()

                self.assertEqual(p1.returncode, 0, "Build 1 failed")
                self.assertEqual(p2.returncode, 0, "Build 2 failed")

                """
                As noted above, with implicit slots, the max number of parallel tasks
                is task_count + num_clients - 1, which in this case is 3.
                The expected behavior is that 3 tasks run in parallel at the beginning as follows:

                |---A---|           (Client 1, 2000ms)
                |---B---|           (Client 1, 2000ms)
                |-------C-------|   (Client 2, 4000ms)

                Once first 2 of 3 tasks finishes (i.e. the short sleeps), the 4th task can start,
                and this should start before the long task finishes.

                |---A---|                    (Client 1, 2000ms)
                |---B---|                    (Client 1, 2000ms)
                |-------C-------|            (Client 2, 4000ms)
                           |-------D-------| (Client 2, 4000ms)

                Make sure D starts after A and B finishes, but before C finishes.
                """

                spans = get_spans_for_multiple_builds(
                    [b1.path, b2.path], [client_tasks, client_tasks]
                )
                spans.sort(key=lambda s: s[0])

                initial_tasks = spans[: task_count + client_tasks - 1]
                last_task = spans[-1]

                earliest_end = min(task[1] for task in initial_tasks)
                latest_end = max(task[1] for task in initial_tasks)

                # Assert last task starts after at least one initial task finishes,
                # but before all initial tasks finish.
                self.assertTrue(
                    last_task[0] >= earliest_end and last_task[0] < latest_end
                )

    def _test_MAKEFLAGS_value(
        self, ninja_args: T.List[str] = [], prefix_args: T.List[str] = []
    ):
        build_plan = r"""
rule print
    command = echo MAKEFLAGS="[$$MAKEFLAGS]"

build all: print
"""
        with BuildDir(build_plan) as b:
            ret = b.ninja_spawn(
                ninja_args + ["--quiet", "all"], prefix_args=prefix_args
            )
            self.assertEqual(ret.returncode, 0)
            output = ret.stdout.strip()
            pos = output.find("MAKEFLAGS=[")
            self.assertNotEqual(pos, -1, "Could not find MAKEFLAGS in output!")
            makeflags, sep, _ = output[pos + len("MAKEFLAGS=[") :].partition("]")
            self.assertEqual(sep, "]", "Missing ] in output!: " + output)
            self.assertTrue(
                "--jobserver-auth=" in makeflags,
                f"Missing --jobserver-auth from MAKEFLAGS [{makeflags}]\nSTDOUT [{ret.stdout}]\nSTDERR [{ret.stderr}]",
            )

    def test_client_passes_MAKEFLAGS(self):
        self._test_MAKEFLAGS_value(
            prefix_args=[sys.executable, "-S", _JOBSERVER_POOL_SCRIPT]
        )


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