#!/usr/bin/env python3
import argparse
import os
import platform
import subprocess
import sys
import tempfile

from ci import git_ls_py_files

basedir = os.path.dirname(os.path.realpath(__file__))


def get_python_version(python):
    hexversion = int(
        subprocess.check_output([python, "-c", "import sys; print(sys.hexversion)"])
    )
    return "{}.{}.{}".format(
        (hexversion >> 24) & 0xFF,
        (hexversion >> 16) & 0xFF,
        (hexversion >> 8) & 0xFF,
    )


def run_in_venv(venv, args):
    subprocess.check_call(
        [os.path.join(basedir, "run-in-venv"), venv, *args], cwd=basedir
    )


def pip_install_in_venv(venv, pip, args):
    run_in_venv(
        venv,
        [
            pip,
            "install",
            "--disable-pip-version-check",
            "--upgrade",
            *args,
        ],
    )


def host2machine(host):
    return {"arm": "armv7l"}.get(host, host)


def setup_build_venv(build, python, host_python, host):
    venv = os.path.join(build, get_python_version(python))
    if host_python is None:
        subprocess.check_call([python, "-m", "venv", venv])
    else:
        bootstrap_venv = os.path.join(venv, "bootstrap")
        subprocess.check_call([python, "-m", "venv", bootstrap_venv])
        pip_install_in_venv(bootstrap_venv, "pip", ["crossenv"])
        run_in_venv(
            bootstrap_venv,
            [
                "python",
                "-m",
                "crossenv",
                f"--machine={host2machine(host)}",
                host_python,
                venv,
            ],
        )
    if host_python is None:
        pip_install_in_venv(
            venv,
            "pip",
            [
                "--requirement=build-requirements.txt",
                "--requirement=cross-requirements.txt",
            ],
        )
    else:
        pip_install_in_venv(venv, "build-pip", ["--requirement=build-requirements.txt"])
        cross_requirements = ["--requirement=cross-requirements.txt"]
        if sys.version_info >= (3, 12):
            cross_requirements.append("setuptools")
        pip_install_in_venv(venv, "pip", cross_requirements)
    return venv


def repair_wheel(venv, whl, wheel_dir, cross_prefix, host):
    with tempfile.TemporaryDirectory(dir=wheel_dir) as tmp_wheel_dir:
        run_in_venv(
            venv,
            [
                f"{cross_prefix}python",
                os.path.join(basedir, "auditwheel"),
                "repair",
                f"--plat=manylinux2014_{host2machine(host)}",
                f"--wheel-dir={tmp_wheel_dir}",
                whl,
            ],
        )
        (tmp_manylinux_whl_name,) = os.listdir(tmp_wheel_dir)
        manylinux_whl = os.path.join(wheel_dir, tmp_manylinux_whl_name)
        os.rename(
            os.path.join(tmp_wheel_dir, tmp_manylinux_whl_name),
            manylinux_whl,
        )
        return manylinux_whl


def build_wheel(
    venv, build_type, build_args, wheel_dir, skip_repair, cross_prefix, host
):
    with tempfile.TemporaryDirectory() as dist_dir:
        run_in_venv(
            venv,
            [
                "python",
                "setup.py",
                "bdist_wheel",
                f"--build-type={build_type}",
                f"--dist-dir={dist_dir}",
                *build_args,
            ],
        )
        (whl_name,) = os.listdir(dist_dir)
        whl = os.path.join(dist_dir, whl_name)
        os.makedirs(wheel_dir, exist_ok=True)
        if skip_repair:
            result = os.path.join(wheel_dir, whl_name)
            os.rename(whl, result)
            return result
        else:
            return repair_wheel(venv, whl, wheel_dir, cross_prefix, host)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--build-type", default="Release")
    parser.add_argument("--python", default="python3")
    parser.add_argument("--host", default=platform.machine())
    parser.add_argument("--host-python")
    parser.add_argument("--skip-tests", action="store_true")
    parser.add_argument("--skip-repair", action="store_true")
    parser.add_argument(
        "--wheel-dir",
        default=os.path.join(basedir, "dist", "wheelhouse"),
    )
    parser.add_argument("--sanitize", action="store_true")
    args, build_args = parser.parse_known_args()
    if args.host != platform.machine():
        build_args.extend(
            (
                f"-DCMAKE_SYSTEM_NAME={platform.system()}",
                f"-DCMAKE_SYSTEM_PROCESSOR={args.host}",
            )
        )
    if args.sanitize:
        build_args.append("-DCMAKE_CXX_FLAGS=-fsanitize=address,undefined")
    build = os.path.join(basedir, "build", f"{platform.system()}-{args.host}")
    build_prefix = "" if args.host_python is None else "build-"
    cross_prefix = "" if args.host_python is None else "cross-"
    venv = setup_build_venv(build, args.python, args.host_python, args.host)
    run_in_venv(venv, [f"{build_prefix}python", "-m", "flake8", *git_ls_py_files()])
    if not args.skip_tests:
        run_in_venv(
            venv,
            [
                "python",
                "setup.py",
                "build",
                f"--build-type={args.build_type}",
                *build_args,
            ],
        )
        run_in_venv(
            venv,
            [
                "./python3-asan" if args.sanitize else "python3",
                "-m",
                "unittest",
                "discover",
            ],
        )
    build_wheel(
        venv,
        args.build_type,
        build_args,
        args.wheel_dir,
        args.skip_repair,
        cross_prefix,
        args.host,
    )


if __name__ == "__main__":
    main()
