#!/usr/bin/env python3

import argparse
import contextlib
import functools
import hashlib
import json
import os
import shutil
import subprocess
import sys
import sysconfig
import tempfile
from pathlib import Path
from textwrap import dedent
from urllib.request import urlopen

import tomllib

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


EMSCRIPTEN_DIR = Path(__file__).parent
CHECKOUT = EMSCRIPTEN_DIR.parent.parent.parent
CONFIG_FILE = EMSCRIPTEN_DIR / "config.toml"

DEFAULT_CROSS_BUILD_DIR = CHECKOUT / "cross-build"
HOST_TRIPLE = "wasm32-emscripten"


@functools.cache
def load_config_toml():
    with CONFIG_FILE.open("rb") as file:
        return tomllib.load(file)


@functools.cache
def required_emscripten_version():
    return load_config_toml()["emscripten-version"]


@functools.cache
def emsdk_cache_root(emsdk_cache):
    required_version = required_emscripten_version()
    return Path(emsdk_cache).absolute() / required_version


@functools.cache
def emsdk_activate_path(emsdk_cache):
    return emsdk_cache_root(emsdk_cache) / "emsdk/emsdk_env.sh"


def get_build_paths(cross_build_dir=None, emsdk_cache=None):
    """Compute all build paths from the given cross-build directory."""
    if cross_build_dir is None:
        cross_build_dir = DEFAULT_CROSS_BUILD_DIR
    cross_build_dir = Path(cross_build_dir).absolute()
    host_triple_dir = cross_build_dir / HOST_TRIPLE
    prefix_dir = host_triple_dir / "prefix"
    if emsdk_cache:
        prefix_dir = emsdk_cache_root(emsdk_cache) / "prefix"

    return {
        "cross_build_dir": cross_build_dir,
        "native_build_dir": cross_build_dir / "build",
        "host_triple_dir": host_triple_dir,
        "host_build_dir": host_triple_dir / "build",
        "host_dir": host_triple_dir / "build" / "python",
        "prefix_dir": prefix_dir,
    }


LOCAL_SETUP = CHECKOUT / "Modules" / "Setup.local"
LOCAL_SETUP_MARKER = b"# Generated by Tools/wasm/emscripten.py\n"


def validate_emsdk_version(emsdk_cache):
    """Validate that the emsdk cache contains the required emscripten version."""
    required_version = required_emscripten_version()
    emsdk_env = emsdk_activate_path(emsdk_cache)
    if not emsdk_env.is_file():
        print(
            f"Required emscripten version {required_version} not found in {emsdk_cache}",
            file=sys.stderr,
        )
        sys.exit(1)
    print(f"✅ Emscripten version {required_version} found in {emsdk_cache}")


def parse_env(text):
    result = {}
    for line in text.splitlines():
        key, val = line.split("=", 1)
        result[key] = val
    return result


@functools.cache
def get_emsdk_environ(emsdk_cache):
    """Returns os.environ updated by sourcing emsdk_env.sh"""
    if not emsdk_cache:
        return os.environ
    env_text = subprocess.check_output(
        [
            "bash",
            "-c",
            f"EMSDK_QUIET=1 source {emsdk_activate_path(emsdk_cache)} && env",
        ],
        text=True,
    )
    return parse_env(env_text)


def updated_env(updates, emsdk_cache):
    """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 | get_emsdk_environ(emsdk_cache) | updates
    env_diff = {}
    for key, value in environment.items():
        if os.environ.get(key) != value:
            env_diff[key] = value

    print("🌎 Environment changes:")
    for key in sorted(env_diff.keys()):
        print(f"  {key}={env_diff[key]}")

    return environment


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

    path_key is a key into context.build_paths, used to resolve the working
    directory at call time.
    """

    def decorator(func):
        @functools.wraps(func)
        def wrapper(context):
            working_dir = context.build_paths[path_key]
            try:
                tput_output = subprocess.check_output(
                    ["tput", "cols"], encoding="utf-8"
                )
                terminal_width = int(tput_output.strip())
            except subprocess.CalledProcessError:
                terminal_width = 80
            print("⎯" * terminal_width)
            print("📁", working_dir)
            if (
                clean_ok
                and getattr(context, "clean", False)
                and working_dir.exists()
            ):
                print("🚮 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, *, quiet, **kwargs):
    """Execute a command.

    If 'quiet' is true, then redirect stdout and stderr to a temporary file.
    """
    print("❯", " ".join(map(str, command)))
    if not quiet:
        stdout = None
        stderr = None
    else:
        stdout = tempfile.NamedTemporaryFile(
            "w",
            encoding="utf-8",
            delete=False,
            prefix="cpython-emscripten-",
            suffix=".log",
        )
        stderr = subprocess.STDOUT
        print(f"📝 Logging output to {stdout.name} (--quiet)...")

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


def build_platform():
    """The name of the build/host platform."""
    # Can also be found via `config.guess`.`
    return sysconfig.get_config_var("BUILD_GNU_TYPE")


def build_python_path(context):
    """The path to the build Python binary."""
    native_build_dir = context.build_paths["native_build_dir"]
    binary = native_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 {native_build_dir}"
            )

    return binary


def install_emscripten(context):
    emsdk_cache = context.emsdk_cache
    if emsdk_cache is None:
        print("install-emscripten requires --emsdk-cache", file=sys.stderr)
        sys.exit(1)
    version = required_emscripten_version()
    emsdk_target = emsdk_cache_root(emsdk_cache) / "emsdk"
    if emsdk_target.exists():
        if not context.quiet:
            print(f"Emscripten version {version} already installed")
        return
    if not context.quiet:
        print(f"Installing emscripten version {version}")
    emsdk_target.mkdir(parents=True)
    call(
        [
            "git",
            "clone",
            "https://github.com/emscripten-core/emsdk.git",
            emsdk_target,
        ],
        quiet=context.quiet,
    )
    call([emsdk_target / "emsdk", "install", version], quiet=context.quiet)
    call([emsdk_target / "emsdk", "activate", version], quiet=context.quiet)
    if not context.quiet:
        print(f"Installed emscripten version {version}")


@subdir("native_build_dir", clean_ok=True)
def configure_build_python(context, working_dir):
    """Configure the build/host Python."""
    if LOCAL_SETUP.exists():
        print(f"👍 {LOCAL_SETUP} exists ...")
    else:
        print(f"📝 Touching {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, quiet=context.quiet)


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

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

    print(f"🎉 {binary} {version}")


def check_shasum(file: str, expected_shasum: str):
    with open(file, "rb") as f:
        digest = hashlib.file_digest(f, "sha256")
    if digest.hexdigest() != expected_shasum:
        raise RuntimeError(f"Unexpected shasum for {file}")


def download_and_unpack(working_dir: Path, url: str, expected_shasum: str):
    with tempfile.NamedTemporaryFile(
        suffix=".tar.gz", delete_on_close=False
    ) as tmp_file:
        with urlopen(url) as response:
            shutil.copyfileobj(response, tmp_file)
        tmp_file.close()
        check_shasum(tmp_file.name, expected_shasum)
        shutil.unpack_archive(tmp_file.name, working_dir)


def should_build_library(prefix, name, config, quiet):
    cached_config = prefix / (name + ".json")
    if not cached_config.exists():
        if not quiet:
            print(
                f"No cached build of {name} version {config['version']} found, building"
            )
        return True

    try:
        with cached_config.open("rb") as f:
            cached_config = json.load(f)
    except json.JSONDecodeError:
        if not quiet:
            print(f"Cached data for {name} invalid, rebuilding")
        return True
    if config == cached_config:
        if not quiet:
            print(
                f"Found cached build of {name} version {config['version']}, not rebuilding"
            )
        return False

    if not quiet:
        print(
            f"Found cached build of {name} version {config['version']} but it's out of date, rebuilding"
        )
    return True


def write_library_config(prefix, name, config, quiet):
    cached_config = prefix / (name + ".json")
    with cached_config.open("w") as f:
        json.dump(config, f)
    if not quiet:
        print(f"Succeded building {name}, wrote config to {cached_config}")


@subdir("host_build_dir", clean_ok=True)
def make_emscripten_libffi(context, working_dir):
    prefix = context.build_paths["prefix_dir"]
    libffi_config = load_config_toml()["libffi"]
    if not should_build_library(
        prefix, "libffi", libffi_config, context.quiet
    ):
        return
    url = libffi_config["url"]
    version = libffi_config["version"]
    shasum = libffi_config["shasum"]
    libffi_dir = working_dir / f"libffi-{version}"
    shutil.rmtree(libffi_dir, ignore_errors=True)
    download_and_unpack(
        working_dir,
        url.format(version=version),
        shasum,
    )
    call(
        [EMSCRIPTEN_DIR / "make_libffi.sh"],
        env=updated_env({"PREFIX": prefix}, context.emsdk_cache),
        cwd=libffi_dir,
        quiet=context.quiet,
    )
    write_library_config(prefix, "libffi", libffi_config, context.quiet)


@subdir("host_build_dir", clean_ok=True)
def make_mpdec(context, working_dir):
    prefix = context.build_paths["prefix_dir"]
    mpdec_config = load_config_toml()["mpdec"]
    if not should_build_library(prefix, "mpdec", mpdec_config, context.quiet):
        return

    url = mpdec_config["url"]
    version = mpdec_config["version"]
    shasum = mpdec_config["shasum"]
    mpdec_dir = working_dir / f"mpdecimal-{version}"
    shutil.rmtree(mpdec_dir, ignore_errors=True)
    download_and_unpack(
        working_dir,
        url.format(version=version),
        shasum,
    )
    call(
        [
            "emconfigure",
            mpdec_dir / "configure",
            "CFLAGS=-fPIC",
            "--prefix",
            prefix,
            "--disable-shared",
        ],
        cwd=mpdec_dir,
        quiet=context.quiet,
        env=updated_env({}, context.emsdk_cache),
    )
    call(
        ["make", "install"],
        cwd=mpdec_dir,
        quiet=context.quiet,
    )
    write_library_config(prefix, "mpdec", mpdec_config, context.quiet)


@subdir("host_dir", clean_ok=True)
def configure_emscripten_python(context, working_dir):
    """Configure the emscripten/host build."""
    paths = context.build_paths
    config_site = os.fsdecode(EMSCRIPTEN_DIR / "config.site-wasm32-emscripten")

    emscripten_build_dir = working_dir.relative_to(CHECKOUT)

    python_build_dir = paths["native_build_dir"] / "build"
    lib_dirs = list(python_build_dir.glob("lib.*"))
    assert len(lib_dirs) == 1, (
        f"Expected a single lib.* directory in {python_build_dir}"
    )
    lib_dir = os.fsdecode(lib_dirs[0])
    pydebug = lib_dir.endswith("-pydebug")
    python_version = lib_dir.removesuffix("-pydebug").rpartition("-")[-1]
    sysconfig_data = (
        f"{emscripten_build_dir}/build/lib.emscripten-wasm32-{python_version}"
    )
    if pydebug:
        sysconfig_data += "-pydebug"

    host_runner = context.host_runner
    if node_version := os.environ.get("PYTHON_NODE_VERSION", None):
        res = subprocess.run(
            [
                "bash",
                "-c",
                f"source ~/.nvm/nvm.sh && nvm which {node_version}",
            ],
            text=True,
            capture_output=True,
        )
        host_runner = res.stdout.strip()
    pkg_config_path_dir = (paths["prefix_dir"] / "lib/pkgconfig/").resolve()
    env_additions = {
        "CONFIG_SITE": config_site,
        "HOSTRUNNER": host_runner,
        "EM_PKG_CONFIG_PATH": str(pkg_config_path_dir),
    }
    build_python = os.fsdecode(build_python_path(context))
    configure = [
        "emconfigure",
        os.path.relpath(CHECKOUT / "configure", working_dir),
        "CFLAGS=-DPY_CALL_TRAMPOLINE -sUSE_BZIP2",
        "PKG_CONFIG=pkg-config",
        f"--host={HOST_TRIPLE}",
        f"--build={build_platform()}",
        f"--with-build-python={build_python}",
        "--without-pymalloc",
        "--disable-shared",
        "--disable-ipv6",
        "--enable-big-digits=30",
        "--enable-wasm-dynamic-linking",
        f"--prefix={paths['prefix_dir']}",
    ]
    if pydebug:
        configure.append("--with-pydebug")
    if context.args:
        configure.extend(context.args)
    call(
        configure,
        env=updated_env(env_additions, context.emsdk_cache),
        quiet=context.quiet,
    )

    shutil.copy(
        EMSCRIPTEN_DIR / "node_entry.mjs", working_dir / "node_entry.mjs"
    )

    node_entry = working_dir / "node_entry.mjs"
    exec_script = working_dir / "python.sh"
    exec_script.write_text(
        dedent(
            f"""\
            #!/bin/sh

            # Macs come with FreeBSD coreutils which doesn't have the -s option
            # so feature detect and work around it.
            if which grealpath > /dev/null 2>&1; then
                # It has brew installed gnu core utils, use that
                REALPATH="grealpath -s"
            elif which realpath > /dev/null 2>&1 && realpath --version > /dev/null 2>&1 && realpath --version | grep GNU > /dev/null 2>&1; then
                # realpath points to GNU realpath so use it.
                REALPATH="realpath -s"
            else
                # Shim for macs without GNU coreutils
                abs_path () {{
                    echo "$(cd "$(dirname "$1")" || exit; pwd)/$(basename "$1")"
                }}
                REALPATH=abs_path
            fi

            # Before node 24, --experimental-wasm-jspi uses different API,
            # After node 24 JSPI is on by default.
            ARGS=$({host_runner} -e "$(cat <<"EOF"
            const major_version = Number(process.version.split(".")[0].slice(1));
            if (major_version === 24) {{
                process.stdout.write("--experimental-wasm-jspi");
            }}
            EOF
            )")

            # We compute our own path, not following symlinks and pass it in so that
            # node_entry.mjs can set sys.executable correctly.
            # Intentionally allow word splitting on NODEFLAGS.
            exec {host_runner} $NODEFLAGS $ARGS {node_entry} --this-program="$($REALPATH "$0")" "$@"
            """
        )
    )
    exec_script.chmod(0o755)
    print(f"🏃‍♀️ Created {exec_script} ... ")
    sys.stdout.flush()


@subdir("host_dir")
def make_emscripten_python(context, working_dir):
    """Run `make` for the emscripten/host build."""
    call(
        ["make", "--jobs", str(cpu_count()), "all"],
        env=updated_env({}, context.emsdk_cache),
        quiet=context.quiet,
    )

    exec_script = working_dir / "python.sh"
    subprocess.check_call([exec_script, "--version"])


def build_target(context):
    """Build one or more targets."""
    steps = []
    if context.target in {"all"}:
        steps.append(install_emscripten)
    if context.target in {"build", "all"}:
        steps.extend([
            configure_build_python,
            make_build_python,
        ])
    if context.target in {"host", "all"}:
        steps.extend([
            make_emscripten_libffi,
            make_mpdec,
            configure_emscripten_python,
            make_emscripten_python,
        ])

    for step in steps:
        step(context)


def clean_contents(context):
    """Delete all files created by this script."""
    if context.target in {"all", "build"}:
        build_dir = context.build_paths["native_build_dir"]
        if build_dir.exists():
            print(f"🧹 Deleting {build_dir} ...")
            shutil.rmtree(build_dir)

    if context.target in {"all", "host"}:
        host_triple_dir = context.build_paths["host_triple_dir"]
        if host_triple_dir.exists():
            print(f"🧹 Deleting {host_triple_dir} ...")
            shutil.rmtree(host_triple_dir)

    if LOCAL_SETUP.exists():
        with LOCAL_SETUP.open("rb") as file:
            if file.read(len(LOCAL_SETUP_MARKER)) == LOCAL_SETUP_MARKER:
                print(f"🧹 Deleting generated {LOCAL_SETUP} ...")


def main():
    default_host_runner = "node"

    parser = argparse.ArgumentParser()
    subcommands = parser.add_subparsers(dest="subcommand")
    install_emscripten_cmd = subcommands.add_parser(
        "install-emscripten",
        help="Install the appropriate version of Emscripten",
    )
    build = subcommands.add_parser("build", help="Build everything")
    build.add_argument(
        "target",
        nargs="?",
        default="all",
        choices=["all", "host", "build"],
        help=(
            "What should be built. 'build' for just the build platform, or "
            "'host' for the host platform, or 'all' for both. Defaults to 'all'."
        ),
    )

    configure_build = subcommands.add_parser(
        "configure-build-python", help="Run `configure` for the build Python"
    )
    make_mpdec_cmd = subcommands.add_parser(
        "make-mpdec",
        help="Clone mpdec repo, configure and build it for emscripten",
    )
    make_libffi_cmd = subcommands.add_parser(
        "make-libffi",
        help="Clone libffi repo, configure and build it for emscripten",
    )
    make_build = subcommands.add_parser(
        "make-build-python", help="Run `make` for the build Python"
    )
    configure_host = subcommands.add_parser(
        "configure-host",
        help="Run `configure` for the host/emscripten (pydebug builds are inferred from the build Python)",
    )
    make_host = subcommands.add_parser(
        "make-host", help="Run `make` for the host/emscripten"
    )
    clean = subcommands.add_parser(
        "clean", help="Delete files and directories created by this script"
    )
    clean.add_argument(
        "target",
        nargs="?",
        default="host",
        choices=["all", "host", "build"],
        help=(
            "What should be cleaned. 'build' for just the build platform, or "
            "'host' for the host platform, or 'all' for both. Defaults to 'host'."
        ),
    )

    for subcommand in (
        install_emscripten_cmd,
        build,
        configure_build,
        make_libffi_cmd,
        make_mpdec_cmd,
        make_build,
        configure_host,
        make_host,
        clean,
    ):
        subcommand.add_argument(
            "--quiet",
            action="store_true",
            default=False,
            dest="quiet",
            help="Redirect output from subprocesses to a log file",
        )
        subcommand.add_argument(
            "--cross-build-dir",
            action="store",
            default=None,
            dest="cross_build_dir",
            help="Path to the cross-build directory "
            f"(default: {DEFAULT_CROSS_BUILD_DIR})",
        )
        subcommand.add_argument(
            "--emsdk-cache",
            action="store",
            default=None,
            dest="emsdk_cache",
            help="Path to emsdk cache directory. If provided, validates that "
            "the required emscripten version is installed.",
        )
    for subcommand in configure_build, configure_host:
        subcommand.add_argument(
            "--clean",
            action="store_true",
            default=False,
            dest="clean",
            help="Delete any relevant directories before building",
        )
    for subcommand in build, configure_build, configure_host:
        subcommand.add_argument(
            "args", nargs="*", help="Extra arguments to pass to `configure`"
        )
    for subcommand in build, configure_host:
        subcommand.add_argument(
            "--host-runner",
            action="store",
            default=default_host_runner,
            dest="host_runner",
            help="Command template for running the emscripten host"
            f"`{default_host_runner}`)",
        )

    context = parser.parse_args()

    if context.emsdk_cache and context.subcommand != "install-emscripten":
        validate_emsdk_version(context.emsdk_cache)
        context.emsdk_cache = Path(context.emsdk_cache).absolute()
    else:
        print("Build will use EMSDK from current environment.")

    context.build_paths = get_build_paths(
        context.cross_build_dir, context.emsdk_cache
    )

    dispatch = {
        "install-emscripten": install_emscripten,
        "make-libffi": make_emscripten_libffi,
        "make-mpdec": make_mpdec,
        "configure-build-python": configure_build_python,
        "make-build-python": make_build_python,
        "configure-host": configure_emscripten_python,
        "make-host": make_emscripten_python,
        "build": build_target,
        "clean": clean_contents,
    }

    if not context.subcommand:
        # No command provided, display help and exit
        print(
            "Expected one of",
            ", ".join(sorted(dispatch.keys())),
            file=sys.stderr,
        )
        parser.print_help(sys.stderr)
        sys.exit(1)
    dispatch[context.subcommand](context)


if __name__ == "__main__":
    main()