#!/usr/bin/env python3
"""Cache a native build: test_native_build.py OUTPUT -- COMPILER FLAGS SOURCES LIBS."""
import fcntl
import hashlib
import json
import os
from pathlib import Path
import shlex
import shutil
import subprocess
import sys
import tempfile
import time


def digest(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def dependencies(text):
    paths = set()
    for rule in text.replace("\\\n", "").splitlines():
        if ": " in rule:
            paths.update(shlex.split(rule.split(": ", 1)[1].replace("$$", "$")))
    return paths


def hashes(paths):
    return {str(Path(p).absolute()): digest(p) for p in sorted(paths)}


def main():
    if len(sys.argv) < 5 or sys.argv[2] != "--":
        print(__doc__, file=sys.stderr)
        return 2
    output = Path(sys.argv[1]).absolute()
    command = sys.argv[3:]
    compiler = Path(shutil.which(command[0]) or command[0]).resolve()
    stat = compiler.stat()
    version = subprocess.check_output([str(compiler), "--version"], timeout=15).decode()
    key = [command, os.getcwd(), str(compiler), version, stat.st_size, stat.st_mtime_ns,
           stat.st_ctime_ns, digest(__file__), {name: os.environ.get(name) for name in
           ("PATH", "CPATH", "CPLUS_INCLUDE_PATH", "C_INCLUDE_PATH", "LIBRARY_PATH",
            "COMPILER_PATH", "GCC_EXEC_PREFIX", "SDKROOT", "SOURCE_DATE_EPOCH")}]
    output.parent.mkdir(parents=True, exist_ok=True)
    manifest = Path(str(output) + ".json")
    with open(str(output) + ".lock", "a") as lock:
        deadline = time.monotonic() + 300
        while True:
            try:
                fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
                break
            except BlockingIOError:
                if time.monotonic() >= deadline:
                    print(f"Native build lock timed out: {output}", file=sys.stderr)
                    return 1
                time.sleep(0.1)
        try:
            cached = json.loads(manifest.read_text())
            if (cached["key"] == key and cached["dependencies"] == hashes(cached["dependencies"])
                    and cached["binary"] == digest(output) and os.access(output, os.X_OK)):
                return 0
        except (OSError, ValueError, KeyError, TypeError):
            pass
        manifest.unlink(missing_ok=True)
        # Multiple translation units overwrite one -MF file. Scan their union before the build.
        scan = subprocess.run(command + ["-M"], stdout=subprocess.PIPE, timeout=120)
        if scan.returncode:
            return scan.returncode
        paths = dependencies(scan.stdout.decode())
        paths.update(arg for arg in command[1:] if Path(arg).is_file())
        for arg in command[1:]:
            if arg.startswith("-l"):
                for suffix in (".so", ".a"):
                    library = subprocess.check_output(
                        [str(compiler), "-print-file-name=lib" + arg[2:] + suffix], timeout=15
                    ).decode().strip()
                    if Path(library).is_file():
                        paths.add(library)
        before = hashes(paths)
        with tempfile.TemporaryDirectory(prefix=output.name + ".", dir=output.parent) as work:
            binary, depfile = Path(work) / "binary", Path(work) / "deps"
            result = subprocess.run(command + ["-MD", "-MF", str(depfile), "-o", str(binary)], timeout=300)
            if result.returncode:
                return result.returncode
            paths.update(dependencies(depfile.read_text()))
            after = hashes(paths)
            if before != after:
                print(f"Native build inputs changed: {output}. Retry the build.", file=sys.stderr)
                return 1
            state = {"key": key, "dependencies": after, "binary": digest(binary)}
            metadata = Path(work) / "manifest"
            metadata.write_text(json.dumps(state))
            os.replace(binary, output)
            os.replace(metadata, manifest)
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except (OSError, subprocess.SubprocessError) as error:
        print(f"Native build failed: {error}", file=sys.stderr)
        sys.exit(1)