import zstandard
import tarfile
import json
import shutil
import subprocess
from pathlib import Path


Import("env")


cpython_build = Dir("cpython_build")


env["bits"] = "64"
if env["headless"]:
    env["pandemonium_binary_download_platform"] = "linux_headless.64"
else:
    env["pandemonium_binary_download_platform"] = "x11.64"
env["cpython_build"] = cpython_build
env["cpython_build_dir"] = cpython_build
env["DIST_SITE_PACKAGES"] = Dir(f"{env['DIST_PLATFORM']}/lib/python3.8/site-packages")


### Build config for pythonscript ###


env.AppendUnique(CFLAGS=["-m64"])
env.AppendUnique(LINKFLAGS=["-m64"])
# Cannot use CPPPATH&LIBPATH here given headers are within `cpython_build` target,
# so Scons consider the headers are a missing target
env.AppendUnique(CFLAGS=[f"-I{cpython_build.abspath}/include/python3.8/"])
env.AppendUnique(LINKFLAGS=[f"-L{cpython_build.abspath}/lib"])
env.AppendUnique(CYTHON_COMPILE_DEPS=[cpython_build])


### Fetch Python prebuild ###


CPYTHON_PREBUILD_URL = "https://github.com/indygreg/python-build-standalone/releases/download/20200822/cpython-3.8.5-x86_64-unknown-linux-gnu-pgo-20200823T0036.tar.zst"
cpython_prebuild_archive = env.Download(
    target=File(CPYTHON_PREBUILD_URL.rsplit("/", 1)[1]), url=CPYTHON_PREBUILD_URL
)
env.NoClean(cpython_prebuild_archive)


### Extract prebuild ###


def extract_cpython_prebuild(target, source, env):
    archive_path = source[0].abspath
    target_path = target[0].abspath
    with open(archive_path, "rb") as fh:
        dctx = zstandard.ZstdDecompressor()
        with dctx.stream_reader(fh) as reader:
            with tarfile.open(mode="r|", fileobj=reader) as tf:
                tf.extractall(target_path)


cpython_prebuild_src = env.Command(
    Dir("cpython_prebuild"), cpython_prebuild_archive, extract_cpython_prebuild
)
env.NoClean(cpython_prebuild_src)


### Generate custom build from the prebuild ###


def generate_cpython_build(target, source, env):
    build = Path(target[0].abspath)
    prebuild = Path(source[0].abspath) / "python"

    conf = json.loads((prebuild / "PYTHON.json").read_text())
    assert conf["version"] == "5"
    assert conf["libpython_link_mode"] == "shared"
    assert conf["target_triple"] == "x86_64-unknown-linux-gnu"

    shutil.copytree(str(prebuild / "install"), str(build), symlinks=True)
    shutil.copytree(str(prebuild / "licenses"), str(build / "licenses"), symlinks=True)

    shutil.rmtree(str(build / "share"))

    # Remove static library stuff
    config = conf["python_stdlib_platform_config"]
    assert config.startswith("install/lib/")
    config = build / config[len("install/") :]
    assert config.exists()
    shutil.rmtree(str(config))

    stdlib_path = build / "lib/python3.8"

    # Remove tests lib (pretty big and basically useless)
    shutil.rmtree(str(stdlib_path / "test"))

    # Also remove __pycache__ & .pyc stuff
    for pycache in stdlib_path.glob("**/__pycache__"):
        shutil.rmtree(str(pycache))

    # Make sure site-packages is empty to avoid including pip (ensurepip should be used instead)
    shutil.rmtree(str(stdlib_path / "site-packages"))

    # Zip the stdlib to save plenty of space \o/
    if env["compressed_stdlib"]:
        tmp_stdlib_path = build / "lib/tmp_python3.8"
        shutil.move(str(stdlib_path), str(tmp_stdlib_path))
        stdlib_path.mkdir()
        shutil.move(str(tmp_stdlib_path / "lib-dynload"), str(stdlib_path / "lib-dynload"))
        shutil.make_archive(
            base_name=build / "lib/python38", format="zip", root_dir=str(tmp_stdlib_path)
        )
        shutil.rmtree(str(tmp_stdlib_path))
        # Oddly enough, os.py must be present (even if empty !) otherwise
        # Python failed to find it home...
        (stdlib_path / "os.py").touch()

    (stdlib_path / "site-packages").mkdir()


env.Command(cpython_build, cpython_prebuild_src, generate_cpython_build)
env.NoClean(cpython_build)