# ===----------------------------------------------------------------------===##
#
# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#
# ===----------------------------------------------------------------------===##

"""
Lit test format for LLVM libc tests.

This format discovers pre-built test executables in the build directory
and runs them. It extends lit's ExecutableTest format.

The lit config sets test_source_root == test_exec_root (both to the build
directory), following the pattern used by llvm/test/Unit/lit.cfg.py.

Test executables are discovered by _isTestExecutable() and run by execute().

Integration tests that require command-line arguments or environment variables
have a sidecar <executable>.params file generated by CMake.  The format is
one arg per line, a "---" separator, then one KEY=VALUE env entry per line.
"""

import os
import shlex
import sys

import lit.formats
import lit.Test
import lit.util

kIsWindows = sys.platform in ["win32", "cygwin"]


class LibcTest(lit.formats.ExecutableTest):
    """
    Test format for libc unit tests.

    Extends ExecutableTest to discover pre-built test executables in the
    build directory rather than the source directory.
    """

    def getTestsInDirectory(self, testSuite, path_in_suite, litConfig, localConfig):
        """
        Discover test executables in the build directory.

        Since test_source_root == test_exec_root (both point to build dir),
        we use getSourcePath() to find test executables.
        """
        source_path = testSuite.getSourcePath(path_in_suite)

        # Look for test executables in the build directory
        if not os.path.isdir(source_path):
            return

        # Sort for deterministic test discovery/output ordering.
        for filename in sorted(os.listdir(source_path)):
            filepath = os.path.join(source_path, filename)

            # Match our test executable pattern
            if self._isTestExecutable(filename, filepath, localConfig):
                # Create a test with the executable name
                yield lit.Test.Test(testSuite, path_in_suite + (filename,), localConfig)

    def _isTestExecutable(self, filename, filepath, localConfig):
        """
        Check if a file is a libc test executable we should run.

        Recognized patterns (all must end with .__build__, optionally followed
        by .exe on Windows):
          libc.test.src.<category>.<test_name>.__build__
          libc.test.src.<category>.<test_name>.__unit__[.<opts>...].__build__
          libc.test.src.<category>.<test_name>.__hermetic__[.<opts>...].__build__
          libc.test.include.<test_name>.__unit__[.<opts>...].__build__
          libc.test.include.<test_name>.__hermetic__[.<opts>...].__build__
          libc.test.integration.<category>.<test_name>.__build__
        """
        test_name = filename
        if kIsWindows and filename.endswith(".exe"):
            test_name = filename[: -len(".exe")]

        if not test_name.endswith(".__build__"):
            return False
        if test_name.startswith("libc.test.src."):
            pass  # Accept all src tests ending in .__build__
        elif test_name.startswith("libc.test.include."):
            if ".__unit__." not in test_name and ".__hermetic__." not in test_name:
                return False
        elif test_name.startswith("libc.test.integration."):
            pass  # Accept all integration tests ending in .__build__
        elif test_name.startswith("libc.test.shared."):
            pass  # Accept all shared tests ending in .__build__
        elif test_name.startswith("libc.test.utils."):
            pass  # Accept all utils tests ending in .__build__
        else:
            return False
        if not os.path.isfile(filepath):
            return False
        # GPU binaries are not host-executable but run via an emulator, so ignore X_OK if emulator is set.
        if (
            not kIsWindows
            and not os.access(filepath, os.X_OK)
            and not getattr(localConfig, "libc_crosscompiling_emulator", None)
        ):
            return False
        return True

    def _getParamsPath(self, test_path):
        params_path = test_path + ".params"
        if os.path.isfile(params_path):
            return params_path

        root, ext = os.path.splitext(test_path)
        if ext.lower() == ".exe":
            params_path = root + ".params"
            if os.path.isfile(params_path):
                return params_path

        return None

    def execute(self, test, litConfig):
        """
        Execute a test by running the test executable.

        Runs from the executable's directory so relative paths (like
        testdata/test.txt) work correctly.

        If a sidecar <executable>.params file exists, it supplies the
        command-line arguments and environment variables for the test.

        Honors test.config.maxIndividualTestTime (set via --timeout) to
        kill tests that exceed the per-test time limit.
        """

        test_path = test.getSourcePath()
        exec_dir = os.path.dirname(test_path)

        # Read optional sidecar .params file generated by CMake for tests that
        # need specific args/env (e.g. integration tests with ARGS/ENV).
        # Format: one arg per line, "---" separator, then KEY=VALUE env lines.
        loader_args = []
        test_args = []
        extra_env = {}
        params_path = self._getParamsPath(test_path)
        if params_path:
            with open(params_path) as f:
                content = f.read()
            sections = content.split("---\n")
            if len(sections) >= 3:
                loader_args = [l for l in sections[0].splitlines() if l]
                test_args = [l for l in sections[1].splitlines() if l]
                env_section = sections[2]
            else:
                loader_args = []
                test_args = [l for l in sections[0].splitlines() if l]
                env_section = sections[1] if len(sections) > 1 else ""

            for line in env_section.splitlines():
                if "=" in line:
                    k, _, v = line.partition("=")
                    extra_env[k] = v

        # Build the environment: inherit the current process environment, then
        # set PWD to exec_dir so getenv("PWD") matches getcwd(), then overlay
        # any test-specific variables from the .params file.
        env = dict(os.environ)
        env["PWD"] = exec_dir
        env.update(extra_env)

        timeout = test.config.maxIndividualTestTime

        test_cmd_template = getattr(test.config, "libc_test_cmd", "")
        if test_cmd_template:
            if "@BINARY@" in test_cmd_template:
                # Insert loader_args before the binary, and test_args after.
                prefix, _, suffix = test_cmd_template.partition("@BINARY@")
                cmd_args = (
                    shlex.split(prefix)
                    + loader_args
                    + [test_path]
                    + shlex.split(suffix)
                    + test_args
                )
            else:
                # Fallback to appending the binary path if @BINARY@ placeholder is missing.
                cmd_args = (
                    shlex.split(test_cmd_template)
                    + loader_args
                    + [test_path]
                    + test_args
                )
            if not cmd_args:
                cmd_args = [test_path]
        else:
            cmd_args = [test_path] + test_args

        try:
            out, err, exit_code = lit.util.executeCommand(
                cmd_args, cwd=exec_dir, env=env, timeout=timeout
            )
        except lit.util.ExecuteCommandTimeoutException as e:
            return (
                lit.Test.TIMEOUT,
                f"{e.out}\n--\n" f"Reached timeout of {timeout} seconds",
            )

        if not exit_code:
            return lit.Test.PASS, ""

        return lit.Test.FAIL, out + err