#include <stdbool.h>
#include <nrvn.h>
#include <acpi.h>
#include <vga.h>

#define data_port 0x60
#define status_reg 0x64
#define cmd_reg 0x64

#define config_rcmd 0x20
#define config_wcmd 0x60
#define ctrl_test_cmd 0xaa

static inline void print_port_status(bool second_port, const char *msg) {
	vga_puts("ps/2 port ");
	vga_puts(second_port ? "2 " : "1 ");
	vga_puts(msg);
}

static inline void send_cmd(uint8_t cmd) {
	while (inb(status_reg) & 0x2);
	outb(cmd_reg, cmd);
}

static inline bool read_data(uint8_t *data_out) {
	size_t counter = 50;
	while (!(inb(status_reg) & 0x1) && --counter > 0);
	if (counter == 0)
		return false;
	*data_out = inb(data_port);
	return true;
}

static inline void send_data(uint8_t data, bool second_device) {
	if (second_device)
		send_cmd(0xd4);
	while (inb(status_reg) & 0x2);
	outb(data_port, data);
}

// TODO: ps/2 mouse interface is weird.
static bool interface_test(bool second_device) {
	uint8_t rply;

	// Test port
	send_cmd(second_device ? 0xa9 : 0xab);
	read_data(&rply);
	switch (rply) {
		case 0x00:
			print_port_status(second_device, "port test PASS\n");
			break;
		case 0x01:
			print_port_status(second_device, "FAIL clock line stuck low\n");
			return false;
		case 0x02:
			print_port_status(second_device, "FAIL clock line stuck high\n");
			return false;
		case 0x03:
			print_port_status(second_device, "FAIL data line stuck low\n");
			return false;
		case 0x04:
			print_port_status(second_device, "FAIL data line stuck data\n");
			return false;
	}

	// Re-enable interface
	send_cmd(second_device ? 0xa8 : 0xae);

	// reset and self test device
	send_data(0xff, second_device);
	if (!read_data(&rply)) {
		print_port_status(second_device, "self-test FAIL: no reply\n");
		return false;
	}

	switch (rply) {
		case 0xFA:
			if (read_data(&rply), rply == 0xaa) {
				print_port_status(second_device, "self-test PASS\n");
			} else {
				print_port_status(second_device, "self-test FAIL");
				vga_putx(rply);
				vga_putc('\n');
				return false;
			}
			(void) inb(data_port); // PS/2 mouses send their id, we need to flush
			break;
		case 0xFC:
			print_port_status(second_device, "self-test FAIL: ");
			vga_putx(rply);
			vga_putc('\n');
			return false;
		default:
			print_port_status(second_device, "self-test unknown reply: ");
			vga_putx(rply);
			vga_putc('\n');
			return false;
	}

	// Disable device
	print_port_status(second_device, "disabling scanning.\n");
	send_data(0xf5, second_device);
	if (!read_data(&rply)) {
		print_port_status( second_device, "failed to disable: ");
		return false;
	}

	if (rply != 0xfa && rply != 0xaa) {
		print_port_status( second_device, "failed to disable: ");
		vga_putx(rply);
		vga_putc('\n');
		return false;
	}

	print_port_status(second_device, "scanning disabled.\n");

	print_port_status(second_device, "detecting device type.\n");
	send_data(0xf2, second_device);
	if (read_data(&rply), rply != 0xfa) {
		print_port_status(second_device, "failed to identify: ");
		vga_putx(rply);
		vga_putc('\n');
		return false;
	}
	print_port_status(second_device, "id ack: ");
	vga_putx((read_data(&rply), rply));
	vga_putc(' ');
	if (read_data(&rply))
		vga_putx(rply);
	vga_putc('\n');

	// Enable device
	send_data(0xf4, second_device);
	if (read_data(&rply), rply != 0xfa) {
		print_port_status(second_device, "failed to enable: ");
		vga_putx(rply);
		vga_putc('\n');
		return false;
	}

	print_port_status(second_device, "scanning enabled.\n");

	return true;
}

bool ps2_init(struct fadt *fadt) {
	if (fadt != NULL && !(fadt->boot_arch_flags & 0x2)) {
		vga_puts("fadt boot arch bit 1 not set: ");
		vga_putx(fadt->boot_arch_flags);
		vga_putc(' ');
		vga_putx(((char*)fadt)[109]);
		vga_putc('\n');
	}

	uint8_t rply;

	// Disable ps/2 devices
	vga_puts("disabling first device\n");
	send_cmd(0xad);
	vga_puts("disabling second device\n");
	send_cmd(0xa7);

	// flush output buffer
	(void) inb(data_port);

	// set controller config, clearing bits 0, 1, and 6
	send_cmd(config_rcmd);
	uint8_t config = (read_data(&rply), rply) & ~0b01000011;

	// check if dual channel is already set
	bool dual_channel = !(config & 0x5);

	send_cmd(config_wcmd);
	send_data(config, false);

	// perform controller self test
	send_cmd(ctrl_test_cmd);
	if (read_data(&rply), rply != 0x55) {
		vga_puts("ps/2 controller self test failed: ");
		vga_putx(rply);
		vga_putc('\n');
		return false;
	}

	// re-set config in case the self tests reset the device
	send_cmd(config_wcmd);
	send_data(config, false);

	// determine if there are two ports
	if (!dual_channel) {
		// attempt to enable second channel
		send_cmd(0xa8);

		// check if configuration show second port
		// and disable it again if it shows.
		send_cmd(0x20);
		if ((dual_channel = !((read_data(&rply), rply) & 0x5)))
			send_cmd(0xa7);
	}

	// enable and reset devices
	if (!interface_test(false))
		vga_puts("failed to test port 1 interface\n");

	if (dual_channel && !interface_test(true))
		vga_puts("failed to test port 2 interface\n");

	// re enabled bits 1 and 2 (interrupts)
	send_cmd(0x20);
	config = (read_data(&rply), rply) | 0x3;

	send_cmd(0x60);
	send_data(config, false);

	return true;
}