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


Import("env")


cpython_build = Dir("cpython_build")


env["bits"] = "64"
env["pandemonium_binary_download_platform"] = "win64.exe"
env["cpython_build"] = cpython_build
env["cpython_build_dir"] = cpython_build
env["DIST_SITE_PACKAGES"] = Dir(f"{env['DIST_PLATFORM']}/Lib/site-packages")


### Build config for pythonscript ###

# 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"])
env.AppendUnique(LINKFLAGS=[f"/LIBPATH:{cpython_build.abspath}/libs"])
env.AppendUnique(CYTHON_COMPILE_DEPS=[cpython_build])


### Fetch Python prebuild ###


CPYTHON_PREBUILD_URL = "https://github.com/indygreg/python-build-standalone/releases/download/20200830/cpython-3.8.5-x86_64-pc-windows-msvc-shared-pgo-20200830T2254.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-pc-windows-msvc"

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

    stdlib_path = build / "Lib"

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

    # Remove .pdb debug symbols
    for pdbfile in (build / "DLLs").glob("*.pdb"):
        pdbfile.unlink()

    # 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"]:
        shutil.make_archive(base_name=build / "python38", format="zip", root_dir=str(stdlib_path))
        shutil.rmtree(str(stdlib_path))
        stdlib_path.mkdir()
        # 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)