/*	$NetBSD$	*/

/* XXX license */

#include <sys/param.h>
#include <sys/disklabel.h>
#include <sys/ioctl.h>
#include <sys/stat.h>

#include <machine/disklabel.h>

#include <ctype.h>
#include <err.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <util.h>

int	getpart(char *cmd);
void	printmap(struct part_map_entry *pme);
void	printpart(struct part_map_entry *pme, int part);
void	addpart(struct part_map_entry *pme);
void	delpart(struct part_map_entry *pme, int part);
int	editpart(struct part_map_entry *pme, int part);
int	checkpart(struct part_map_entry *pme, int part,
		  u_int32_t start, u_int32_t size, int subjunctive);
void	writemap(struct part_map_entry *pme, char *path);
void	init_disk(char *dev);

#if 0
#define DEFAULT_PART_STATUS ( PART_STATUS_VALID | PART_STATUS_ALLOCATED | \
			      PART_STATUS_IN_USE | PART_STATUS_READABLE | \
			      PART_STATUS_WRITABLE )
#endif

char *defparttype[] = {
	PART_TYPE_PARTMAP,
	PART_TYPE_UNIX,
	PART_TYPE_MAC,
	PART_TYPE_SCRATCH,
};
int ndefparttypes = sizeof(defparttype) / sizeof(char *);

struct disklabel dl;

int main(int argc, char **argv)
{
	int fd, part;
	struct drvr_map dm;
	struct part_map_entry tmp, *pme;
	char path[MAXPATHLEN], rpath[MAXPATHLEN], *cmd;

	if (argc != 2) {
		fprintf(stderr, "usage: mdisk disk\n");
		exit(1);
	}

	sprintf(path, "/dev/%sc", argv[1]);
	sprintf(rpath, "/dev/r%sc", argv[1]);

	fd = open(path, O_RDONLY);
	if (fd == -1)
		err(1, "%s: Could not open", path);

	/* Read disklabel. */
	if (ioctl(fd, DIOCGDINFO, &dl) < 0)
		err(1, "%s: Could not read disklabel", path);

	/* Read and check block 0. */
	if (read(fd, &dm, sizeof(dm)) < 0)
		err(1, "%s: couldn't read driver map", path);
	if (dm.sbSig != DRIVER_MAP_MAGIC)
		init_disk(rpath);

	/*
	 * Read first partition map entry to get the number of
	 * partitions.
	 */
	lseek(fd, dl.d_secsize, SEEK_SET);
	read(fd, &tmp, sizeof(tmp));

	/*
	 * Now read the rest of the entries. We allocate pme to be
	 * as large as the whole partition map partition, but only
	 * read in the entries that are in use.
	 */
	pme = malloc(tmp.pmPartBlkCnt * sizeof(tmp));
	lseek(fd, dl.d_secsize, SEEK_SET);
	read(fd, pme, tmp.pmMapBlkCnt * sizeof(tmp));
	close(fd);

	/* Check the partition map. */
	for (part = 1; part < pme[0].pmMapBlkCnt; part++) {
		checkpart(pme, part + 1, pme[part].pmPyPartStart,
			  pme[part].pmPartBlkCnt, 0);
	}

	/* And do the main loop. */
	for (;;) {
		printf("mdisk: ");
		cmd = fparseln(stdin, NULL, NULL, "\0\0\0", 0);
		if (!cmd)
			break;

		switch (*cmd) {
		case '?':
			printf("Commands:\n");
			printf("  ? - help\n");
			printf("  p - print current partition map\n");
			printf("  a - add partition\n\n");
			printf("  d - delete partition\n");
			printf("  e - edit partition\n\n");
			printf("  w - write partition map to disk\n");
			printf("  q - quit\n\n");
			break;

		case 'p':
			printmap(pme);
			break;

		case 'a':
			addpart(pme);
			break;

		case 'd':
			part = getpart(cmd);
			if (part)
				delpart(pme, part);
			break;

		case 'e':
			part = getpart(cmd);
			if (part)
				editpart(pme, part);
			break;

		case 'w':
			writemap(pme, path);
			break;

		case 'q':
			return (0);

		default:
			printf("Unrecognized command. Type '?' for help.\n");
		}
		free(cmd);
	}

	return (0);
}

int getpart(char *cmd)
{
	int part;
	char *p, *num;

	if (cmd[1] == ' ') {
		part = strtoul(&cmd[2], &p, 10);
		if (!*p && part)
			return part;
	}

	for (;;) {
		printf("partition number: ");
		num = fparseln(stdin, NULL, NULL, "\0\0\0", 0);
		if (!num)
			return 0;
		part = strtoul(num, &p, 10);
		free(num);
		if (!*p && part)
			return part;
	}
}

void printmap(struct part_map_entry *pme)
{
	int i;

	printf("Disk has %d blocks.\n\n", dl.d_secperunit);

	for (i = 0; i < pme[0].pmMapBlkCnt; i++)
		printpart(pme, i + 1);
	printf("\n");
}

static char *unix_type(struct blockzeroblock *bzb)
{
	if (bzb->bzbMagic == BZB_MAGIC) {
		if (bzb->bzbFlags & BZB_ROOTFS)
			return "root";
		else if (bzb->bzbFlags & BZB_USRFS)
			return "generic";
		else if (bzb->bzbType == BZB_TYPESWAP)
			return "swap";
	}
	return "unknown";
}

void printpart(struct part_map_entry *pme, int part)
{
	struct part_map_entry *p = pme + (part - 1);
	double size;
	char units, *name;

	size = p->pmPartBlkCnt / 2.0;
	if (size < 1024)
		units = 'k';
	else {
		size /= 1024;
		if (size < 1024)
			units = 'M';
		else {
			size /= 1024;
			units = 'G';
		}
	}

	name = p->pmPartName;
	printf("%2d: %.32s (%.32s)", part, *name ? name : "[No Name]",
	       p->pmPartType);
	if (!strcmp(p->pmPartType, PART_TYPE_UNIX)) {
		struct blockzeroblock *bzb;

		bzb = (struct blockzeroblock *)&p->pmBootArgs;
		printf(" NetBSD %s", unix_type(bzb));
	}
	printf("\n    start %d, size %d (%.1f %c)\n",
	       p->pmPyPartStart, p->pmPartBlkCnt,
	       size, units);

#if 0
	if (p->pmPartStatus) {
		printf("    ");
		if (p->pmPartStatus & PART_STATUS_VALID)
			printf("valid ");
		if (p->pmPartStatus & PART_STATUS_ALLOCATED)
			printf("allocated ");
		if (p->pmPartStatus & PART_STATUS_IN_USE)
			printf("in_use ");
		if (p->pmPartStatus & PART_STATUS_BOOTABLE)
			printf("bootable ");
		if (p->pmPartStatus & PART_STATUS_READABLE)
			printf("readable ");
		if (p->pmPartStatus & PART_STATUS_WRITABLE)
			printf("writable ");
		if (p->pmPartStatus & PART_STATUS_BOOT_PIC)
			printf("boot_pic ");
		printf("\n");
	}

	if (p->pmPartStatus & PART_STATUS_BOOTABLE) {
		printf("    boot from %d, %d bytes, to 0x%x, entry 0x%x\n",
		       p->pmLgBootStart, p->pmBootSize,
		       p->pmBootLoad, p->pmBootEntry);
	}
#endif
}

void addpart(struct part_map_entry *pme)
{
	struct blockzeroblock *bzb;
	int n;

	n = pme[0].pmMapBlkCnt;
	if (n >= pme[0].pmPartBlkCnt) {
		printf("Partition map is full.\n");
		return;
	}

	memset(&pme[n], 0, sizeof(struct part_map_entry));
	strcpy(pme[n].pmPartName, "NetBSD");
	strcpy(pme[n].pmPartType, PART_TYPE_UNIX);
	bzb = (struct blockzeroblock *)&pme[n].pmBootArgs;
	bzb->bzbMagic = BZB_MAGIC;
	bzb->bzbFlags = BZB_USRFS;
	bzb->bzbType = 0;
	pme[n].pmPyPartStart =
		pme[n - 1].pmPyPartStart + pme[n - 1].pmPartBlkCnt;
	if (dl.d_secperunit > pme[n].pmPyPartStart + 1)
		pme[n].pmPartBlkCnt = dl.d_secperunit - pme[n].pmPyPartStart + 1;
	else
		pme[n].pmPartBlkCnt = 0;

	if (!editpart(pme, n))
		return;

	pme[n].pmSig = PART_ENTRY_MAGIC;
	pme[n].pmMapBlkCnt = n + 1;
	pme[n].pmDataCnt = pme[n].pmPartBlkCnt;
#if 0
	pme[n].pmPartStatus = DEFAULT_PART_STATUS;
#endif
	strcpy(pme[n].pmProcessor, "PowerPC");		/* XXX */

	while (--n)
		pme[n].pmMapBlkCnt++;
}

int editpart(struct part_map_entry *pme, int part)
{
	char *resp, name[32], type[32], *p;
	u_int32_t start, count;
	struct blockzeroblock bzb;

	printf("Partition %d\n\n", part);

	printf("Partition name (default: %s): ",
	       *pme[part - 1].pmPartName ?
	       (char *)pme[part - 1].pmPartName : "[No Name]");
	resp = fparseln(stdin, NULL, NULL, "\0\0\0", 0);
	if (!resp)
		return 0;
	/*
	 * pme.pmPartName is NUL-terminated iff it is less than
	 * 32 characters long, so just strncpy is actually correct
	 * here (and for pmPartType below).
	 */
	if (*resp)
		strncpy(name, resp, sizeof(name));
	else
		strncpy(name, pme[part - 1].pmPartName, sizeof(name));
	free(resp);

	*type = '\0';
	do {
		int i;

		printf("Partition type (or `?': default: %s): ",
		       pme[part - 1].pmPartType);
		resp = fparseln(stdin, NULL, NULL, "\0\0\0", 0);
		if (!resp)
			return 0;
		if (!strcmp(resp, "?")) {
			printf("Type a number to get a default type, or "
			       "any string for a non-standard type.\n");
			for (i = 0; i < ndefparttypes; i++)
				printf("  %2d. %s\n", i + 1, defparttype[i]);
		} else if (isdigit(*resp)) {
			i = atoi(resp);
			if (i <= ndefparttypes)
				strncpy(type, defparttype[i - 1],
					sizeof(type));
			else
				printf("No such type partition type %d.\n", i);
		} else if (*resp)
			strncpy(type, resp, sizeof(type));
		else
			strncpy(type, pme[part - 1].pmPartType, sizeof(type));
		free(resp);
	} while (!*type);

	if (!strcmp(type, PART_TYPE_UNIX)) {
		memcpy(&bzb, pme[part - 1].pmBootArgs, sizeof(bzb));
		if (bzb.bzbMagic != BZB_MAGIC) {
			memset(&bzb, 0, sizeof(bzb));
			bzb.bzbMagic = BZB_MAGIC;
			bzb.bzbFlags = BZB_USRFS;
			bzb.bzbType = 0;
		}
	tryagain:
		printf("NetBSD partition type (or `?': default: %s): ",
		       unix_type(&bzb));
		resp = fparseln(stdin, NULL, NULL, "\0\0\0", 0);
		if (!resp)
			return 0;
		if (!strcmp(resp, "?")) {
			printf("   1. generic\n   2. root\n   3. swap\n");
			goto tryagain;
		} else if (!strcmp(resp, "1") || !strcmp(resp, "generic")) {
			bzb.bzbFlags |= BZB_USRFS;
			bzb.bzbFlags &= ~BZB_ROOTFS;
			bzb.bzbType = 0;
		} else if (!strcmp(resp, "2") || !strcmp(resp, "root")) {
			bzb.bzbFlags |= BZB_ROOTFS;
			bzb.bzbFlags &= ~BZB_USRFS;
			bzb.bzbType = 0;
		} else if (!strcmp(resp, "3") || !strcmp(resp, "swap")) {
			bzb.bzbFlags &= ~(BZB_ROOTFS | BZB_USRFS);
			bzb.bzbType = BZB_TYPESWAP;
		} else if (strcmp(resp, "unknown") && *resp)
			goto tryagain;
	}

	printf("First block (default: %d): ", pme[part - 1].pmPyPartStart);
	resp = fparseln(stdin, NULL, NULL, "\0\0\0", 0);
	if (!resp)
		return 0;
	if (*resp)
		start = strtoul(resp, NULL, 10);
	else
		start = pme[part - 1].pmPyPartStart;

	printf("Size of partition (in b[locks], k, M, or G) (default: %d b): ",
	       pme[part - 1].pmPartBlkCnt);
	resp = fparseln(stdin, NULL, NULL, "\0\0\0", 0);
	if (!resp)
		return 0;
	if (*resp) {
		count = strtoul(resp, &p, 10);
		while (*p && isspace(*p))
			p++;
		switch (tolower(*p)) {
		case 'b':
		case '\0':
			break;

		case 'k':
			count /= 2;
			break;

		case 'm':
			count /= 2 * 1024;
			break;

		case 'g':
			count /= 2 * 1024 * 1024;
			break;

		default:
			printf("Unknown modifier '%c'\n", *p);
			return 0;
		}
	} else
		count = pme[part - 1].pmPartBlkCnt;

	if (!checkpart(pme, part, start, count, 1))
		return 0;

	strncpy(pme[part - 1].pmPartName, name, sizeof(name));
	strncpy(pme[part - 1].pmPartType, type, sizeof(type));
	if (!strcmp(type, PART_TYPE_UNIX))
		memcpy(pme[part - 1].pmBootArgs, &bzb, sizeof(bzb));
	pme[part - 1].pmPyPartStart = start;
	pme[part - 1].pmPartBlkCnt = count;

	return 1;
}

void delpart(struct part_map_entry *pme, int part)
{
	int i, nparts;

	if (part == 1) {
		printf("Nope\n");
		return;
	}

	nparts = pme[0].pmMapBlkCnt;
	if (part != nparts)
		memcpy(&pme[part - 1], &pme[part],
		       nparts * sizeof(*pme));

	pme[0].pmMapBlkCnt--;
	for (i = 1; i < pme[0].pmMapBlkCnt; i++)
		pme[i].pmMapBlkCnt--;

	printf("Deleted partition %d\n", part);
}

/*
 * Make sure that the new bounds for a partition don't overlap
 * any other partitions.
 */
int checkpart(struct part_map_entry *pme, int part,
	      u_int32_t start, u_int32_t size, int subjunctive)
{
	int i, retval = 1;

	if (start + size > dl.d_secperunit) {
		printf("Partition %d %s off the end of the disk.\n", part,
		       subjunctive ? "would extend" : "extends");
		retval = 0;
	}

	for (i = 0; i < pme[0].pmMapBlkCnt; i++) {
		if (i == part - 1)
			continue;

		if ((start >= pme[i].pmPyPartStart &&
		     start <= pme[i].pmPyPartStart + pme[i].pmPartBlkCnt - 1) ||
		    (start + size - 1 >= pme[i].pmPyPartStart &&
		     start + size - 1 <= pme[i].pmPyPartStart + pme[i].pmPartBlkCnt - 1) ||
		    (pme[i].pmPyPartStart >= start &&
		     pme[i].pmPyPartStart <= start + size - 1)) {
			printf("Partition %d %s partition %d.\n", part,
			       subjunctive ? "would overlap" : "overlaps",
			       i + 1);
			retval = 0;
		}
	}
	return retval;
}

void writemap(struct part_map_entry *pme, char *path)
{
	int fd;

	fd = open(path, O_WRONLY);
	if (!fd)
		err(1, "%s: Could not open", path);

	lseek(fd, dl.d_secsize, SEEK_SET);
	write(fd, pme, pme[0].pmMapBlkCnt * sizeof(struct part_map_entry));
	close(fd);
}	

void init_disk(char *dev)
{
	/* XXX */
	printf("%s isn't initialized\n", dev);
}