#include <stdint.h>

#include <hardware/clocks.h>
#include <hardware/flash.h>
#include <hardware/gpio.h>
#include <hardware/timer.h>
#include <pico/multicore.h>

#include <bsp/board_api.h>
#include <tusb.h>

#include "ff/ff.h"
#include "ff/diskio.h"

#include "6502.h"
#include "6522.h"
#include "drive.h"

const uint32_t GPIO_CLK = 2;
const uint32_t GPIO_DATA = 3;
const uint32_t GPIO_DATA2 = 5;
const uint32_t GPIO_ATN = 4;
const uint32_t GPIO_LED = 25;

const uint8_t DEVICE = 8;

#define FS_BYTES 1572864
static uint8_t __in_flash("fs") __attribute__((aligned(4096))) FS[FS_BYTES];

static FATFS FAT_FS;

static uint8_t ram[2048];
static uint8_t dos1541[16384];
static uint8_t disk[174848];

uint32_t ts;
cpu_t cpu;
static via_t via1;
static via_t via2;
drive_t drive;

int8_t track_step;

void via_pa_read(via_t *via, uint8_t cycle) {
	if (via == &via2) {
		via->ira = drive.latched[(drive.current_idx + cycle) & (DRIVE_LOOKAHEAD - 1)];
	} else {
		via->ira = via->ora;
	}
}

void via_pb_read(via_t *via, uint8_t cycle) {
	if (via == &via1) {
		while (ts + cycle > time_us_32()) {}
		uint8_t clk = gpio_get(GPIO_CLK) ? 0 : 0b00000100;
		uint8_t data = gpio_get(GPIO_DATA) ? 0 : 0b00000001;
		uint8_t atn = gpio_get(GPIO_ATN) ? 0 : 0b10000000;
		uint8_t dev = (DEVICE - 8) << 5;
		uint8_t atna = via->orb & 0b00010000;

		via->irb = clk | data | atn | dev | atna;
	} else {
		uint8_t low = via->orb & 0b00001111;
		uint8_t sync = drive.sync[(drive.current_idx + cycle) & (DRIVE_LOOKAHEAD - 1)] ? 0 : 0b10000000;

		via->irb = low | sync | 0b00010000;
	}
}

void via_pa_write(via_t *via, uint8_t cycle) {
}

void via_pb_write(via_t *via, uint8_t cycle) {
	if (via == &via1) {
		while (ts + cycle > time_us_32()) {}
		gpio_set_dir(GPIO_CLK, (via->orb & 0b00001000) > 0 && (via->ddrb & 0b00001000) > 0);
		gpio_set_dir(GPIO_DATA, (via->orb & 0b00000010) > 0 && (via->ddrb & 0b00000010) > 0);
		if ((via->orb & 0b00010000) && (via->ddrb & 0b00010000)) {
			gpio_set_dir(GPIO_DATA2, gpio_get(GPIO_ATN));
		} else {
			gpio_set_dir(GPIO_DATA2, !gpio_get(GPIO_ATN));
		}
	} else if (via == &via2) {
		if ((via->ddrb & 3) == 3) {
			int8_t step = via->orb & 3;
			int8_t r = (step - track_step) % 4;
			uint8_t mod = r < 0 ? r + 4 : r;
			if (mod == 1) {
				drive_step_up(&drive);
			} else if (mod == 3) {
				drive_step_down(&drive);
			}
			track_step = step;
		}

		drive.spinning = ((via->ddrb & via->orb) & 4) > 0;

		gpio_put(GPIO_LED, (via->orb & 0b00001000) > 0 && (via->ddrb & 0b00001000) > 0);
	} 
}

uint8_t bus_read(uint16_t addr, uint8_t cycle) {
	if (addr < 2048) {
		return ram[addr];
	} else if (addr >= 0xc000) {
		return dos1541[addr - 0xc000];
	} else if (addr >= 0x1800 && addr < 0x1810) {
		uint8_t v = 0;
		via_read(&via1, addr - 0x1800, &v, cycle);
		return v;
	} else if (addr >= 0x1c00 && addr < 0x1c10) {
		uint8_t v = 0;
		via_read(&via2, addr - 0x1c00, &v, cycle);
		return v;
	} else {
		return 0;
	}
}

void bus_write(uint16_t addr, uint8_t v, uint8_t cycle) {
	if (addr < 2048) {
		ram[addr] = v;
	} else if (addr >= 0x1800 && addr < 0x1810) {
		via_write(&via1, addr - 0x1800, v, cycle);
	} else if (addr >= 0x1c00 && addr < 0x1c10) {
		via_write(&via2, addr - 0x1c00, v, cycle);
	}
}

void atn_irq_callback(uint pin, uint32_t flags) {
	if (flags & GPIO_IRQ_EDGE_RISE) {
		if (!(via1.pcr & 1)) {
			via1.ifr |= 0b10000010;
		}
		if ((via1.orb & 0b00010000) && (via1.ddrb & 0b00010000)) {
			gpio_set_dir(GPIO_DATA2, 1);
		} else {
			gpio_set_dir(GPIO_DATA2, 0);
		}
	} else {
		if (via1.pcr & 1) {
			via1.ifr |= 0b10000010;
		}
		if ((via1.orb & 0b00010000) && (via1.ddrb & 0b00010000)) {
			gpio_set_dir(GPIO_DATA2, 0);
		} else {
			gpio_set_dir(GPIO_DATA2, 1);
		}
	}
}

static bool ejected = false;

void tud_msc_inquiry_cb(uint8_t lun, uint8_t vid[8], uint8_t pid[16], uint8_t rev[4]) {
	const char v[] = "TinyUSB";
	const char p[] = "pico1541";
	const char r[] = "1.0";

	memcpy(vid, v, strlen(v));
	memcpy(pid, p, strlen(p));
	memcpy(rev, r, strlen(r));
}

bool tud_msc_test_unit_ready_cb(uint8_t lun) {
	if (ejected) {
		tud_msc_set_sense(lun, SCSI_SENSE_NOT_READY, 0x3a, 0x00);
		return false;
	}
	return true;
}

void tud_msc_capacity_cb(uint8_t lun, uint32_t *block_count, uint16_t *block_size) {
	*block_count = FS_BYTES / 512;
	*block_size = 512;
}

bool tud_msc_start_stop_cb(uint8_t lun, uint8_t power_condition, bool start, bool load_eject) {
	if (load_eject && !start) {
		ejected = true;
	}
	return true;
}

int32_t tud_msc_read10_cb(uint8_t lun, uint32_t lba, uint32_t offset, void *buffer, uint32_t bufsize) {
	if (lba >= FS_BYTES / 512) {
		return -1;
	}

	uint8_t *addr = FS + 512 * lba + offset;
	memcpy(buffer, addr, bufsize);

	return (int32_t)bufsize;
}

// FIXME This is pretty naive....
int32_t tud_msc_write10_cb(uint8_t lun, uint32_t lba, uint32_t offset, uint8_t *buffer, uint32_t bufsize) {
	if (lba >= FS_BYTES / 512) {
		return -1;
	}

	uint32_t flash_sect = ((lba * 512 + offset) / 4096) * 4096;
	uint32_t sect_ofs = lba * 512 + offset - flash_sect;
	for (uint32_t i = 0; i < bufsize; i += 4096) {
		uint8_t flash_buf[4096];
		memcpy(flash_buf, FS + flash_sect + i, 4096);

		if (i == 0) {
			if (bufsize < 4096 - sect_ofs) {
				memcpy(flash_buf + sect_ofs, buffer, bufsize);
			} else {
				memcpy(flash_buf + sect_ofs, buffer, 4096 - sect_ofs);
			}
		} else if (bufsize - i < 4096) {
			memcpy(flash_buf, buffer + i - sect_ofs, bufsize - i);
		} else {
			memcpy(flash_buf, buffer + i - sect_ofs, 4096);
		}

		flash_range_erase((uint32_t)FS + flash_sect - 0x10000000, 4096);
		flash_range_program((uint32_t)FS + flash_sect - 0x10000000, flash_buf, 4096);	
	}

	return (int32_t)bufsize;
}

bool tud_msc_is_writable_cb(uint8_t lun) {
	return true;
}

int32_t tud_msc_scsi_cb(uint8_t lun, const uint8_t *scsi_cmd, void *buffer, uint16_t bufsize) {
	tud_msc_set_sense(lun, SCSI_SENSE_ILLEGAL_REQUEST, 0x20, 0x00);
	return -1;
}

DSTATUS disk_status(BYTE pdrv) {
	return 0;
}

DSTATUS disk_initialize(BYTE pdrv) {
	return 0;
}

DRESULT disk_read(BYTE pdrv, BYTE *buff, LBA_t sector, UINT count) {
	memcpy(buff, FS + 512 * sector, 512 * count);
	return RES_OK;
}

DRESULT disk_write(BYTE pdrv, const BYTE *buff, LBA_t sector, UINT count) {
	// FIXME implement
	return RES_ERROR;
}

DRESULT disk_ioctl(BYTE pdrv, BYTE cmd, void *buff) {
	switch (cmd) {
		case CTRL_SYNC:
			return RES_OK;
		case GET_SECTOR_COUNT:
			*(LBA_t *)buff = (FS_BYTES / 512);
			return RES_OK;
		case GET_SECTOR_SIZE:
			*(WORD *)buff = 512;
			return RES_OK;
		case GET_BLOCK_SIZE:
			*(DWORD *)buff = 4096;
			return RES_OK;
		case CTRL_TRIM:
			return RES_OK;
	}

	return RES_PARERR;
}

void drive_thread() {
	drive_init(&drive);
	drive.image = disk;
	
	multicore_fifo_push_blocking(1);

	while (multicore_fifo_pop_blocking_inline()) {
		drive_cycle(&drive);
	}
}

int main() {
	board_init();
	tud_init(BOARD_TUD_RHPORT);
	if (board_init_after_tusb) {
		board_init_after_tusb();
	}

	while (!ejected) {
		tud_task();
	}
	tud_disconnect();

	f_mount(&FAT_FS, "0:", 0);
	
	FIL file;
	FRESULT res;
	uint32_t btr;
	uint32_t br;

	res = f_open(&file, "0:dos1541.bin", FA_READ);
	if (res) {
		return 0;
	}
	btr = 16384;
	br = 0;
	while (btr > 0) {
		UINT brr;
		if (btr < 4096) {
			f_read(&file, dos1541 + br, btr, &brr);
		} else {
			f_read(&file, dos1541 + br, 4096, &brr);
		}
		br += brr;
		btr -= brr;
	}
	f_close(&file);

	res = f_open(&file, "0:image.d64", FA_READ);
	if (res) {
		return 0;
	}
	btr = 174848;
	br = 0;
	while (btr > 0) {
		UINT brr;
		if (btr < 4096) {
			f_read(&file, disk + br, btr, &brr);
		} else {
			f_read(&file, disk + br, 4096, &brr);
		}
		br += brr;
		btr -= brr;
	}
	f_close(&file);

	f_unmount("0:");
	
	set_sys_clock_khz(250000, true);
	multicore_launch_core1(drive_thread);

	gpio_init(GPIO_CLK);
	gpio_init(GPIO_DATA);
	gpio_init(GPIO_ATN);
	gpio_init(GPIO_DATA2);
	gpio_init(25);

	gpio_put(GPIO_CLK, 0);
	gpio_put(GPIO_DATA, 0);
	gpio_put(GPIO_LED, 0);
	gpio_put(GPIO_DATA2, 0);

	gpio_set_dir(GPIO_CLK, GPIO_IN);
	gpio_set_dir(GPIO_DATA, GPIO_IN);
	gpio_set_dir(GPIO_ATN, GPIO_IN);
	gpio_set_dir(GPIO_LED, GPIO_OUT);
	gpio_set_dir(GPIO_DATA2, GPIO_IN);

	gpio_set_pulls(GPIO_CLK, true, false);
	gpio_set_pulls(GPIO_DATA, true, false);
	gpio_set_pulls(GPIO_DATA2, true, false);
	gpio_set_pulls(GPIO_ATN, true, false);

	gpio_set_irq_enabled_with_callback(GPIO_ATN, GPIO_IRQ_EDGE_RISE | GPIO_IRQ_EDGE_FALL, true, &atn_irq_callback);

	track_step = 0;
	multicore_fifo_pop_blocking();

	cpu_reset(&cpu);
	via_init(&via1);
	via_init(&via2);

	ts = time_us_32();
	uint16_t osc2 = 0xffff;
	while (!cpu.jammed) {
		if (ts + 100 < time_us_32()) {
			break;
		}

		uint8_t cycles = cpu_step(&cpu);
		if (drive_advance(&drive, cycles) && (via2.pcr & 0b1110) == 0b1110) {
			cpu.flags |= 1 << FLAG_V_BIT;
		}
		multicore_fifo_push_blocking_inline(1);

		via_cycle(&via1, cycles);
		via_cycle(&via2, cycles);
		cpu.irq = via_irq(&via1) | via_irq(&via2);

		ts += cycles;
		while (ts > time_us_32()) {}
	}

	return 0;
}