#!/usr/bin/env python3

__lazy_modules__ = ["shutil", "sys", "tempfile", "tomllib"]

import contextlib
import functools
import os
import pathlib
import shutil
import subprocess
import sys
import sysconfig
import tempfile
import tomllib

try:
    from os import process_cpu_count as cpu_count
except ImportError:
    from os import cpu_count


CHECKOUT = HERE = pathlib.Path(__file__).parent

while CHECKOUT != CHECKOUT.parent:
    if (CHECKOUT / "configure").is_file():
        break
    CHECKOUT = CHECKOUT.parent
else:
    raise FileNotFoundError(
        "Unable to find the root of the CPython checkout by looking for 'configure'"
    )

CROSS_BUILD_DIR = CHECKOUT / "cross-build"
# Build platform can also be found via `config.guess`.
BUILD_DIR = CROSS_BUILD_DIR / sysconfig.get_config_var("BUILD_GNU_TYPE")

LOCAL_SETUP = CHECKOUT / "Modules" / "Setup.local"
LOCAL_SETUP_MARKER = (
    b"# Generated by Platforms/WASI .\n"
    b"# Required to statically build extension modules."
)

WASMTIME_VAR_NAME = "WASMTIME"
WASMTIME_HOST_RUNNER_VAR = f"{{{WASMTIME_VAR_NAME}}}"


def separator():
    """Print a separator line across the terminal width."""
    try:
        tput_output = subprocess.check_output(
            ["tput", "cols"], encoding="utf-8"
        )
    except subprocess.CalledProcessError:
        terminal_width = 80
    else:
        terminal_width = int(tput_output.strip())
    print("⎯" * terminal_width)


def log(emoji, message, *, spacing=None):
    """Print a notification with an emoji.

    If 'spacing' is None, calculate the spacing based on the number of code points
    in the emoji as terminals "eat" a space when the emoji has multiple code points.
    """
    if spacing is None:
        spacing = " " if len(emoji) == 1 else "  "
    print("".join([emoji, spacing, message]))


def updated_env(updates={}):
    """Create a new dict representing the environment to use.

    The changes made to the execution environment are printed out.
    """
    env_defaults = {}
    # https://reproducible-builds.org/docs/source-date-epoch/
    git_epoch_cmd = ["git", "log", "-1", "--pretty=%ct"]
    try:
        epoch = subprocess.check_output(
            git_epoch_cmd, encoding="utf-8"
        ).strip()
        env_defaults["SOURCE_DATE_EPOCH"] = epoch
    except subprocess.CalledProcessError:
        pass  # Might be building from a tarball.
    # This layering lets SOURCE_DATE_EPOCH from os.environ takes precedence.
    environment = env_defaults | os.environ | updates

    env_diff = {}
    for key, value in environment.items():
        if os.environ.get(key) != value:
            env_diff[key] = value

    env_vars = [
        f"\n     {key}={item}" for key, item in sorted(env_diff.items())
    ]
    log("🌎", f"Environment changes:{''.join(env_vars)}")

    return environment


def subdir(working_dir, *, clean_ok=False):
    """Decorator to change to a working directory."""

    def decorator(func):
        @functools.wraps(func)
        def wrapper(context):
            nonlocal working_dir

            if callable(working_dir):
                working_dir = working_dir(context)
            separator()
            log("📁", os.fsdecode(working_dir))
            if (
                clean_ok
                and getattr(context, "clean", False)
                and working_dir.exists()
            ):
                log("🚮", "Deleting directory (--clean)...")
                shutil.rmtree(working_dir)

            working_dir.mkdir(parents=True, exist_ok=True)

            with contextlib.chdir(working_dir):
                return func(context, working_dir)

        return wrapper

    return decorator


def call(command, *, context=None, quiet=False, **kwargs):
    """Execute a command.

    If 'quiet' is true, then redirect stdout and stderr to a temporary file.
    """
    if context is not None:
        quiet = context.quiet

    log("❯", " ".join(map(str, command)), spacing="  ")
    if not quiet:
        stdout = None
        stderr = None
    else:
        if (logdir := getattr(context, "logdir", None)) is None:
            logdir = pathlib.Path(tempfile.gettempdir())
        stdout = tempfile.NamedTemporaryFile(
            "w",
            encoding="utf-8",
            delete=False,
            dir=logdir,
            prefix="cpython-wasi-",
            suffix=".log",
        )
        stderr = subprocess.STDOUT
        log("📝", f"Logging output to {stdout.name} (--quiet)...")

    subprocess.check_call(command, **kwargs, stdout=stdout, stderr=stderr)


def build_python_path():
    """The path to the build Python binary."""
    binary = BUILD_DIR / "python"
    if not binary.is_file():
        binary = binary.with_suffix(".exe")
        if not binary.is_file():
            raise FileNotFoundError(
                f"Unable to find `python(.exe)` in {BUILD_DIR}"
            )

    return binary


def build_python_is_pydebug():
    """Find out if the build Python is a pydebug build."""
    test = "import sys, test.support; sys.exit(test.support.Py_DEBUG)"
    result = subprocess.run(
        [build_python_path(), "-c", test],
        capture_output=True,
    )
    return bool(result.returncode)


@subdir(BUILD_DIR, clean_ok=True)
def configure_build_python(context, working_dir):
    """Configure the build/host Python."""
    if LOCAL_SETUP.exists():
        if LOCAL_SETUP.read_bytes() == LOCAL_SETUP_MARKER:
            log("👍", f"{LOCAL_SETUP} exists ...")
        else:
            log("⚠️", f"{LOCAL_SETUP} exists, but has unexpected contents")
    else:
        log("📝", f"Creating {LOCAL_SETUP} ...")
        LOCAL_SETUP.write_bytes(LOCAL_SETUP_MARKER)

    configure = [os.path.relpath(CHECKOUT / "configure", working_dir)]
    if context.args:
        configure.extend(context.args)

    call(configure, context=context)


@subdir(BUILD_DIR)
def make_build_python(context, working_dir):
    """Make/build the build Python."""
    call(["make", "--jobs", str(cpu_count()), "all"], context=context)

    binary = build_python_path()
    cmd = [
        binary,
        "-c",
        "import sys; "
        "print(f'{sys.version_info.major}.{sys.version_info.minor}')",
    ]
    version = subprocess.check_output(cmd, encoding="utf-8").strip()

    log("🎉", f"{binary} {version}")


def wasi_sdk(context):
    """Find the path to the WASI SDK."""
    if wasi_sdk_path := context.wasi_sdk_path:
        if not wasi_sdk_path.exists():
            raise ValueError(
                "WASI SDK not found; "
                "download from "
                "https://github.com/WebAssembly/wasi-sdk and/or "
                "specify via $WASI_SDK_PATH or --wasi-sdk"
            )
        return wasi_sdk_path

    with (HERE / "config.toml").open("rb") as file:
        config = tomllib.load(file)
    wasi_sdk_version = config["targets"]["wasi-sdk"]

    if wasi_sdk_path_env_var := os.environ.get("WASI_SDK_PATH"):
        wasi_sdk_path = pathlib.Path(wasi_sdk_path_env_var)
        if not wasi_sdk_path.exists():
            raise ValueError(
                f"WASI SDK not found at $WASI_SDK_PATH ({wasi_sdk_path})"
            )
    else:
        opt_path = pathlib.Path("/opt")
        # WASI SDK versions have a ``.0`` suffix, but it's a constant; the WASI SDK team
        # has said they don't plan to ever do a point release and all of their Git tags
        # lack the ``.0`` suffix.
        # Starting with WASI SDK 23, the tarballs went from containing a directory named
        # ``wasi-sdk-{WASI_SDK_VERSION}.0`` to e.g.
        # ``wasi-sdk-{WASI_SDK_VERSION}.0-x86_64-linux``.
        potential_sdks = [
            path
            for path in opt_path.glob(f"wasi-sdk-{wasi_sdk_version}.0*")
            if path.is_dir()
        ]
        if len(potential_sdks) == 1:
            wasi_sdk_path = potential_sdks[0]
        elif (default_path := opt_path / "wasi-sdk").is_dir():
            wasi_sdk_path = default_path

    # Starting with WASI SDK 25, a VERSION file is included in the root
    # of the SDK directory that we can read to warn folks when they are using
    # an unsupported version.
    if wasi_sdk_path and (version_file := wasi_sdk_path / "VERSION").is_file():
        version_details = version_file.read_text(encoding="utf-8")
        found_version = version_details.splitlines()[0]
        # Make sure there's a trailing dot to avoid false positives if somehow the
        # supported version is a prefix of the found version (e.g. `25` and `2567`).
        if not found_version.startswith(f"{wasi_sdk_version}."):
            major_version = found_version.partition(".")[0]
            log(
                "⚠️",
                f" Found WASI SDK {major_version}, "
                f"but WASI SDK {wasi_sdk_version} is the supported version",
            )

    # Cache the result.
    context.wasi_sdk_path = wasi_sdk_path
    return wasi_sdk_path


def wasi_sdk_env(context):
    """Calculate environment variables for building with wasi-sdk."""
    wasi_sdk_path = wasi_sdk(context)
    sysroot = wasi_sdk_path / "share" / "wasi-sysroot"
    env = {
        "CC": "clang",
        "CPP": "clang-cpp",
        "CXX": "clang++",
        "AR": "llvm-ar",
        "RANLIB": "ranlib",
    }

    for env_var, binary_name in list(env.items()):
        env[env_var] = os.fsdecode(wasi_sdk_path / "bin" / binary_name)

    if not wasi_sdk_path.name.startswith("wasi-sdk"):
        for compiler in ["CC", "CPP", "CXX"]:
            env[compiler] += f" --sysroot={sysroot}"

    env["PKG_CONFIG_PATH"] = ""
    env["PKG_CONFIG_LIBDIR"] = os.pathsep.join(
        map(
            os.fsdecode,
            [sysroot / "lib" / "pkgconfig", sysroot / "share" / "pkgconfig"],
        )
    )
    env["PKG_CONFIG_SYSROOT_DIR"] = os.fsdecode(sysroot)

    env["WASI_SDK_PATH"] = os.fsdecode(wasi_sdk_path)
    env["WASI_SYSROOT"] = os.fsdecode(sysroot)

    env["PATH"] = os.pathsep.join([
        os.fsdecode(wasi_sdk_path / "bin"),
        os.environ["PATH"],
    ])

    return env


def host_triple(context):
    """Determine the target triple for the WASI host build."""
    if context.host_triple:
        return context.host_triple

    with (HERE / "config.toml").open("rb") as file:
        config = tomllib.load(file)

    # Cache the result.
    context.host_triple = config["targets"]["host-triple"]
    return context.host_triple


@subdir(lambda context: CROSS_BUILD_DIR / host_triple(context), clean_ok=True)
def configure_wasi_python(context, working_dir):
    """Configure the WASI/host build."""
    config_site = os.fsdecode(HERE / "config.site-wasm32-wasi")

    wasi_build_dir = working_dir.relative_to(CHECKOUT)

    args = {
        "WASMTIME": "wasmtime",
        "ARGV0": f"/{wasi_build_dir}/python.wasm",
        "CHECKOUT": os.fsdecode(CHECKOUT),
        "WASMTIME_CONFIG_PATH": os.fsdecode(HERE / "wasmtime.toml"),
    }
    # Check dynamically for wasmtime in case it was specified manually via
    # `--host-runner`.
    if "{WASMTIME}" in context.host_runner:
        if wasmtime := shutil.which("wasmtime"):
            args["WASMTIME"] = wasmtime
        else:
            raise FileNotFoundError(
                "wasmtime not found; download from "
                "https://github.com/bytecodealliance/wasmtime"
            )
    host_runner = context.host_runner.format_map(args)
    env_additions = {"CONFIG_SITE": config_site, "HOSTRUNNER": host_runner}
    build_python = os.fsdecode(build_python_path())
    # The path to `configure` MUST be relative, else `python.wasm` is unable
    # to find the stdlib due to Python not recognizing that it's being
    # executed from within a checkout.
    configure = [
        os.path.relpath(CHECKOUT / "configure", working_dir),
        f"--host={host_triple(context)}",
        f"--build={BUILD_DIR.name}",
        f"--with-build-python={build_python}",
    ]
    if build_python_is_pydebug():
        configure.append("--with-pydebug")
    if context.args:
        configure.extend(context.args)
    call(
        configure,
        env=updated_env(env_additions | wasi_sdk_env(context)),
        context=context,
    )

    python_wasm = working_dir / "python.wasm"
    exec_script = working_dir / "python.sh"
    with exec_script.open("w", encoding="utf-8") as file:
        file.write(f'#!/bin/sh\nexec {host_runner} {python_wasm} "$@"\n')
    exec_script.chmod(0o755)
    log("🏃", f"Created {exec_script} (--host-runner)... ")
    sys.stdout.flush()


@subdir(lambda context: CROSS_BUILD_DIR / host_triple(context))
def make_wasi_python(context, working_dir):
    """Run `make` for the WASI/host build."""
    call(
        ["make", "--jobs", str(cpu_count()), "all"],
        env=updated_env(),
        context=context,
    )

    exec_script = working_dir / "python.sh"
    call([exec_script, "--version"], quiet=False)
    log(
        "🎉",
        f"Use `{exec_script.relative_to(pathlib.Path().absolute())}` "
        "to run CPython w/ the WASI host specified by --host-runner",
    )


def clean_contents(context):
    """Delete all files created by this script."""
    if CROSS_BUILD_DIR.exists():
        log("🧹", f"Deleting {CROSS_BUILD_DIR} ...")
        shutil.rmtree(CROSS_BUILD_DIR)

    if LOCAL_SETUP.exists():
        if LOCAL_SETUP.read_bytes() == LOCAL_SETUP_MARKER:
            log("🧹", f"Deleting generated {LOCAL_SETUP} ...")