# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

# NB: This script requires that the underlying docker installation
# has `buildx` support, and also that the underlying system has
# QEMU installed so that Docker can run images of a different
# architecture.

# This script can be run via `pants run`, but it needs to be run
# with the `--no-watch-filesystem` flag, because the contents of the
# `awslambda` backend depend on the resources generated by this script.
# If run without the `--no-watch-filesystem` flag, the script will
# loop infinitely, because it restarts itself whenever the output files
# change (FIXME #21243).
#
# Command: pants --no-watch-filesystem run build-support/bin/generate_faas_complete_platforms.py

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path

from pants.backend.awslambda.python.target_types import PythonAwsLambdaRuntime
from pants.backend.google_cloud_function.python.target_types import PythonGoogleCloudFunctionRuntime
from pants.backend.python.util_rules.faas import FaaSArchitecture, PythonFaaSRuntimeField
from pants.base.build_environment import get_buildroot

# GCF images throw permissions errors when trying to install pex in the default site-packages.
# Additionally, some of the images are missing a `pip` alias, so we need to use `python -m pip`.
COMMAND = (
    "python -m pip install --target=/tmp/subdir pex 1>&2 && "
    + "PYTHONPATH=/tmp/subdir /tmp/subdir/bin/pex3 interpreter inspect --markers --tags"
)


def extract_complete_platform(repo: str, architecture: FaaSArchitecture, tag: str) -> object:
    image = f"{repo}:{tag}"
    docker_platform = "linux/amd64" if architecture == FaaSArchitecture.X86_64 else "linux/arm64"
    print(
        f"Extracting complete platform for {image} on platform {docker_platform}", file=sys.stderr
    )
    result = subprocess.run(
        [
            "docker",
            "run",
            "--platform",
            docker_platform,
            "--entrypoint",
            "/bin/sh",
            image,
            "-c",
            COMMAND,
        ],
        check=True,
        stdout=subprocess.PIPE,
    )
    return json.loads(result.stdout)


def run(runtime_field: type[PythonFaaSRuntimeField], python_base: Path) -> None:
    cp_dir = python_base / runtime_field.known_runtimes_complete_platforms_module().replace(
        ".", "/"
    )
    print(f"Generating for {runtime_field.__name__}, writing to {cp_dir}", file=sys.stderr)
    for rt in runtime_field.known_runtimes:
        cp = extract_complete_platform(
            rt.docker_repo,
            FaaSArchitecture(rt.architecture) if rt.architecture else FaaSArchitecture.X86_64,
            rt.tag,
        )

        fname = cp_dir / rt.file_name()
        with fname.open("w") as f:
            json.dump(cp, f, indent=2)


def create_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Generates the complete platform JSON files for AWS Lambda and/or GCF"
    )
    parser.add_argument(
        "--runtime",
        choices=["lambda", "gcf", "all"],
        default="all",
        help="Choose which runtime(s) to generate complete platform files for",
    )
    return parser


def main() -> None:
    args = create_parser().parse_args()

    build_root = Path(get_buildroot()) / "src/python"

    # Type declaration needed for mypy
    selected_runtimes: list[type[PythonFaaSRuntimeField]]
    if args.runtime == "lambda":
        selected_runtimes = [PythonAwsLambdaRuntime]
    elif args.runtime == "gcf":
        selected_runtimes = [PythonGoogleCloudFunctionRuntime]
    else:
        selected_runtimes = [PythonAwsLambdaRuntime, PythonGoogleCloudFunctionRuntime]

    for runtime_field in selected_runtimes:
        run(runtime_field, build_root)


if __name__ == "__main__":
    main()
