/*
 * CHU - driver for CH375 in disk mode
 * Copyright (C) 2024 Javier S. Pedro
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.

 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <dos.h>
#include <dosfunc.h>
#include <io.h>

#include "inc/mbr.h"
#include "inc/dosdrv.h"

#include "chudisk.h"

uint16_t ch375_port;

bool read_only;
bool show_fat32;

uint8_t drive;
uint32_t partition_start;
uint32_t partition_len;

#define MAX_PARTITIONS 32
struct partinfo {
	uint32_t start_sector;
	uint32_t num_sectors;
	uint8_t type;
} partinfo[MAX_PARTITIONS];
unsigned int num_parts;

static int ioctl_read_control(uint8_t drive, void *buffer, int len)
{
	union REGS regs;
	regs.h.ah = DOS_IOCTL;
	regs.h.al = 0x04; // Read control string from block device
	regs.x.bx = drive;
	regs.x.cx = len;
	regs.x.dx = FP_OFF(buffer);

	intdos(&regs, &regs);

	if (regs.x.cflag) {
		return -1;
	} else {
		return regs.x.ax;
	}
}

static int ioctl_write_control(uint8_t drive, void *buffer, int len)
{
	union REGS regs;
	regs.h.ah = DOS_IOCTL;
	regs.h.al = 0x05; // Write control string to block device
	regs.x.bx = drive;
	regs.x.cx = len;
	regs.x.dx = FP_OFF(buffer);

	intdos(&regs, &regs);

	if (regs.x.cflag) {
		return -1;
	} else {
		return regs.x.ax;
	}
}

static int ioctl_read_abs_sector(uint8_t drive, void *buffer, uint32_t sector, int sectors)
{
	CHUDISK_IOCTL_DISK_DATA ciod;
	union REGS regs;

	ciod.rwabs.buffer = buffer;
	ciod.rwabs.sector = sector;
	ciod.rwabs.sectors = sectors;

	regs.h.ah = DOS_IOCTL;
	regs.h.al = 0x0D; // IOCTL for block devices
	regs.x.bx = drive;
	regs.h.ch = DOS_IOCTL_DISK;
	regs.h.cl = CHUDISK_IOCTL_DISK_READ_ABS_SECTOR;
	regs.x.dx = FP_OFF(&ciod);

	intdos(&regs, &regs);

	return regs.x.cflag ? -1 : 0;
}

static void dos_disk_reset()
{
	union REGS regs;
	regs.h.ah = 0x0D; // DOS DISK RESET

	intdos(&regs, &regs);
}

static int parse_chudisk_control_string(const char *b, int len)
{
	// Not sure why I'm doing it this way instead of just sharing a pointer.
	int i = 0;
#if 0
	printf("control string:\n");
	for (i = 0; i < len; i+=8) {
		printf(" %hx %hx %hx %hx %hx %hx %hx %hx\n",
		       b[i+0], b[i+1], b[i+2], b[i+3], b[i+4], b[i+5], b[i+6], b[i+7]);
	}
	i = 0;
#endif
	while (i < len) {
		char c = b[i++];
		switch (c) {
		case 'B':
			memcpy(&ch375_port, &b[i], sizeof(ch375_port));
			i+=sizeof(ch375_port);
			break;
		case '3':
			show_fat32 = b[i++];
			break;
		case 'R':
			read_only = b[i++];
			break;
		case 'S':
			memcpy(&partition_start, &b[i], sizeof(partition_start));
			i+=sizeof(partition_start);
			break;
		case 'L':
			memcpy(&partition_len, &b[i], sizeof(partition_len));
			i+=sizeof(partition_len);
			break;
		case '\0':
			return 0;
		default:
			fprintf(stderr, "Unidentified control string parameter: '%c'\n", c);
			return -1;
		}
	}

	return 0;
}

static bool is_selectable_partition(const MBR_PARTENTRY *partentry)
{
	return partentry->type == MBR_FAT12 || partentry->type == MBR_FAT16_32MB
	        || partentry->type == MBR_FAT16 || partentry->type == MBR_FAT16_LBA
	        || (show_fat32 && (partentry->type == MBR_FAT32 || partentry->type == MBR_FAT32_LBA));
}

static const char *get_partition_type(const MBR_PARTENTRY *partentry)
{
	static char buffer[8];
	switch (partentry->type) {
	case MBR_EMPTY:
		return "invalid";
	case MBR_FAT12:
		return "fat12";
	case MBR_FAT16_32MB:
	case MBR_FAT16:
	case MBR_FAT16_LBA:
		return "fat16";
	case MBR_IFS:
		return "ntfs";
	case MBR_FAT32:
	case MBR_FAT32_LBA:
		return "fat32";
	case MBR_EXT:
	case MBR_EXT_LBA:
		return "extended";
	case MBR_LINUX:
	case MBR_LINUX_SWAP:
		return "linux";
	default:
		sprintf(buffer, "%#X", partentry->type);
		return buffer;
	}
}


static const char * show_partition_size(const MBR_PARTENTRY *partentry)
{
	static char buffer[12];
	uint32_t size = partentry->lbaLen / 2;
	STATIC_ASSERT(SECTOR_SIZE == 512);

	if (size <= 10000) {
		sprintf(buffer, "%lu KB", size);
		return buffer;
	}

	size /= 1024;
	if (size <= 10000) {
		sprintf(buffer, "%lu MB", size);
		return buffer;
	}

	size /= 1024;

	sprintf(buffer, "%lu GB", size);
	return buffer;
}

static const char *get_partition_label(const MBR_PARTENTRY *partentry, const uint32_t start_sector)
{
	static char buffer[SECTOR_SIZE];
	DOS_BPB *bpb = (DOS_BPB*) &buffer[DOS_BPB_OFFSET_ON_DISK];

	switch (partentry->type) {
	// Focus on well-known types for now
	case MBR_FAT12:
	case MBR_FAT16_32MB:
	case MBR_FAT16:
	case MBR_FAT16_LBA:
	case MBR_FAT32:
	case MBR_FAT32_LBA:
		break;
	default:
		return "           ";
	}

	if (ioctl_read_abs_sector(drive, buffer, start_sector, 1) != 0) {
		fprintf(stderr, "while reading partition boot record at %lu: %s\n", start_sector, strerror(errno));
		return "  ????     ";
	}

	if (bpb->sectorsPerFAT != 0) {
		// Assume FAT16
		return bpb->fat16.volumeLabel;
	} else {
		// Assume FAT32
		return bpb->fat32.volumeLabel;
	}
}

static void print_partition(const uint32_t mbr_sector, const MBR_PARTENTRY *partentry, bool selectable)
{
	const bool ebr = mbr_sector != 0;
	const uint32_t start_sector = mbr_sector + partentry->lbaFirst;
	const bool is_current = partition_start == start_sector;
	const char *typeLabel = get_partition_type(partentry);
	const char *volumeLabel = get_partition_label(partentry, start_sector);
	const char *sizeLabel = show_partition_size(partentry);

	if (partentry->type == MBR_EMPTY) {
		// Skip empty partitions no matter what
		return;
	}
	if (ebr && (partentry->type == MBR_EXT || partentry->type == MBR_EXT_LBA)) {
		// Skip displaying EXT partitions EBR, noisy
		return;
	}

	if (selectable) {
		const unsigned index = num_parts + 1;

		printf("   %c[%2u]  %8s  %-11.11s  %10s \n",
		       is_current ? '*' : ' ', index, typeLabel, volumeLabel, sizeLabel);
	} else {
		printf("   %c[  ]  %8s  %-11.11s  %10s \n",
		       is_current ? '*' : ' ', typeLabel, volumeLabel, sizeLabel);
	}

}

static int scan_disk_partitions()
{
	static char buffer[SECTOR_SIZE];
	uint32_t next_sector = 0;

	printf("   Index  %8s  %-11.11s  %10s \n",
	       "Type", "Label", "Size");

	// Start with MBR at sector 0
	next_sector = 0;
	do {
		const uint32_t this_sector = next_sector;
		const MBR_PARTENTRY *partentry;
		int i;

		// Try to read sector
		if (ioctl_read_abs_sector(drive, buffer, this_sector, 1) != 0) {
			fprintf(stderr, "while reading %s at disk sector %lu: %s\n",
			        this_sector == 0 ? "MBR" : "EBR",
			        this_sector, strerror(errno));
			// Fatal for MBR, non-fatal for EBR
			if (this_sector == 0) {
				fprintf(stderr, "is the disk present?\n");
				return -1;
			}
			break;
		}

		// Check for signature
		if (*(uint16_t*)&buffer[MBR_SIGNATURE_OFFSET_ON_DISK] != 0xAA55U) {
			fprintf(stderr, "on disk sector %lu: invalid MBR signature %x\n",
			        this_sector, *(uint16_t*)&buffer[MBR_SIGNATURE_OFFSET_ON_DISK]);
			if (this_sector == 0) return -1;
			break;
		}

		// Parse the MBR
		partentry = (MBR_PARTENTRY*)&buffer[MBR_PARTENTRY_OFFSET_ON_DISK];

		next_sector = 0; // If there is another EBR we will store it here
		for (i = 0; i < 4 && num_parts < MAX_PARTITIONS; i++) {
			bool is_selectable = is_selectable_partition(&partentry[i]);

			print_partition(this_sector, &partentry[i], is_selectable);

			if (is_selectable) {
				partinfo[num_parts].type = partentry[i].type;
				partinfo[num_parts].start_sector = this_sector + partentry[i].lbaFirst;
				partinfo[num_parts].num_sectors = partentry[i].lbaLen;
				num_parts++;
			}
			if (partentry[i].type == MBR_EXT || partentry[i].type == MBR_EXT_LBA) {
				// Go to the EBR next
				next_sector = this_sector + partentry[i].lbaFirst;
			}
		}
	} while (next_sector != 0 && num_parts < MAX_PARTITIONS);

	return 0;
}

static void usage()
{
	printf("Usage: chud X:, where X is a drive letter handled by CHUDISK.SYS\n");
}

int main(int argc, const char *argv[])
{
	static char buffer[16];
	int i;
	unsigned part;

	for (i = 1; i < argc; i++) {
		// Drive letter?
		if (argv[i][0] >= 'A' && argv[i][0] <= 'Z' && argv[i][1] == ':') {
			drive = 1 + argv[i][0] - 'A';
		} else if (argv[i][0] >= 'a' && argv[i][0] <= 'z' && argv[i][1] == ':') {
			drive = 1 + argv[i][0] - 'a';
		}
	}

	// Now read the driver settings and current partition
	i = ioctl_read_control(drive, buffer, sizeof(buffer));
	if (i < 0) {
		perror("ioctl_read_control");
		usage();
		return EXIT_FAILURE;
	}

	if (parse_chudisk_control_string(buffer, i)) {
		fprintf(stderr, "wrong IOCTL string returned by driver\n");
		usage();
		return EXIT_FAILURE;
	}

	if (scan_disk_partitions() != 0) {
		return EXIT_FAILURE;
	}
	if (num_parts == 0) {
		fprintf(stderr, "No partitions found on the disk\n");
		return EXIT_FAILURE;
	}

	printf("Select partition in [1..%u] range: ", num_parts);
	i = scanf("%u", &part);

	if (i != 1) {
		// Assume cancel
		return EXIT_FAILURE;
	}

	if (part == 0 || part > num_parts) {
		fprintf(stderr, "Invalid partition %u\n", part);
		return EXIT_FAILURE;
	}

	printf("Switching drive to partition %u at offset %lu size %lu\n", part,
	       partinfo[part-1].start_sector, partinfo[part-1].num_sectors);

	dos_disk_reset(); // Flush caches before we swap the drive...

	i = 0;
	buffer[i++] = 'S';
	memcpy(&buffer[i], &partinfo[part-1].start_sector, sizeof(uint32_t));
	i+=sizeof(uint32_t);
	buffer[i++] = 'L';
	memcpy(&buffer[i], &partinfo[part-1].num_sectors, sizeof(uint32_t));
	i+=sizeof(uint32_t);
	buffer[i] = '\0';

	if (ioctl_write_control(drive, buffer, i) < i) {
		perror("ioctl_write_control");
		return EXIT_FAILURE;
	}

	// The above IOCTL will also cause the next mediacheck from DOS to return changed.

	return EXIT_SUCCESS;
}