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"] = "osx.64"
env["pandemonium_binary_download_zip_path"] = "Pandemonium.app/Contents/MacOS/Pandemonium"
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"])


### Fetch Python prebuild ###


CPYTHON_PREBUILD_URL = "https://github.com/indygreg/python-build-standalone/releases/download/20200823/cpython-3.8.5-x86_64-apple-darwin-pgo-20200823T2228.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-apple-darwin"

    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))

    # Patch binaries to load libpython3.x.dylib with a relative path
    # Lib paths are hardcoded into the executable, and if the lib is not found at the path, then it craps out.
    # Unfortunately compiling python will hardcode the absolute path of libpython.dylib into the executable,
    # so if you move it around it will break.
    # the solution here is to modify the executable and make sure the lib path is not an absolute path,
    # but an path relative to @loader_path, which is a special symbol that points to the executable.
    # See: http://joaoventura.net/blog/2016/embeddable-python-osx-from-src/
    # and https://stackoverflow.com/questions/7880454/python-executable-not-finding-libpython-shared-library
    prebuild_shared_lib_path = conf["build_info"]["core"]["shared_lib"]
    path, _ = prebuild_shared_lib_path.rsplit("/", 1)
    assert path == "install/lib"  # Make sure libpython.so is on lib folder
    binary = build / "bin/python3.8"
    assert binary.is_file()
    dylib = build / "lib/libpython3.8.dylib"
    cmd = f"install_name_tool -id @rpath/{dylib.name} {dylib}"
    subprocess.run(cmd.split(), check=True)

    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)