#!/usr/bin/env python3
# Download the pinned LLVM release and unpack it into ext/llvm, mirroring the old
# premake5.lua `fetch-llvm` action. Run this once before `meson setup` on a
# machine without a suitable LLVM install:
#
#   python3 tools/fetch_llvm.py
#
# It is a no-op when ext/llvm already carries the LLDB headers.

import os
import platform
import shutil
import sys
import tarfile
import urllib.request

LLVM_VERSION = "22.1.0"

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
EXT = os.path.join(REPO_ROOT, "ext")
VENDORED = os.path.join(EXT, "llvm")
STAGING = os.path.join(EXT, "llvm-staging")


def is_arm():
    machine = platform.machine().lower()
    return machine in ("arm64", "aarch64")


def release_asset():
    host = platform.system()

    if host == "Darwin":
        # LLVM only publishes an Apple-silicon macOS binary
        return f"LLVM-{LLVM_VERSION}-macOS-ARM64.tar.xz" if is_arm() else None

    if host == "Linux":
        arch = "ARM64" if is_arm() else "X64"
        return f"LLVM-{LLVM_VERSION}-Linux-{arch}.tar.xz"

    if host == "Windows":
        arch = "aarch64" if is_arm() else "x86_64"
        return f"clang+llvm-{LLVM_VERSION}-{arch}-pc-windows-msvc.tar.xz"

    return None


def download(url, dest):
    print("downloading " + url)
    reported = [0]

    def hook(block, block_size, total):
        current = block * block_size
        if total > 0 and current - reported[0] >= total / 20:
            reported[0] = current
            pct = min(100, int(current / total * 100))
            print(f"  {pct}% ({current / 1048576:.0f} mb)")

    tmp = dest + ".part"
    try:
        urllib.request.urlretrieve(url, tmp, hook)
    except OSError as error:
        if os.path.isfile(tmp):
            os.remove(tmp)
        raise SystemExit(f"download failed: {error}")
    os.replace(tmp, dest)


def extract(archive, dest):
    print("extracting " + os.path.basename(archive))
    try:
        with tarfile.open(archive, "r:xz") as tar:
            tar.extractall(dest)
        return
    except (tarfile.TarError, OSError, EOFError) as error:
        # some stripped-down Pythons lack the lzma module; fall back to system tar
        print(f"  python extraction failed ({error}), trying system tar")

    if os.system(f'tar -xf "{archive}" -C "{dest}"') != 0:
        raise SystemExit(
            f"could not extract {archive}; unpack it into {VENDORED} by hand"
        )


def main():
    if os.path.isdir(os.path.join(VENDORED, "include", "lldb")):
        print(f"{VENDORED} already has an llvm with lldb headers")
        return 0

    asset = release_asset()
    if asset is None:
        raise SystemExit(
            f"no pinned llvm release for {platform.system()} "
            f"{platform.machine()}; install llvm yourself and set LLVM_PREFIX"
        )

    archive = os.path.join(EXT, asset)
    if not os.path.isfile(archive):
        url = (
            "https://github.com/llvm/llvm-project/releases/download/"
            f"llvmorg-{LLVM_VERSION}/{asset}"
        )
        download(url, archive)

    shutil.rmtree(STAGING, ignore_errors=True)
    os.makedirs(STAGING)
    extract(archive, STAGING)

    unpacked = [
        os.path.join(STAGING, name)
        for name in os.listdir(STAGING)
        if os.path.isdir(os.path.join(STAGING, name))
    ]
    if not unpacked:
        raise SystemExit(f"{archive} did not contain an llvm directory")

    shutil.rmtree(VENDORED, ignore_errors=True)
    os.rename(unpacked[0], VENDORED)

    shutil.rmtree(STAGING, ignore_errors=True)
    os.remove(archive)

    print(f"llvm {LLVM_VERSION} is ready in {VENDORED}")
    return 0


if __name__ == "__main__":
    sys.exit(main())