/*
 * Copyright 2005-2007 Ingo Weinhold, bonefish@users.sf.net
 * Copyright 2005-2013 Axel Dörfler, axeld@pinc-software.de
 * Copyright 2009 Jonas Sundström, jonas@kirilla.se
 *
 * All rights reserved. Distributed under the terms of the MIT License.
 */


#include <set>
#include <string>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <termios.h>

#include <Application.h>
#include <Path.h>
#include <String.h>
#include <fs_volume.h>

#include <DiskDevice.h>
#include <DiskDevicePrivate.h>
#include <DiskDeviceRoster.h>
#include <DiskDeviceTypes.h>
#include <DiskDeviceList.h>
#include <Partition.h>

#include <tracker_private.h>


using std::set;
using std::string;

extern const char* __progname;


typedef set<string> StringSet;

// usage
static const char* kUsage =
        "Usage: %s <options> [ <volume name> ... ]\n\n"
        "Mounts the volume with name <volume name>, if given. Lists info about\n"
        "mounted and mountable volumes and mounts/unmounts volumes.\n"
        "\n"
        "The terminology is actually not quite correct: By volumes only partitions\n"
        "living on disk devices are meant.\n"
        "\n"
        "Options:\n"
        "[general]\n"
        "  -s                    - silent; don't print info about (un)mounting\n"
        "  -h, --help            - print this info text\n"
        "\n"
        "[mounting]\n"
        "  -all                  - mount all mountable volumes\n"
        "  -allbfs               - mount all mountable BFS volumes\n"
        "  -allhfs               - mount all mountable HFS volumes\n"
        "  -alldos               - mount all mountable DOS volumes\n"
        "  -ro, -readonly        - mount volumes read-only\n"
        "  -u, -unmount <volume> - unmount the volume with the name <volume>\n"
        "  -open                 - opens the mounted volumes in Tracker\n"
        "\n"
        "[info]\n"
        "  -p, -l                - list all mounted and mountable volumes\n"
        "  -lh                   - list all existing volumes (incl. not-mountable "
                "ones)\n"
        "  -dd                   - list all disk existing devices\n"
        "\n"
        "[obsolete]\n"
        "  -r                    - ignored\n"
        "  -publishall           - ignored\n"
        "  -publishbfs           - ignored\n"
        "  -publishhfs           - ignored\n"
        "  -publishdos           - ignored\n";


const char* kAppName = __progname;

static int sVolumeNameWidth = B_OS_NAME_LENGTH;
static int sFSNameWidth = 25;


static void
print_usage(bool error)
{
        fprintf(error ? stderr : stdout, kUsage, kAppName);
}


static void
print_usage_and_exit(bool error)
{
        print_usage(error);
        exit(error ? 0 : 1);
}


static const char*
size_string(int64 size)
{
        double blocks = size;
        static char string[64];

        if (size < 1024)
                sprintf(string, "%" B_PRId64, size);
        else {
                const char* units[] = {"K", "M", "G", NULL};
                int32 i = -1;

                do {
                        blocks /= 1024.0;
                        i++;
                } while (blocks >= 1024 && units[i + 1]);

                snprintf(string, sizeof(string), "%.1f%s", blocks, units[i]);
        }

        return string;
}


static status_t
open_in_tracker(BPartition* partition)
{
        BPath mountPoint;
        status_t status = partition->GetMountPoint(&mountPoint);
        if (status != B_OK)
                return status;

        entry_ref ref;
        status = get_ref_for_path(mountPoint.Path(), &ref);
        if (status != B_OK)
                return status;

        BMessage refs(B_REFS_RECEIVED);
        refs.AddRef("refs", &ref);
        return BMessenger(kTrackerSignature).SendMessage(&refs);
}


//      #pragma mark -


struct MountVisitor : public BDiskDeviceVisitor {
        MountVisitor()
                :
                silent(false),
                mountAll(false),
                mountBFS(false),
                mountHFS(false),
                mountDOS(false),
                readOnly(false),
                openInTracker(false)
        {
        }

        virtual bool Visit(BDiskDevice* device)
        {
                return Visit(device, 0);
        }

        virtual bool Visit(BPartition* partition, int32 level)
        {
                // get name and type
                BString name = partition->ContentName();
                if (name.IsEmpty())
                        name = partition->Name();
                const char* type = partition->ContentType();

                // check whether to mount
                bool mount = false;
                if (name && toMount.find(name.String()) != toMount.end()) {
                        toMount.erase(name.String());
                        if (!partition->IsMounted())
                                mount = true;
                        else if (!silent)
                                fprintf(stderr, "Volume `%s' already mounted.\n", name.String());
                } else if (mountAll) {
                        mount = true;
                } else if (mountBFS && type != NULL
                        && strcmp(type, kPartitionTypeBFS) == 0) {
                        mount = true;
                } else if (mountHFS && type != NULL
                        && strcmp(type, kPartitionTypeHFS) == 0) {
                        mount = true;
                } else if (mountDOS && type != NULL
                        && (strcmp(type, kPartitionTypeFAT12) == 0
                                || strcmp(type, kPartitionTypeFAT32) == 0)) {
                        mount = true;
                }

                // don't try to mount a partition twice
                if (partition->IsMounted())
                        mount = false;

                // check whether to unmount
                bool unmount = false;
                if (name && toUnmount.find(name.String()) != toUnmount.end()) {
                        toUnmount.erase(name.String());
                        if (partition->IsMounted()) {
                                unmount = true;
                                mount = false;
                        } else if (!silent)
                                fprintf(stderr, "Volume `%s' not mounted.\n", name.String());
                }

                // mount/unmount
                if (mount) {
                        status_t error = partition->Mount(NULL,
                                readOnly ? B_MOUNT_READ_ONLY : 0);
                        if (!silent) {
                                if (error >= B_OK) {
                                        BPath mountPoint;
                                        partition->GetMountPoint(&mountPoint);
                                        printf("Volume `%s' mounted successfully at '%s'.\n", name.String(),
                                                mountPoint.Path());
                                } else {
                                        fprintf(stderr, "Failed to mount volume `%s': %s\n",
                                                name.String(), strerror(error));
                                }
                        }
                        if (openInTracker && error == B_OK)
                                open_in_tracker(partition);
                } else if (unmount) {
                        status_t error = partition->Unmount();
                        if (!silent) {
                                if (error == B_OK) {
                                        printf("Volume `%s' unmounted successfully.\n", name.String());
                                } else {
                                        fprintf(stderr, "Failed to unmount volume `%s': %s\n",
                                                name.String(), strerror(error));
                                }
                        }
                }

                return false;
        }

        bool            silent;
        StringSet       toMount;
        StringSet       toUnmount;
        bool            mountAll;
        bool            mountBFS;
        bool            mountHFS;
        bool            mountDOS;
        bool            readOnly;
        bool            openInTracker;
};


struct PrintPartitionsVisitor : public BDiskDeviceVisitor {
        PrintPartitionsVisitor()
                : listMountablePartitions(false),
                  listAllPartitions(false)
        {
        }

        bool IsUsed()
        {
                return listMountablePartitions || listAllPartitions;
        }

        virtual bool Visit(BDiskDevice* device)
        {
                return Visit(device, 0);
        }

        virtual bool Visit(BPartition* partition, int32 level)
        {
                // get name and type
                BString name = partition->ContentName();
                if (name.IsEmpty()) {
                        name = partition->Name();
                        if (name.IsEmpty()) {
                                if (partition->ContainsFileSystem())
                                        name = "<unnamed>";
                                else
                                        name = "";
                        }
                }
                const char* type = partition->ContentType();
                if (type == NULL)
                        type = "<unknown>";

                // shorten known types for display
                if (!strcmp(type, kPartitionTypeMultisession))
                        type = "Multisession";
                else if (!strcmp(type, kPartitionTypeIntelExtended))
                        type = "Intel Extended";

                BPath path;
                partition->GetPath(&path);

                // cut off beginning of the device path (if /dev/disk/)
                int32 skip = strlen("/dev/disk/");
                if (strncmp(path.Path(), "/dev/disk/", skip))
                        skip = 0;

                BPath mountPoint;
                if (partition->IsMounted())
                        partition->GetMountPoint(&mountPoint);

                printf("%-*s %-*s %8s %s%s(%s)\n", sVolumeNameWidth, name.String(),
                        sFSNameWidth, type, size_string(partition->Size()),
                        partition->IsMounted() ? mountPoint.Path() : "",
                        partition->IsMounted() ? "  " : "",
                        path.Path() + skip);
                return false;
        }

        bool listMountablePartitions;
        bool listAllPartitions;
};


//      #pragma mark -


class MountVolume : public BApplication {
public:
                                                MountVolume();
        virtual                         ~MountVolume();

        virtual void            RefsReceived(BMessage* message);
        virtual void            ArgvReceived(int32 argc, char** argv);
        virtual void            ReadyToRun();
};


MountVolume::MountVolume()
        :
        BApplication("application/x-vnd.haiku-mountvolume")
{
}


MountVolume::~MountVolume()
{
}


void
MountVolume::RefsReceived(BMessage* message)
{
        status_t status;
        int32 refCount;
        type_code typeFound;

        status = message->GetInfo("refs", &typeFound, &refCount);
        if (status != B_OK || refCount < 1) {
                fprintf(stderr, "Failed to get info from entry_refs BMessage: %s\n",
                        strerror(status));
                exit(1);
        }

        entry_ref ref;
        BPath path;

        int32 argc = refCount + 2;
        char** argv = new char*[argc + 1];
        argv[0] = strdup(kAppName);
        argv[1] = strdup("-open");

        for (int32 i = 0; i < refCount; i++) {
                message->FindRef("refs", i, &ref);
                status = path.SetTo(&ref);
                if (status != B_OK) {
                        fprintf(stderr, "Failed to get a path (%s) from entry (%s): %s\n",
                                path.Path(), ref.name, strerror(status));
                }
                argv[2 + i] = strdup(path.Path());
        }
        argv[argc] = NULL;

        ArgvReceived(argc, argv);
}


void
MountVolume::ArgvReceived(int32 argc, char** argv)
{
        MountVisitor mountVisitor;
        PrintPartitionsVisitor printPartitionsVisitor;
        bool listAllDevices = false;

        if (argc < 2)
                printPartitionsVisitor.listMountablePartitions = true;

        // parse arguments

        for (int argi = 1; argi < argc; argi++) {
                const char* arg = argv[argi];

                if (arg[0] != '\0' && arg[0] != '-') {
                        mountVisitor.toMount.insert(arg);
                } else if (strcmp(arg, "-s") == 0) {
                        mountVisitor.silent = true;
                } else if (strcmp(arg, "-h") == 0 || strcmp(arg, "--help") == 0) {
                        print_usage_and_exit(false);
                } else if (strcmp(arg, "-all") == 0) {
                        mountVisitor.mountAll = true;
                } else if (strcmp(arg, "-allbfs") == 0) {
                        mountVisitor.mountBFS = true;
                } else if (strcmp(arg, "-allhfs") == 0) {
                        mountVisitor.mountHFS = true;
                } else if (strcmp(arg, "-alldos") == 0) {
                        mountVisitor.mountDOS = true;
                } else if (strcmp(arg, "-ro") == 0 || strcmp(arg, "-readonly") == 0) {
                        mountVisitor.readOnly = true;
                } else if (strcmp(arg, "-u") == 0 || strcmp(arg, "-unmount") == 0) {
                        argi++;
                        if (argi >= argc)
                                print_usage_and_exit(true);
                        mountVisitor.toUnmount.insert(argv[argi]);
                } else if (strcmp(arg, "-open") == 0) {
                        mountVisitor.openInTracker = true;
                } else if (strcmp(arg, "-p") == 0 || strcmp(arg, "-l") == 0) {
                        printPartitionsVisitor.listMountablePartitions = true;
                } else if (strcmp(arg, "-lh") == 0) {
                        printPartitionsVisitor.listAllPartitions = true;
                } else if (strcmp(arg, "-dd") == 0) {
                        listAllDevices = true;
                } else if (strcmp(arg, "-r") == 0 || strcmp(arg, "-publishall") == 0
                        || strcmp(arg, "-publishbfs") == 0
                        || strcmp(arg, "-publishhfs") == 0
                        || strcmp(arg, "-publishdos") == 0) {
                        // obsolete: ignore
                } else
                        print_usage_and_exit(true);
        }

        // get a disk device list
        BDiskDeviceList deviceList;
        status_t error = deviceList.Fetch();
        if (error != B_OK) {
                fprintf(stderr, "Failed to get the list of disk devices: %s",
                        strerror(error));
                exit(1);
        }

        // mount/unmount volumes
        deviceList.VisitEachMountablePartition(&mountVisitor);

        BDiskDeviceRoster roster;

        // try mount file images
        for (StringSet::iterator iterator = mountVisitor.toMount.begin();
                        iterator != mountVisitor.toMount.end();) {
                const char* name = (*iterator).c_str();
                iterator++;

                BEntry entry(name, true);
                if (!entry.Exists())
                        continue;

                // TODO: improve error messages
                BPath path;
                if (entry.GetPath(&path) != B_OK)
                        continue;

                partition_id id = -1;
                BDiskDevice device;
                BPartition* partition;

                if (!strncmp(path.Path(), "/dev/", 5)) {
                        // seems to be a device path
                        if (roster.GetPartitionForPath(path.Path(), &device, &partition)
                                        != B_OK)
                                continue;
                } else {
                        // a file with this name exists, so try to mount it
                        id = roster.RegisterFileDevice(path.Path());
                        if (id < 0)
                                continue;

                        if (roster.GetPartitionWithID(id, &device, &partition) != B_OK) {
                                roster.UnregisterFileDevice(id);
                                continue;
                        }
                }

                status_t status = partition->Mount(NULL,
                        mountVisitor.readOnly ? B_MOUNT_READ_ONLY : 0);
                if (!mountVisitor.silent) {
                        if (status >= B_OK) {
                                BPath mountPoint;
                                partition->GetMountPoint(&mountPoint);
                                printf("%s \"%s\" mounted successfully at \"%s\".\n",
                                        id < 0 ? "Device" : "Image", name, mountPoint.Path());
                        }
                }
                if (status >= B_OK) {
                        if (mountVisitor.openInTracker)
                                open_in_tracker(partition);

                        // remove from list
                        mountVisitor.toMount.erase(name);
                } else if (id >= 0)
                        roster.UnregisterFileDevice(id);
        }

        // TODO: support unmounting images by path!

        // print errors for the volumes to mount/unmount, that weren't found
        if (!mountVisitor.silent) {
                for (StringSet::iterator it = mountVisitor.toMount.begin();
                                it != mountVisitor.toMount.end(); it++) {
                        fprintf(stderr, "Failed to mount volume `%s': Volume not found.\n",
                                (*it).c_str());
                }
                for (StringSet::iterator it = mountVisitor.toUnmount.begin();
                                it != mountVisitor.toUnmount.end(); it++) {
                        fprintf(stderr, "Failed to unmount volume `%s': Volume not "
                                "found.\n", (*it).c_str());
                }
        }

        // update the disk device list
        error = deviceList.Fetch();
        if (error != B_OK) {
                fprintf(stderr, "Failed to update the list of disk devices: %s",
                        strerror(error));
                exit(1);
        }

        // print information

        if (listAllDevices) {
                // TODO
        }

        // determine width of the terminal in order to shrink the columns if needed
        if (isatty(STDOUT_FILENO)) {
                winsize size;
                if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &size, sizeof(winsize)) == 0) {
                        if (size.ws_col < 95) {
                                sVolumeNameWidth -= (95 - size.ws_col) / 2;
                                sFSNameWidth -= (95 - size.ws_col) / 2;
                        }
                }
        }

        if (printPartitionsVisitor.IsUsed()) {
                printf("%-*s %-*s     Size Mounted At (Device)\n",
                        sVolumeNameWidth, "Volume", sFSNameWidth, "File System");
                BString separator;
                separator.SetTo('-', sVolumeNameWidth + sFSNameWidth + 35);
                puts(separator.String());

                if (printPartitionsVisitor.listAllPartitions)
                        deviceList.VisitEachPartition(&printPartitionsVisitor);
                else
                        deviceList.VisitEachMountablePartition(&printPartitionsVisitor);
        }

        exit(0);
}


void
MountVolume::ReadyToRun()
{
        // We will only get here if we were launched without any arguments or
        // startup messages

        extern int __libc_argc;
        extern char** __libc_argv;

        ArgvReceived(__libc_argc, __libc_argv);
}


//      #pragma mark -


int
main()
{
        MountVolume mountVolume;
        mountVolume.Run();
        return 0;
}