#!/usr/bin/env python3

import subprocess
import sys
import os
import shutil
import re
import stat
import tarfile
import requests
import hashlib

run_script_linux = """#!/bin/bash

set -e

script_path=`readlink -f $0`
script_dir=`dirname $script_path`

program_full_name="$PROGRAM_FULL_NAME"
$DOWNLOAD_DEPENDENCIES_COMMAND
"$script_dir/$SO_LOADER" --library-path "$script_dir/libs":$HOME/.local/lib/sibs/"$program_full_name":/usr/lib/sibs/"$program_full_name" "$script_dir/$PROGRAM_NAME" "$@"
"""

def get_executable_dynamic_libraries(filepath):
    libs = []
    process = subprocess.Popen(["ldd", filepath], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (stdout, stderr) = process.communicate()
    if process.returncode != 0:
        raise RuntimeError("Failed to execute ldd on executable %s, error:\n%s" % (filepath, stderr))
    lines = stdout.splitlines()
    for line in lines:
        s = line.split()
        if len(s) >= 3 and b"=>" in s:
            lib_path = " ".join(x.decode("UTF-8") for x in s[2:-1])
            if os.path.exists(lib_path):
                libs.append([os.path.realpath(lib_path), os.path.basename(lib_path)])
    return libs

def make_executable(filepath):
    mode = os.stat(filepath).st_mode
    os.chmod(filepath, mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)

def get_libnss_files():
    libnss_pattern = re.compile("libnss.*\\.so.*")
    files = []
    for filepath in os.listdir("/usr/lib"):
        if libnss_pattern.match(filepath):
            files.append(os.path.join("/usr/lib", filepath))
    return files

def file_get_sha1(filepath):
    with open(filepath, "rb") as f:
        sha1 = hashlib.sha1()
        sha1.update(f.read())
        return sha1.hexdigest()
    raise RuntimeError("No such file: %s" % filepath)

def usage():
    print("usage: package.py executable_path program_version destination_path <--bundle|--bundle-install>")
    exit(1)

def main():
    if len(sys.argv) <= 4:
        usage()

    script_dir = os.path.dirname(os.path.realpath(sys.argv[0]))
    executable_path = sys.argv[1]
    program_version = sys.argv[2]
    destination_path = sys.argv[3]
    package_type = sys.argv[4]

    if package_type != "--bundle" and package_type != "--bundle-install":
        usage()

    os.makedirs(destination_path, exist_ok=True)
    libs = get_executable_dynamic_libraries(executable_path)

    so_loader_pattern = re.compile("ld-linux-x86-64\\.so.*")
    so_loader = None
    for lib in libs:
        if so_loader_pattern.match(lib[1]):
            if so_loader != None:
                print("Unexpected error: found multiple so loaders, unable to recover")
                exit(2)
            so_loader = lib[1]
    
    if not so_loader:
        print("Unexpected error: no so loader found, unable to recover")
        exit(5)

    print("So loader: %s" % so_loader)

    mapped_libs = []
    dependencies = []
    if package_type == "--bundle-install":
        new_libs = []
        # Remove libraries from the package that can be downloaded remotely (and which user most likely has if they have another program that uses sibs on their computer)
        for lib in libs:
            lib_name = os.path.basename(lib[0])
            r = requests.get("https://raw.githubusercontent.com/DEC05EBA/libraries/master/linux/x86_64/" + lib_name + ".sha1")
            # TODO: Remove check if it's so_loader or not, we can use so loader in sibs downloaded lib directory
            if r.ok and lib[1] != so_loader:
                external_checksum = r.text.splitlines()[0]
                file_checksum = file_get_sha1(lib[0])
                if external_checksum == file_checksum:
                    mapped_libs.append([ lib_name, lib[1] ])
                    dependencies.append([ lib_name, external_checksum ])
                else:
                    new_libs.append(lib)
                print("Checksum for file %s: %s, equals? %s" % (lib_name, external_checksum, "yes" if external_checksum == file_checksum else "no"))
            else:
                new_libs.append(lib)
        libs = new_libs
    so_loader = os.path.join("libs", so_loader)

    libnss_files = get_libnss_files()
    for idx, libnss_file in enumerate(libnss_files):
        new_file_path = os.path.join(destination_path, os.path.basename(libnss_file))
        libnss_files[idx] = new_file_path
        if os.path.islink(libnss_file):
            target_name = os.path.basename(os.readlink(libnss_file))
            os.symlink(target_name, new_file_path)
        else:
            shutil.copyfile(libnss_file, new_file_path)

    executable_filename = os.path.basename(executable_path)
    new_executable_path = os.path.join(destination_path, executable_filename)
    shutil.copyfile(executable_path, new_executable_path)
    make_executable(new_executable_path)

    print("Patching executable")
    process = subprocess.Popen(["patchelf", "--set-interpreter", so_loader, new_executable_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (stdout, stderr) = process.communicate()
    if process.returncode != 0:
        print("Failed to execute patchelf --set-interpreter on executable %s, error: %s" % (new_executable_path, stderr))
        exit(3)
    
    #process = subprocess.Popen(["patchelf", "--set-rpath", "libs", new_executable_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    #(stdout, stderr) = process.communicate()
    #if process.returncode != 0:
    #    print("Failed to execute patchelf --set-rpath on executable %s, error: %s" % (new_executable_path, stderr))
    #    exit(4)

    run_script_path = os.path.join(destination_path, "run.sh")
    with open(run_script_path, "wb") as run_script:
        run_script.write(run_script_linux.replace("$SO_LOADER", so_loader)\
            .replace("$PROGRAM_NAME", "program")\
            .replace("$PROGRAM_FULL_NAME", executable_filename + "." + program_version)\
            .replace("$DOWNLOAD_DEPENDENCIES_COMMAND", "\"$script_dir/download_dependencies.sh\" \"$program_full_name\"" if package_type == "--bundle-install" else "")\
            .encode("UTF-8"))
    make_executable(run_script_path)

    urls_filepath = os.path.join(destination_path, "urls.conf")
    with open(urls_filepath, "wb") as urls_file:
        urls_file.write("https://raw.githubusercontent.com/DEC05EBA/libraries/master/linux/x86_64/\n".encode("UTF-8"))

    library_mapping_filepath = os.path.join(destination_path, "libmap.conf")
    with open(library_mapping_filepath, "wb") as library_mapping_file:
        for mapped_lib in mapped_libs:
            library_mapping_file.write((mapped_lib[0] + "\n" + mapped_lib[1] + "\n").encode("UTF-8"))

    dependencies_filepath = os.path.join(destination_path, "dependencies.conf")
    with open(dependencies_filepath, "wb") as dependencies_file:
        for dependency in dependencies:
            dependencies_file.write((dependency[0] + "\n" + dependency[1] + "\n").encode("UTF-8"))

    package_name = new_executable_path + ".tar.gz"
    print("Creating archive %s" % os.path.basename(package_name))
    with tarfile.open(package_name, "w:gz") as tar:
        for lib in libs:
            print("Adding shared library %s to package" % lib[0])
            tar.add(lib[0], arcname=os.path.join("libs", lib[1]))

        for libnss_file in libnss_files:
            print("Adding libnss shared library %s to package" % libnss_file)
            libnss_name = os.path.basename(libnss_file)
            tar.add(libnss_file, arcname=os.path.join("libs", libnss_name))

        print("Adding executable %s to package" % executable_path)
        tar.add(new_executable_path, arcname="program")
        print("Adding run script %s to package" % run_script_path)
        tar.add(run_script_path, arcname=executable_filename)
        if package_type == "--bundle-install":
            print("Adding urls file %s to package" % urls_filepath)
            tar.add(urls_filepath, arcname="urls.conf")
            print("Adding library mapping file %s to package" % library_mapping_filepath)
            tar.add(library_mapping_filepath, arcname="libmap.conf")
            print("Adding dependencies file %s to package" % dependencies_filepath)
            tar.add(dependencies_filepath, arcname="dependencies.conf")
            download_dependencies_script_filepath = os.path.join(script_dir, "download_dependencies.sh")
            print("Adding download dependencies script file %s to package" % download_dependencies_script_filepath)
            tar.add(download_dependencies_script_filepath, arcname="download_dependencies.sh")
    
    print("Removing temporary files")
    os.remove(new_executable_path)
    os.remove(run_script_path)
    os.remove(urls_filepath)
    os.remove(library_mapping_filepath)
    os.remove(dependencies_filepath)
    for libnss_file in libnss_files:
        os.remove(libnss_file)
    print("Package has been created at %s" % package_name)

if __name__ == "__main__":
    main()