#include "drivers/atapi/ata.h"

#include "lib/print.h"

#include <stddef.h>

uint8_t selected_drive_primary = 0;
uint8_t selected_drive_secondary = 0;

uint8_t primary_master_info_buffer[512];
uint8_t primary_slave_info_buffer[512];
uint8_t secondary_master_info_buffer[512];
uint8_t secondary_slave_info_buffer[512];

ide_device_t ide_devices[4];

void fix_ata_string(char* buf, size_t sz) {
    if ((sz % 2) != 0) return;
    for (size_t i = 0; i < sz; i += 2) {
        char tmp = buf[i];
        buf[i] = buf[i + 1];
        buf[i + 1] = tmp;
    }
    buf[sz-1] = '\0'; 
}

void scan_drives() {
    uint8_t type = ident_device(IDE_SECONDARY_IO_BASE, IDE_DRIVE_MASTER, secondary_master_info_buffer);
    if (type == 2) {
        ide_devices[0].bus = IDE_SECONDARY_IO_BASE;
        ide_devices[0].device = IDE_DRIVE_MASTER;
        ide_devices->type = type;

        atapi_identify_t* ident_result = (atapi_identify_t*)secondary_master_info_buffer;
        ide_devices->ident_info_struct = (void*)ident_result;

        kprintf("Assigned ATAPI master device on secondary bus to drive %d\n", 0);
        uint8_t buf[2048];
        int rtn = read_atapi(IDE_SECONDARY_IO_BASE, IDE_DRIVE_MASTER, 0x10, 1, (uint16_t*)buf);
        if (rtn < 0) {
            if (rtn == -1) {
                kprintf("Read error, failed to send command\n");
            }
            
            if (rtn == -2) {
                kprintf("Read error, failed to read sector\n");
            }

            kprintf("Read error\n");
            return;
        }

        kprintf("Read first sector of drive\n");

        if (buf[0] != 0x01) {
            kprintf("First byte not valid for ISO drive\n");
            return;
        }

        if (buf[1] != 'C' || buf[2] != 'D') {
            kprintf("Magic bytes are invalid for ISO drive\n");
            return;
        }

        kprintf("ISO filesystem detected on drive\n");
    }
}

extern void outb(uint8_t port, uint8_t data);
extern uint8_t inb(uint8_t port);
extern void atapi_wait(uint16_t port);

int read_atapi(uint16_t bus, uint8_t drive, uint32_t lba, uint32_t sectors, uint16_t* buf) {
    volatile uint8_t read_cmd[12] = {0xA8, 0,
	                                 (lba >> 0x18) & 0xFF, (lba >> 0x10) & 0xFF, (lba >> 0x08) & 0xFF,
	                                 (lba >> 0x00) & 0xFF,
	                                 (sectors >> 0x18) & 0xFF, (sectors >> 0x10) & 0xFF, (sectors >> 0x08) & 0xFF,
	                                 (sectors >> 0x00) & 0xFF,
	                                 0, 0};

    outb(bus + IO_DRIVE_SELECT_REG, drive);   // select drive
    atapi_wait(bus);
    outb(bus + IO_FEATURE_REG, 0x00);
    outb(bus + IO_LBA_MID_REG, 2048 & 0xFF);
    outb(bus + IO_LBA_HI_REG, 2048 >> 8);
    outb(bus + IO_CMD_REG, 0xA0);

    atapi_wait(bus);
    atapi_wait(bus);

    while (1) {
        kprintf("waiting bsy\n");
		uint8_t status = inb(bus + IO_STATUS_REG);
        if (status & STATUS_REG_BSY) {
            atapi_wait(bus);
            continue;
        }

		if (status & STATUS_REG_ERR)
			return -1;

		if (status & STATUS_REG_DRQ)
			break;

		atapi_wait(bus);
	}

    send_scsi_packet(bus, (void*)read_cmd, 6);

    for (uint32_t i = 0; i < sectors; i++) {
        while (1) {
            uint8_t status = inb(bus + IO_STATUS_REG);
            if (status & STATUS_REG_ERR)
                return -2;
            if (!(status & STATUS_REG_BSY) && (status & STATUS_REG_DRQ))
                break;
            atapi_wait(bus);
        }
        int bytes = inb(bus + IO_LBA_HI_REG) << 8 | inb(bus + IO_LBA_MID_REG);
        int words = bytes / 2;

        read_data_words(bus, ((uint8_t*)buf + i * 0x800), words);
    }

    return 0;
}