/*
 * Copyright 2013, Ingo Weinhold, ingo_weinhold@gmx.de.
 * Distributed under the terms of the MIT License.
 */


#include <errno.h>
#include <getopt.h>
#include <pwd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <Architecture.h>
#include <Path.h>
#include <PathFinder.h>
#include <StringList.h>


extern const char* __progname;
const char* kCommandName = __progname;


static const char* kUsage =
        "Usage: %s [-hl]\n"
        "       %s [-p] <architecture> [ <command> ... ]\n"
        "Executes the given command or, by default, a shell with a PATH\n"
        "environment variable modified such that commands for the given\n"
        "architecture will be preferred, respectively used exclusively in case of\n"
        "the primary architecture.\n"
        "\n"
        "Options:\n"
        "  -h, --help\n"
        "    Print this usage info.\n"
        "  -l, --list-architectures\n"
        "    List all architectures.\n"
        "  -p, --print-path\n"
        "    Only print the modified PATH variable value; don't execute any\n"
        "    command.\n"
;


static void
print_usage_and_exit(bool error)
{
        fprintf(error ? stderr : stdout, kUsage, kCommandName, kCommandName);
        exit(error ? 1 : 0);
}


static bool
is_primary_architecture(const char* architecture)
{
        return strcmp(architecture, get_primary_architecture()) == 0;
}


static void
get_bin_directories(const char* architecture, BStringList& _directories)
{
        status_t error = BPathFinder::FindPaths(architecture,
                B_FIND_PATH_BIN_DIRECTORY, NULL, 0, _directories);
        if (error != B_OK) {
                fprintf(stderr, "Error: Failed to get bin directories for architecture "
                        "%s: %s\n", architecture, strerror(error));
                exit(1);
        }
}


static void
compute_new_paths(const char* architecture, BStringList& _paths)
{
        // get the primary architecture bin paths
        BStringList primaryBinDirectories;
        get_bin_directories(get_primary_architecture(), primaryBinDirectories);

        // get the bin paths to insert
        BStringList binDirectoriesToInsert;
        if (!is_primary_architecture(architecture))
                get_bin_directories(architecture, binDirectoriesToInsert);

        // split the PATH variable
        char* pathVariableValue = getenv("PATH");
        BStringList paths;
        if (pathVariableValue != NULL
                && !BString(pathVariableValue).Split(":", true, paths)) {
                fprintf(stderr, "Error: Out of memory!\n");
                exit(1);
        }

        // Filter the paths, removing any path that isn't associated with the
        // primary architecture. Also find the insertion index for the architecture
        // bin paths.
        int32 insertionIndex = -1;
        int32 count = paths.CountStrings();
        for (int32 i = 0; i < count; i++) {
                // We always keep relative paths. Filter absolute ones only.
                const char* path = paths.StringAt(i);
                if (path[0] == '/') {
                        // try to normalize the path
                        BPath normalizedPath;
                        if (normalizedPath.SetTo(path, NULL, true) == B_OK)
                                path = normalizedPath.Path();

                        // Check, if this is a primary bin directory. If not, determine the
                        // path's architecture.
                        int32 index = primaryBinDirectories.IndexOf(path);
                        if (index >= 0) {
                                if (insertionIndex < 0)
                                        insertionIndex = _paths.CountStrings();
                        } else if (!is_primary_architecture(
                                        guess_architecture_for_path(path))) {
                                // a non-primary architecture path -- skip
                                continue;
                        }
                }

                if (!_paths.Add(paths.StringAt(i))) {
                        fprintf(stderr, "Error: Out of memory!\n");
                        exit(1);
                }
        }

        // Insert the paths for the specified architecture, if any.
        if (!binDirectoriesToInsert.IsEmpty()) {
                if (!(insertionIndex < 0
                                ? _paths.Add(binDirectoriesToInsert)
                                : _paths.Add(binDirectoriesToInsert, insertionIndex))) {
                        fprintf(stderr, "Error: Out of memory!\n");
                        exit(1);
                }
        }
}


int
main(int argc, const char* const* argv)
{
        bool printPath = false;
        bool listArchitectures = false;

        while (true) {
                static struct option sLongOptions[] = {
                        { "help", no_argument, 0, 'h' },
                        { "list-architectures", no_argument, 0, 'l' },
                        { "print-path", no_argument, 0, 'p' },
                        { 0, 0, 0, 0 }
                };

                opterr = 0; // don't print errors
                int c = getopt_long(argc, (char**)argv, "+hlp",
                        sLongOptions, NULL);
                if (c == -1)
                        break;

                switch (c) {
                        case 'h':
                                print_usage_and_exit(false);
                                break;

                        case 'l':
                                listArchitectures = true;
                                break;

                        case 'p':
                                printPath = true;
                                break;

                        default:
                                print_usage_and_exit(true);
                                break;
                }
        }

        // only one of listArchitectures, printPath may be specified
        if (listArchitectures && printPath)
                print_usage_and_exit(true);

        // get architectures
        BStringList architectures;
        status_t error = get_architectures(architectures);
        if (error != B_OK) {
                fprintf(stderr, "Error: Failed to get architectures: %s\n",
                        strerror(error));
                exit(1);
        }

        // list architectures
        if (listArchitectures) {
                if (optind != argc)
                        print_usage_and_exit(true);

                int32 count = architectures.CountStrings();
                for (int32 i = 0; i < count; i++)
                        printf("%s\n", architectures.StringAt(i).String());
                return 0;
        }

        // The remaining arguments are the architecture and optionally the command
        // to execute.
        if (optind >= argc)
                print_usage_and_exit(true);
        const char* architecture = optind < argc ? argv[optind++] : NULL;

        int commandArgCount = argc - optind;
        const char* const* commandArgs = commandArgCount > 0 ? argv + optind : NULL;

        if (printPath && commandArgs != NULL)
                print_usage_and_exit(true);

        // check the architecture
        if (!architectures.HasString(architecture)) {
                fprintf(stderr, "Error: Unsupported architecture \"%s\"\n",
                        architecture);
                exit(1);
        }

        // get the new paths
        BStringList paths;
        compute_new_paths(architecture, paths);

        BString pathVariableValue = paths.Join(":");
        if (!paths.IsEmpty() && pathVariableValue.IsEmpty())
                fprintf(stderr, "Error: Out of memory!\n");

        if (printPath) {
                printf("%s\n", pathVariableValue.String());
                return 0;
        }

        // set PATH
        if (setenv("PATH", pathVariableValue, 1) != 0) {
                fprintf(stderr, "Error: Failed to set PATH: %s\n", strerror(errno));
                exit(1);
        }

        // if no command is given, get the user's shell
        const char* shellCommand[3];
        if (commandArgs == NULL) {
                struct passwd* pwd = getpwuid(geteuid());
                shellCommand[0] = pwd != NULL ? pwd->pw_shell : "/bin/sh";
                shellCommand[1] = "-l";
                shellCommand[2] = NULL;
                commandArgs = shellCommand;
                commandArgCount = 2;
        }

        // exec the command
        execvp(commandArgs[0], (char* const*)commandArgs);

        fprintf(stderr, "Error: Executing \"%s\" failed: %s\n", commandArgs[0],
                strerror(errno));
        return 1;
}