#!/usr/bin/env python3
# mypy: disable-error-code=no-any-unimported
#
# Originally come from MagicStack's Network Server Performance Benchmarking Toolbench
# https://github.com/MagicStack/vmbench

from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import sys
from itertools import chain
from pathlib import Path
from typing import Any, Literal

import build as pyproject_build
import docker

ROOT_DIR = Path(__file__).parent

SUPPORTED_PYTHON_VERSIONS = {
    "3.11": "3.11.14",
    "3.12": "3.12.12",
    "3.13": "3.13.11",
    "3.14": "3.14.2",
}


def _build_sdist_for_docker_image() -> Path:
    src_dir = ROOT_DIR.parent
    output_dir = ROOT_DIR / "dist"
    distribution: Literal["sdist"] = "sdist"
    config_settings: dict[str, Any] = {}

    shutil.rmtree(output_dir, ignore_errors=True)

    print("Building source distribution for docker image...")
    with pyproject_build.env.DefaultIsolatedEnv() as env:
        builder = pyproject_build.ProjectBuilder.from_isolated_env(env, src_dir)
        env.install(builder.build_system_requires)
        env.install(builder.get_requires_for_build(distribution, config_settings))
        sdist = Path(builder.build(distribution, output_dir, config_settings))
        print(f"-> Sucessfully built {sdist.relative_to(ROOT_DIR)}")
        return sdist


def _build_docker_image(client: docker.DockerClient, sdist_path: Path, tag: str, python_version: str) -> None:
    sdist_path = sdist_path.relative_to(ROOT_DIR)
    print("Building docker image...")
    # NOTE: Do not use docker-py API to build the, because BuildKit is not supported
    # c.f. https://github.com/docker/docker-py/issues/2230
    docker_build_command_line = [
        "docker",
        "build",
        "--rm=true",
        "--force-rm",
        "--progress=plain",
        f"--tag={tag}",
    ]
    buildargs = {
        "PYTHON_VERSION": SUPPORTED_PYTHON_VERSIONS[python_version],
        "EASYNETWORK_SDIST": os.fspath(sdist_path),
    }
    docker_build_command_line += chain.from_iterable(["--build-arg", f"{key}={value}"] for key, value in buildargs.items())
    docker_build_command_line += [os.fspath(ROOT_DIR)]

    with subprocess.Popen(
        docker_build_command_line,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    ) as process:
        assert process.stdout is not None
        while c := process.stdout.read(1):
            sys.stdout.write(c)
        sys.stdout.flush()

    if (retcode := process.returncode) != 0:
        sys.exit(retcode)

    client.images.prune(filters={"dangling": True})
    print(f"-> Sucessfully built docker image {tag!r}")


def main() -> None:
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        "-t",
        "--tag",
        default="easynetwork/benchmark",
        help="Built image tag",
    )
    parser.add_argument(
        "-p",
        "--python-version",
        default=list(SUPPORTED_PYTHON_VERSIONS)[-1],
        choices=list(SUPPORTED_PYTHON_VERSIONS),
        help="The CPython version to use for the servers (in major.minor format)",
    )
    args = parser.parse_args()

    python_version: str = args.python_version
    print(f"Selected Python version: {python_version} (Docker image tag: python:{SUPPORTED_PYTHON_VERSIONS[python_version]})")

    client = docker.from_env()
    _build_docker_image(client, _build_sdist_for_docker_image(), args.tag, python_version)


if __name__ == "__main__":
    main()
