#!/usr/bin/env python3
from subprocess import run, check_output
import sys
import os
import argparse
import shutil

SEPARATOR = " => "
LIBRARY_STORE = os.path.join("/", "gnu", "store")
CURRENT_PROFILE = os.path.join("/", "run", "current-system", "profile")
USER_PROFILE = os.path.join(os.environ["HOME"], ".guix-profile")

LOADER = os.path.join(USER_PROFILE, "lib", "ld-linux-x86-64.so.2")

SEARCH_PATHS = [
    os.path.join(USER_PROFILE, "lib"),
    os.path.join(CURRENT_PROFILE, "lib"),
    LIBRARY_STORE
]

libraries = {}

def file_type(path):
    return check_output(["file", "-L", path]).decode().strip()

def scan_input_binary(input_binary):
    not_found = []
    discovered = {}
    for line in check_output(["ldd", input_binary]).decode().split("\n"):
        line = line.strip()
        if not SEPARATOR in line:
            continue
        (library, path) = line.split(SEPARATOR)
        if path == "not found":
            not_found.append(library)
        else:
            discovered[library] = path.split(" ")[0]
    return (discovered, not_found)

def find_library(library):
    for path in SEARCH_PATHS:
        print(">> >> >> Looking for {0} in: {1}".format(library, path), file=sys.stderr)
        found_library = find_library_in_search_path(library, path)
        if found_library:
            return found_library

def find_library_in_search_path(library, path):
    for (dirpath, dirnames, filenames) in os.walk(path):
        for name in ["sbin", "bin", "share", "include", "etc", "site-packages", "src", "vendor_ruby", "kernel"]:
            if name in dirnames:
                dirnames.remove(name)

        if library in filenames:
            library_path = os.path.join(dirpath, library)
            if "ELF 64-bit" in file_type(library_path):
                return library_path

def find_libraries(already_found, to_find):
    resolved = {}
    for library in to_find:
        print(">> Looking for: {0}".format(library), file=sys.stderr)
        if library in already_found:
            print(">> >> Using already-found path for {0}: {1}".format(
                library, already_found[library]
            ), file=sys.stderr)
            resolved[library] = already_found[library]
        else:
            print(">> >> Performing search for: {0}".format(library), file=sys.stderr)
            library_path = find_library(library)
            if library_path:
                print(">> >> >> >> Resolved {0} as: {1}".format(library, library_path), file=sys.stderr)
                resolved[library] = library_path
            else:
                print(">> >> >> >> Failed to resolve {0}".format(library), file=sys.stderr)
                raise Exception("Failed to resolve {0}".format(library))
    return resolved

def make_ld_library_path(paths):
    return ":".join(set(paths))

def generate_wrapper(paths, command):
    return """#!/bin/sh
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:{path}"
{command} $@""".format(path=make_ld_library_path(paths), command=" ".join(command))

def print_paths(paths):
    print("---")
    print(":".join(set(paths)))

parser = argparse.ArgumentParser()
parser.add_argument("--wrap", action="store_true", help="output wrapper script to stdout")
parser.add_argument("--run", nargs="*", help="run command with specified args")
parser.add_argument("--command", help="command to wrap or run (default to the first input binary)")
parser.add_argument("--arguments", nargs="+", default=[], help="Additional arguments to pass to command")
parser.add_argument("--inputs", nargs="+", default=[], help="input binaries to scan for dependencies")
parser.add_argument("--paths", nargs="*", default=[], help="Additional paths to add to library search path")
parser.add_argument("--patch", help="create a patched copy of command with the given name")
args = parser.parse_args()

paths = args.paths
pre_command = [LOADER]
command = args.command

for input_binary in args.inputs:
    (discovered, not_found) = scan_input_binary(input_binary)
    libraries.update(discovered)
    found = find_libraries(libraries, not_found)
    libraries.update(found)
    paths = paths + [os.path.dirname(path) for path in found.values()]

if (not command) and args.inputs:
    command = os.path.abspath(args.inputs[0])

#command[0] = os.path.abspath(command[0])

full_command = pre_command + [command] + args.arguments
if args.run is not None:
    environment = {"LD_LIBRARY_PATH": make_ld_library_path(paths)}
    environment.update(os.environ)
    completed_process = run(full_command + args.run, env=environment)
    sys.exit(completed_process.returncode)
elif args.wrap:
    print(generate_wrapper(paths, full_command))
elif args.patch is not None:
    shutil.copy(command, args.patch)
    run([
        "patchelf",
        "--set-interpreter", LOADER,
        "--set-rpath", make_ld_library_path(paths),
        args.patch
    ])
else:
    print_paths(paths)