/*
 * Copyright (c)2004 The DragonFly Project.  All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 
 *   Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 * 
 *   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.
 * 
 *   Neither the name of the DragonFly Project nor the names of its
 *   contributors may be used to endorse or promote products derived
 *   from this software without specific prior written permission. 
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS 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
 * COPYRIGHT HOLDERS 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. 
 */

/*
 * diskutilng.c
 * Disk utility functions for installer.
 * $Id: diskutilng.c,v 1.1 2004/07/09 23:25:40 cpressey Exp $
 */

#include <sys/stat.h>
#include <sys/param.h>
#include <sys/ucred.h>
#include <sys/mount.h>
#include <sys/disklabel.h>

#include <machine/pc/msdos.h>

#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <disks.h>

#include <dfui/fspred.h>

#include "pathnames.h"
#include "diskutil.h"
#include "commands.h"
#include "functions.h"

/** STORAGE DESCRIPTORS **/

struct system_storage *
system_storage_new(void)
{
	struct system_storage *s;

	if ((s = malloc(sizeof(struct system_storage))) == NULL)
		return(NULL);

	s->disk_head = NULL;
	s->disk_tail = NULL;
	s->selected_disk = NULL;
	s->selected_slice = NULL;
	s->ram = -1;

	return(s);
}

void
system_storage_free(struct system_storage *s)
{
	disks_free(s);
	free(s);
}

char *unusedisks[] = {
	"md",
	"cd",
	NULL
};

/*
 * Create a new disk description structure.
 */
struct disk *
disk_new(struct system_storage *s, char *desc)
{
	struct disklabel *dl;
	struct disk *d;
	char **disks;

	for (disks = unusedisks; *disks != NULL; disks++) {
		if (strncmp(*disks, desc, strlen(*disks)) == 0)
			return(NULL);
	}

	if ((d = malloc(sizeof(struct disk))) == NULL)
		return(NULL);

	/* XXX: eirik */
	d->number = desc[2] - '0';
	d->device = strdup(desc);
	if (d->device == NULL) {
		free(d);
		return(NULL);
	}
	dl = read_disklabel(desc);
	if (!dl) {
		asprintf(&d->desc, "%s [???M]", desc);
		d->capacity = 0;
	} else {
		asprintf(&d->desc, "%s [%uM]", desc, DLSIZEM(dl));
		d->capacity = DLSIZEM(dl);
		free_disklabel(dl);
	}
	d->we_formatted = 0;

	d->slice_head = NULL;
	d->slice_tail = NULL;

	d->next = NULL;
	if (s->disk_head == NULL) {
		s->disk_head = d;
	} else {
		s->disk_tail->next = d;
	}
	d->prev = s->disk_tail;
	s->disk_tail = d;

	return(d);
}

/*
 * Find a disk description structure from the device name.
 */
struct disk *
disk_find(struct system_storage *s, char *device)
{
	struct disk *d = s->disk_head;

	while (d != NULL) {
		if (!strcmp(device, d->device))
			return(d);
		d = d->next;
	}
	
	return(NULL);
}

/*
 * Free the memory allocated to hold the set of disk descriptions.
 */
void
disks_free(struct system_storage *s)
{
	struct disk *d = s->disk_head, *next;

	while (d != NULL) {
		next = d->next;
		slices_free(d->slice_head);
		free(d->desc);
		free(d->device);
		free(d);
		d = next;
	}

	s->disk_head = NULL;
	s->disk_tail = NULL;
}

/*
 * Look through the dmesg log created at boot-time to find out
 * what IDE and SCSI disks are attached to this computer.
 */
void
discover_disks(struct system_storage *s)
{
	struct disk *d;
	char **disks, **curdisk;

	disks_free(s);

	disks = list_disks();
	for (curdisk = disks; *curdisk != NULL; curdisk++) {
		if ((d = disk_new(s, *curdisk)) != NULL)
			discover_slices(d);
	}
	free_disks(disks);
}

void
discover_storage(struct system_storage *ss)
{
	measure_physical_memory(ss);
	discover_disks(ss);
}

/*
 * Create a new slice description and add it to a disk description.
 * The string `desc' must be allocated prior to calling this function.
 */
struct slice *
slice_new(struct disk *d, int number, char *desc, long capacity)
{
	struct slice *s;

#ifdef DEBUG
	printf("** adding slice %d `%s' to disk %s\n", number, desc, 
	    d->device);
#endif
	
	if ((s = malloc(sizeof(struct slice))) == NULL)
		return(NULL);

	s->subpartition_head = NULL;
	s->subpartition_tail = NULL;
	s->number = number;
	s->desc = desc;
	s->capacity = capacity;

	s->next = NULL;
	if (d->slice_head == NULL) {
		d->slice_head = s;
	} else {
		d->slice_tail->next = s;
	}
	s->prev = d->slice_tail;
	d->slice_tail = s;

	return(s);
}

/*
 * Find a slice description on a given disk description given the
 * slice number.
 */
struct slice *
slice_find(struct disk *d, int number)
{
	struct slice *s = d->slice_head;

	while (s != NULL) {
		if (s->number == number)
			return(s);
		s = s->next;
	}
	
	return(NULL);
}

/*
 * Free all memory for a list of slice descriptions.
 */
void
slices_free(struct slice *head)
{
	struct slice *next;

	while (head != NULL) {
		next = head->next;
		subpartitions_free(head);
		free(head->desc);
		free(head);
		head = next;
	}
}

#define PREFIX "The data for partition "
#define UNUSED "<UNUSED>"

#define OUT_OF_MEMORY "Out of memory"
/*
 * Figure out what slices are available on a given disk by asking fdisk.
 */
void
discover_slices(struct disk *d)
{
	struct dos_partition parts[4];
	int i;

#ifdef DEBUG
	printf("** discovering primary partitions for %s\n", d->desc);
#endif

	/*
	 * Create device entry before trying to use fdisk on it.
	 */
	maybe_makedev(NULL, d->device);
	if (get_slices(d->device, parts) < 0)
		return;
	for (i = 0; i < 4; i++) {
		char *str;
		uint32_t secsize;
		if (parts[i].dp_size == 0)
			continue;
		secsize = get_sectorsize(d->device);
		if (asprintf(&str, "%ss%d [%luM]  sysid: 0x%x", d->device, i,
		             (secsize * parts[i].dp_size / 1024 / 1024),
		             parts[i].dp_typ)
		    == -1)
			str = OUT_OF_MEMORY;
		slice_new(d, i + 1, str, 
		          secsize * parts[i].dp_size / 1024 / 1024);
	}
}

int
subpartition_check(struct slice *s, char *mountpoint, char *capstring,
		   int softupdates, long fsize, long bsize, char **error)
{
	char unit;
	long capacity = 0;

	if (!(!strcasecmp(mountpoint, "swap") || mountpoint[0] == '/')) {
		*error = "Mount point must be either 'swap', or it must "
			 "start with a '/'.";
		return(0);
	}

	if (strpbrk(mountpoint, " \\\"'`") != NULL) {
		*error = "Mount point may not contain the following "
			 "characters: blank space, backslash, or "
			 "single, double, or back quotes.";
		return(0);
	}
	
	if (strlen(capstring) == 0) {
		*error = "A capacity must be specified.";
		return(0);
	}

	unit = capstring[strlen(capstring) - 1];
	if (!strcmp(capstring, "*")) {
		capacity = -1;
	} else if (unit == 'm' || unit == 'M') {
		capacity = atol(capstring);
	} else if (unit == 'g' || unit == 'G') {
		capacity = atol(capstring) * 1024;
	} else {
		*error = "Capacity must be either a '*' symbol to indicate "
			 "'use the rest of the primary partition', or it "
			 "must be a series of decimal digits ending with a "
			 "'M' (indicating megabytes) or a 'G' (indicating "
			 "gigabytes.)";
		return(0);
	}

	/*
	 * If we made it through that obstacle course, all is well.
	 */

	return(1);
}

/*
 * To be safe, subpartition_check() should be called first on the
 * same arguments.
 *
 * fsize and/or bsize may both be -1, indicating
 * "choose a reasonable default."
 */
struct subpartition *
subpartition_new(struct slice *s, char *mountpoint, char *capstring,
		 int softupdates, long fsize, long bsize)
{
	struct subpartition *sp;
	long capacity = 0;
	char unit;

	unit = capstring[strlen(capstring) - 1];
	if (!strcmp(capstring, "*")) {
		capacity = -1;
	} else if (unit == 'm' || unit == 'M') {
		capacity = atol(capstring);
	} else if (unit == 'g' || unit == 'G') {
		capacity = atol(capstring) * 1024;
	}

	if ((sp = malloc(sizeof(struct subpartition))) == NULL) {
		return(NULL);
	}

	if (s->subpartition_tail == NULL)
		sp->letter = 'a';
	else if (s->subpartition_tail->letter == 'b')
		sp->letter = 'd';
	else
		sp->letter = (char)(s->subpartition_tail->letter + 1);

	sp->mountpoint = strdup(mountpoint);
	sp->capstring = strdup(capstring);
	sp->capacity = capacity;

	if (fsize == -1) {
		if (sp->capacity < 1024) {
			sp->fsize = 1024;
		} else {
			sp->fsize = 4096;
		}
	} else {
		sp->fsize = fsize;
	}

	if (bsize == -1) {
		if (sp->capacity < 1024) {
			sp->bsize = 8192;
		} else {
			sp->bsize = 16384;
		}
	} else {
		sp->bsize = bsize;
	}

	if (softupdates == -1) {
		if (strcmp(mountpoint, "/") == 0) {
			sp->softupdates = 0;
		} else {
			sp->softupdates = 1;
		}
	} else {
		sp->softupdates = softupdates;
	}

	sp->is_swap = 0;
	if (strcasecmp(mountpoint, "swap") == 0)
		sp->is_swap = 1;

	sp->next = NULL;
	if (s->subpartition_head == NULL) {
		s->subpartition_head = sp;
	} else {
		s->subpartition_tail->next = sp;
	}
	sp->prev = s->subpartition_tail;
	s->subpartition_tail = sp;

	return(sp);
}

struct subpartition *
subpartition_find(struct slice *s, char *fmt, ...)
{
	struct subpartition *sp = s->subpartition_head;
	char *mountpoint;
	va_list args;

	va_start(args, fmt);
	vasprintf(&mountpoint, fmt, args);
	va_end(args);

	while (sp != NULL) {
		if (!strcmp(mountpoint, sp->mountpoint)) {
			free(mountpoint);
			return(sp);
		}
		sp = sp->next;
	}

	free(mountpoint);
	return(NULL);
}

struct subpartition *
subpartition_find_capacity(struct slice *s, long capacity)
{
	struct subpartition *sp = s->subpartition_head;

	while (sp != NULL) {
		if (sp->capacity == capacity)
			return(sp);
		sp = sp->next;
	}

	return(NULL);
}

int
subpartition_count(struct slice *s)
{
	struct subpartition *sp = s->subpartition_head;
	int count = 0;

	while (sp != NULL) {
		count++;
		sp = sp->next;
	}

	return(count);
}

void
subpartitions_free(struct slice *s)
{
	struct subpartition *sp = s->subpartition_head, *next;

	while (sp != NULL) {
		next = sp->next;
		free(sp->mountpoint);
		free(sp->capstring);
		free(sp);
		sp = next;
	}

	s->subpartition_head = NULL;
	s->subpartition_tail = NULL;
}

/*
 * Utility functions.
 */

/*
 * Round a number up to the nearest power of two.
 */
long
next_power_of_two(long n)
{
	long p;

	for (p = 1; p < n; p <<= 1)
		;

	return(p);
}

/*
 * Look through the dmesg log created at boot-time to find out
 * how much physical memory this machine has.  Return it in MB.
 */
long
measure_physical_memory(struct system_storage *ss)
{
	FILE *p;
	char line[256];
	char *mems;

	p = popen("egrep '^real memory.*=' " DMESG_BOOT, "r");
	while (!feof(p)) {
		fgets(line, 255, p);
		if (feof(p))
			break;
		for (mems = line; *mems != '=' && *mems != '\0'; mems++)
			;
		if (*mems == '=') {
			ss->ram = next_power_of_two(atol(++mems)) /
			    (1024 * 1024);
		}
	}
	pclose(p);

	return(ss->ram);
}

long
measure_activated_swap(void)
{
	FILE *p;
	char line[256];
	char *word;
	long swap = 0;

	p = popen("swapinfo -k", "r");
	while (!feof(p)) {
		fgets(line, 255, p);
		if (feof(p))
			break;
		word = strtok(line, " \t");
		if (strcmp(word, "Device") == 0)
			continue;
		word = strtok(NULL, " \t");
		swap += atol(word);
	}
	pclose(p);

	return(swap / 1024);
}

long
measure_activated_swap_from_slice(struct disk *d, struct slice *s)
{
	FILE *p;
	char line[256], dev[256];
	char *word;
	long swap = 0;

	snprintf(dev, 256, "/dev/%ss%d", d->device, s->number);
	p = popen("swapinfo -k", "r");
	while (!feof(p)) {
		fgets(line, 255, p);
		if (feof(p))
			break;
		word = strtok(line, " \t");
		if (strcmp(word, "Device") == 0)
			continue;
		if (strstr(word, dev) != word)
			continue;
		word = strtok(NULL, " \t");
		swap += atol(word);
	}
	pclose(p);

	return(swap / 1024);
}

char *
capstring(int megs)
{
	static char string[256];

	if (megs == -1)
		strlcpy(string, "*", 2);
	else
		snprintf(string, 255, "%dM", megs);

	return(string);
}
