/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2011 Nathan Whitehorn
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#include <sys/param.h>

#include <bsddialog.h>
#include <err.h>
#include <errno.h>
#include <fstab.h>
#include <inttypes.h>
#include <libgeom.h>
#include <libutil.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sysexits.h>

#include "diskmenu.h"
#include "partedit.h"

struct pmetadata_head part_metadata;
static int sade_mode = 0;

static int apply_changes(struct gmesh *mesh);
static void apply_workaround(struct gmesh *mesh);
static struct partedit_item *read_geom_mesh(struct gmesh *mesh, int *nitems);
static void add_geom_children(struct ggeom *gp, int recurse,
    struct partedit_item **items, int *nitems);
static void init_fstab_metadata(void);
static void get_mount_points(struct partedit_item *items, int nitems);
static int validate_setup(void);

static void
sigint_handler(int sig)
{
        struct gmesh mesh;

        /* Revert all changes and exit dialog-mode cleanly on SIGINT */
        if (geom_gettree(&mesh) == 0) {
                gpart_revert_all(&mesh);
                geom_deletetree(&mesh);
        }

        bsddialog_end();

        exit(1);
}

int
main(int argc, const char **argv)
{
        struct partition_metadata *md;
        const char *progname, *prompt;
        struct partedit_item *items = NULL;
        struct gmesh mesh;
        int i, op, nitems;
        int error;
        struct bsddialog_conf conf;

        progname = getprogname();
        if (strcmp(progname, "sade") == 0)
                sade_mode = 1;

        TAILQ_INIT(&part_metadata);

        init_fstab_metadata();

        if (bsddialog_init() == BSDDIALOG_ERROR)
                err(1, "%s", bsddialog_geterror());
        bsddialog_initconf(&conf);
        if (!sade_mode)
                bsddialog_backtitle(&conf, OSNAME " Installer");
        i = 0;

        /* Revert changes on SIGINT */
        signal(SIGINT, sigint_handler);

        if (strcmp(progname, "autopart") == 0) { /* Guided */
                prompt = "Please review the disk setup. When complete, press "
                    "the Finish button.";
                /* Experimental ZFS autopartition support */
                if (argc > 1 && strcmp(argv[1], "zfs") == 0) {
                        part_wizard("zfs");
                } else {
                        part_wizard("ufs");
                }
        } else if (strcmp(progname, "scriptedpart") == 0) {
                error = scripted_editor(argc, argv);
                prompt = NULL;
                if (error != 0) {
                        bsddialog_end();
                        return (error);
                }
        } else {
                prompt = "Create partitions for " OSNAME ", F1 for help.\n"
                    "No changes will be made until you select Finish.";
        }

        /* Show the part editor either immediately, or to confirm wizard */
        while (prompt != NULL) {
                bsddialog_clear(0);
                if (!sade_mode)
                        bsddialog_backtitle(&conf, OSNAME " Installer");

                error = geom_gettree(&mesh);
                if (error == 0)
                        items = read_geom_mesh(&mesh, &nitems);
                if (error || items == NULL) {
                        conf.title = "Error";
                        bsddialog_msgbox(&conf, "No disks found. If you need "
                            "to install a kernel driver, choose Shell at the "
                            "installation menu.", 0, 0);
                        break;
                }

                get_mount_points(items, nitems);

                if (i >= nitems)
                        i = nitems - 1;
                op = diskmenu_show("Partition Editor", prompt, items, nitems,
                    &i);

                switch (op) {
                case BUTTON_CREATE:
                        gpart_create((struct gprovider *)(items[i].cookie),
                            NULL, NULL, NULL, NULL, 1);
                        break;
                case BUTTON_DELETE:
                        gpart_delete((struct gprovider *)(items[i].cookie));
                        break;
                case BUTTON_MODIFY:
                        gpart_edit((struct gprovider *)(items[i].cookie));
                        break;
                case BUTTON_REVERT:
                        gpart_revert_all(&mesh);
                        while ((md = TAILQ_FIRST(&part_metadata)) != NULL) {
                                if (md->fstab != NULL) {
                                        free(md->fstab->fs_spec);
                                        free(md->fstab->fs_file);
                                        free(md->fstab->fs_vfstype);
                                        free(md->fstab->fs_mntops);
                                        free(md->fstab->fs_type);
                                        free(md->fstab);
                                }
                                if (md->newfs != NULL)
                                        free(md->newfs);
                                free(md->name);

                                TAILQ_REMOVE(&part_metadata, md, metadata);
                                free(md);
                        }
                        init_fstab_metadata();
                        break;
                case BUTTON_AUTO:
                        part_wizard("ufs");
                        break;
                }

                error = 0;
                if (op == BUTTON_FINISH) {
                        conf.button.ok_label = "Commit";
                        conf.button.with_extra = true;
                        conf.button.extra_label = "Revert & Exit";
                        conf.button.cancel_label = "Back";
                        conf.title = "Confirmation";
                        op = bsddialog_yesno(&conf, "Your changes will now be "
                            "written to disk. If you have chosen to overwrite "
                            "existing data, it will be PERMANENTLY ERASED. Are "
                            "you sure you want to commit your changes?", 0, 0);
                        conf.button.ok_label = NULL;
                        conf.button.with_extra = false;
                        conf.button.extra_label = NULL;
                        conf.button.cancel_label = NULL;

                        if (op == BSDDIALOG_OK && validate_setup()) { /* Save */
                                error = apply_changes(&mesh);
                                if (!error)
                                        apply_workaround(&mesh);
                                break;
                        } else if (op == BSDDIALOG_EXTRA) { /* Quit */
                                gpart_revert_all(&mesh);
                                error = -1;
                                break;
                        }
                }

                geom_deletetree(&mesh);
                free(items);
        }
        
        if (prompt == NULL) {
                error = geom_gettree(&mesh);
                if (error == 0) {
                        if (validate_setup()) {
                                error = apply_changes(&mesh);
                        } else {
                                gpart_revert_all(&mesh);
                                error = -1;
                        }
                        geom_deletetree(&mesh);
                }
        }

        bsddialog_end();

        return (error);
}

struct partition_metadata *
get_part_metadata(const char *name, int create)
{
        struct partition_metadata *md;

        TAILQ_FOREACH(md, &part_metadata, metadata) 
                if (md->name != NULL && strcmp(md->name, name) == 0)
                        break;

        if (md == NULL && create) {
                md = calloc(1, sizeof(*md));
                md->name = strdup(name);
                TAILQ_INSERT_TAIL(&part_metadata, md, metadata);
        }

        return (md);
}
        
void
delete_part_metadata(const char *name)
{
        struct partition_metadata *md;

        TAILQ_FOREACH(md, &part_metadata, metadata) {
                if (md->name != NULL && strcmp(md->name, name) == 0) {
                        if (md->fstab != NULL) {
                                free(md->fstab->fs_spec);
                                free(md->fstab->fs_file);
                                free(md->fstab->fs_vfstype);
                                free(md->fstab->fs_mntops);
                                free(md->fstab->fs_type);
                                free(md->fstab);
                        }
                        if (md->newfs != NULL)
                                free(md->newfs);
                        free(md->name);

                        TAILQ_REMOVE(&part_metadata, md, metadata);
                        free(md);
                        break;
                }
        }
}

static int
validate_setup(void)
{
        struct partition_metadata *md, *root = NULL;
        int button;
        struct bsddialog_conf conf;

        TAILQ_FOREACH(md, &part_metadata, metadata) {
                if (md->fstab != NULL && strcmp(md->fstab->fs_file, "/") == 0)
                        root = md;

                /* XXX: Check for duplicate mountpoints */
        }

        bsddialog_initconf(&conf);

        if (root == NULL) {
                conf.title = "Error";
                bsddialog_msgbox(&conf, "No root partition was found. "
                    "The root " OSNAME " partition must have a mountpoint "
                    "of '/'.", 0, 0);
                return (false);
        }

        /*
         * Check for root partitions that we aren't formatting, which is 
         * usually a mistake
         */
        if (root->newfs == NULL && !sade_mode) {
                conf.button.default_cancel = true;
                conf.title = "Warning";
                button = bsddialog_yesno(&conf, "The chosen root partition "
                    "has a preexisting filesystem. If it contains an existing "
                    OSNAME " system, please update it with freebsd-update "
                    "instead of installing a new system on it. The partition "
                    "can also be erased by pressing \"No\" and then deleting "
                    "and recreating it. Are you sure you want to proceed?",
                    0, 0);
                if (button == BSDDIALOG_CANCEL)
                        return (false);
        }

        return (true);
}

static int
mountpoint_sorter(const void *xa, const void *xb)
{
        struct partition_metadata *a = *(struct partition_metadata **)xa;
        struct partition_metadata *b = *(struct partition_metadata **)xb;

        if (a->fstab == NULL && b->fstab == NULL)
                return 0;
        if (a->fstab == NULL)
                return 1;
        if (b->fstab == NULL)
                return -1;

        return strcmp(a->fstab->fs_file, b->fstab->fs_file);
}

static int
apply_changes(struct gmesh *mesh)
{
        struct partition_metadata *md;
        char message[512];
        int i, nitems, error, *miniperc;
        const char **minilabel;
        const char *fstab_path;
        FILE *fstab;
        char *command;
        struct bsddialog_conf conf;

        nitems = 1; /* Partition table changes */
        TAILQ_FOREACH(md, &part_metadata, metadata) {
                if (md->newfs != NULL)
                        nitems++;
        }
        minilabel = calloc(nitems, sizeof(const char *));
        miniperc  = calloc(nitems, sizeof(int));
        minilabel[0] = "Writing partition tables";
        miniperc[0]  = BSDDIALOG_MG_INPROGRESS;
        i = 1;
        TAILQ_FOREACH(md, &part_metadata, metadata) {
                if (md->newfs != NULL) {
                        char *item;

                        asprintf(&item, "Initializing %s", md->name);
                        minilabel[i] = item;
                        miniperc[i]  = BSDDIALOG_MG_PENDING;
                        i++;
                }
        }

        i = 0;
        bsddialog_initconf(&conf);
        conf.title = "Initializing";
        bsddialog_mixedgauge(&conf,
            "Initializing file systems. Please wait.", 0, 0, i * 100 / nitems,
            nitems, minilabel, miniperc);
        gpart_commit(mesh);
        miniperc[i] = BSDDIALOG_MG_COMPLETED;
        i++;

        if (getenv("BSDINSTALL_LOG") == NULL) 
                setenv("BSDINSTALL_LOG", "/dev/null", 1);

        TAILQ_FOREACH(md, &part_metadata, metadata) {
                if (md->newfs != NULL) {
                        miniperc[i] = BSDDIALOG_MG_INPROGRESS;
                        bsddialog_mixedgauge(&conf,
                            "Initializing file systems. Please wait.", 0, 0,
                            i * 100 / nitems, nitems, minilabel, miniperc);
                        asprintf(&command, "(echo %s; %s) >>%s 2>>%s",
                            md->newfs, md->newfs, getenv("BSDINSTALL_LOG"),
                            getenv("BSDINSTALL_LOG"));
                        error = system(command);
                        free(command);
                        miniperc[i] = (error == 0) ?
                            BSDDIALOG_MG_COMPLETED : BSDDIALOG_MG_FAILED;
                        i++;
                }
        }
        bsddialog_mixedgauge(&conf, "Initializing file systems. Please wait.",
            0, 0, i * 100 / nitems, nitems, minilabel, miniperc);

        for (i = 1; i < nitems; i++)
                free(__DECONST(char *, minilabel[i]));

        free(minilabel);
        free(miniperc);

        /* Sort filesystems for fstab so that mountpoints are ordered */
        {
                struct partition_metadata **tobesorted;
                struct partition_metadata *tmp;
                int nparts = 0;
                TAILQ_FOREACH(md, &part_metadata, metadata)
                        nparts++;
                tobesorted = malloc(sizeof(struct partition_metadata *)*nparts);
                nparts = 0;
                TAILQ_FOREACH_SAFE(md, &part_metadata, metadata, tmp) {
                        tobesorted[nparts++] = md;
                        TAILQ_REMOVE(&part_metadata, md, metadata);
                }
                qsort(tobesorted, nparts, sizeof(tobesorted[0]),
                    mountpoint_sorter);

                /* Now re-add everything */
                while (nparts-- > 0)
                        TAILQ_INSERT_HEAD(&part_metadata,
                            tobesorted[nparts], metadata);
                free(tobesorted);
        }

        if (getenv("PATH_FSTAB") != NULL)
                fstab_path = getenv("PATH_FSTAB");
        else
                fstab_path = "/etc/fstab";
        fstab = fopen(fstab_path, "w+");
        if (fstab == NULL) {
                snprintf(message, sizeof(message),
                    "Cannot open fstab file %s for writing (%s)\n",
                    getenv("PATH_FSTAB"), strerror(errno));
                conf.title = "Error";
                bsddialog_msgbox(&conf, message, 0, 0);
                return (-1);
        }
        fprintf(fstab, "# Device\tMountpoint\tFStype\tOptions\tDump\tPass#\n");
        TAILQ_FOREACH(md, &part_metadata, metadata) {
                if (md->fstab != NULL)
                        fprintf(fstab, "%s\t%s\t\t%s\t%s\t%d\t%d\n",
                            md->fstab->fs_spec, md->fstab->fs_file,
                            md->fstab->fs_vfstype, md->fstab->fs_mntops,
                            md->fstab->fs_freq, md->fstab->fs_passno);
        }
        fclose(fstab);

        return (0);
}

static void
apply_workaround(struct gmesh *mesh)
{
        struct gclass *classp;
        struct ggeom *gp;
        struct gconfig *gc;
        const char *scheme = NULL, *modified = NULL;
        struct bsddialog_conf conf;

        LIST_FOREACH(classp, &mesh->lg_class, lg_class) {
                if (strcmp(classp->lg_name, "PART") == 0)
                        break;
        }

        if (strcmp(classp->lg_name, "PART") != 0) {
                bsddialog_initconf(&conf);
                conf.title = "Error";
                bsddialog_msgbox(&conf, "gpart not found!", 0, 0);
                return;
        }

        LIST_FOREACH(gp, &classp->lg_geom, lg_geom) {
                LIST_FOREACH(gc, &gp->lg_config, lg_config) {
                        if (strcmp(gc->lg_name, "scheme") == 0) {
                                scheme = gc->lg_val;
                        } else if (strcmp(gc->lg_name, "modified") == 0) {
                                modified = gc->lg_val;
                        }
                }

                if (scheme && strcmp(scheme, "GPT") == 0 &&
                    modified && strcmp(modified, "true") == 0) {
                        if (getenv("WORKAROUND_LENOVO"))
                                gpart_set_root(gp->lg_name, "lenovofix");
                        if (getenv("WORKAROUND_GPTACTIVE"))
                                gpart_set_root(gp->lg_name, "active");
                }
        }
}

static struct partedit_item *
read_geom_mesh(struct gmesh *mesh, int *nitems)
{
        struct gclass *classp;
        struct ggeom *gp;
        struct partedit_item *items;

        *nitems = 0;
        items = NULL;

        /*
         * Build the device table. First add all disks (and CDs).
         */
        
        LIST_FOREACH(classp, &mesh->lg_class, lg_class) {
                if (strcmp(classp->lg_name, "DISK") != 0 &&
                    strcmp(classp->lg_name, "MD") != 0)
                        continue;

                /* Now recurse into all children */
                LIST_FOREACH(gp, &classp->lg_geom, lg_geom) 
                        add_geom_children(gp, 0, &items, nitems);
        }

        return (items);
}

static void
add_geom_children(struct ggeom *gp, int recurse, struct partedit_item **items,
    int *nitems)
{
        struct gconsumer *cp;
        struct gprovider *pp;
        struct gconfig *gc;

        if (strcmp(gp->lg_class->lg_name, "PART") == 0 &&
            !LIST_EMPTY(&gp->lg_config)) {
                LIST_FOREACH(gc, &gp->lg_config, lg_config) {
                        if (strcmp(gc->lg_name, "scheme") == 0)
                                (*items)[*nitems-1].type = gc->lg_val;
                }
        }

        if (LIST_EMPTY(&gp->lg_provider)) 
                return;

        LIST_FOREACH(pp, &gp->lg_provider, lg_provider) {
                if (strcmp(gp->lg_class->lg_name, "LABEL") == 0)
                        continue;

                /* Skip WORM media */
                if (strncmp(pp->lg_name, "cd", 2) == 0)
                        continue;

                *items = realloc(*items,
                    (*nitems+1)*sizeof(struct partedit_item));
                (*items)[*nitems].indentation = recurse;
                (*items)[*nitems].name = pp->lg_name;
                (*items)[*nitems].size = pp->lg_mediasize;
                (*items)[*nitems].mountpoint = NULL;
                (*items)[*nitems].type = "";
                (*items)[*nitems].cookie = pp;

                LIST_FOREACH(gc, &pp->lg_config, lg_config) {
                        if (strcmp(gc->lg_name, "type") == 0)
                                (*items)[*nitems].type = gc->lg_val;
                }

                /* Skip swap-backed MD devices */
                if (strcmp(gp->lg_class->lg_name, "MD") == 0 &&
                    strcmp((*items)[*nitems].type, "swap") == 0)
                        continue;

                (*nitems)++;

                LIST_FOREACH(cp, &pp->lg_consumers, lg_consumers)
                        add_geom_children(cp->lg_geom, recurse+1, items,
                            nitems);

                /* Only use first provider for acd */
                if (strcmp(gp->lg_class->lg_name, "ACD") == 0)
                        break;
        }
}

static void
init_fstab_metadata(void)
{
        struct fstab *fstab;
        struct partition_metadata *md;

        setfsent();
        while ((fstab = getfsent()) != NULL) {
                md = calloc(1, sizeof(struct partition_metadata));

                md->name = NULL;
                if (strncmp(fstab->fs_spec, "/dev/", 5) == 0)
                        md->name = strdup(&fstab->fs_spec[5]);

                md->fstab = malloc(sizeof(struct fstab));
                md->fstab->fs_spec = strdup(fstab->fs_spec);
                md->fstab->fs_file = strdup(fstab->fs_file);
                md->fstab->fs_vfstype = strdup(fstab->fs_vfstype);
                md->fstab->fs_mntops = strdup(fstab->fs_mntops);
                md->fstab->fs_type = strdup(fstab->fs_type);
                md->fstab->fs_freq = fstab->fs_freq;
                md->fstab->fs_passno = fstab->fs_passno;

                md->newfs = NULL;
                
                TAILQ_INSERT_TAIL(&part_metadata, md, metadata);
        }
}

static void
get_mount_points(struct partedit_item *items, int nitems)
{
        struct partition_metadata *md;
        int i;
        
        for (i = 0; i < nitems; i++) {
                TAILQ_FOREACH(md, &part_metadata, metadata) {
                        if (md->name != NULL && md->fstab != NULL &&
                            strcmp(md->name, items[i].name) == 0) {
                                items[i].mountpoint = md->fstab->fs_file;
                                break;
                        }
                }
        }
}