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

ADD_PACKAGE_PREFIX = " + "

def log (message, *args):
    print(message.format(*args), file=sys.stderr)

def get_all_package_generations ():
    package_generations = {}
    for line in check_output(["guix", "package", "--list-generations"]).decode().split("\n"):
        if not line.startswith(ADD_PACKAGE_PREFIX):
            continue
            
        (package, version, output, path) = line[len(ADD_PACKAGE_PREFIX):].split("\t")
        if not package in package_generations:
            package_generations[package] = []
        
        package_generations[package].append({
            "path": path,
            "version": version,
            "output": output,
            "package": package
        })
        
    return package_generations

def run_with_packages (packages, command):
    env = {}
    env.update(os.environ)
    env.update({
        "PATH": ":".join(["{}/bin".format(package['path']) for package in packages] + [env.get("PATH", "")]),
        "LD_LIBRARY_PATH": ":".join(["{}/lib".format(package['path']) for package in packages] + [env.get("PATH", "")])
    })
    call(args.run, env=env)    
    
parser = argparse.ArgumentParser()
parser.add_argument("--package", help="package name", required=True)
parser.add_argument("--version", help="package version")
parser.add_argument("--index", help="", type=int, default=0)
parser.add_argument("--run", nargs="*", help="run command")
args = parser.parse_args()

all_package_generations = get_all_package_generations()
package_generations_for_package = all_package_generations[args.package]
package_generations_for_package.reverse()

log("Found {} generations for package {}", len(package_generations_for_package), args.package)

if args.version:
    package_generations_for_package = [generation for generation in package_generations_for_package if generation['version'] == args.version]
    log("Found {} generations for version {}", len(package_generations_for_package), args.version)

selected_package = package_generations_for_package[args.index]

log("Selected package: {}", selected_package)

if args.run:
    run_with_packages([selected_package], args.run)
else:
    print(selected_package)