#include <alloc.h>
#include <ata.h>
#include <hardDrive.h>
#include <ports.h>
#include <stdio.h>

#define IDE_DATA_REGISTER 0
#define IDE_ERROR_REGISTER 1
#define IDE_SECTOR_COUNT_REGISTER 2
#define IDE_SECTOR_NUMBER_REGISTER 3
#define IDE_CYLINDER_LOW_REGISTER 4
#define IDE_CYLINDER_HIGH_REGISTER 5
// sector size, drive, head
#define IDE_SDH_REGISTER 6
#define IDE_STATUS_REGISTER 7
#define IDE_COMMAND_REGISTER 7
#define IDE_ALTERNATE_STATUS_CONTROL_REGISTER 0
#define IDE_DRIVE_ADDRESS_CONTROL_REGISTER 1

#define IDE_STATUS_BUSY 0x80
#define IDE_STATUS_READY 0x40
#define IDE_STATUS_WRITE_FAULT 0x20
#define IDE_STATUS_SEEK_COMPLETE 0x10
#define IDE_STATUS_DATA_REQUEST 0x08
#define IDE_STATUS_CORRECTED_DATA 0x04
#define IDE_STATUS_INDEX 0x02
#define IDE_STATUS_ERROR 0x01

#define IDE_IDENTIFY_COMMAND 0xEC

typedef struct {
    uint16_t base, control, busMaster;
} IdeChannels;

typedef struct {
    IdeChannels *channels;
    uint8_t channel, drive, cylinders, heads, sectors;
    uint16_t capabilities, signature, type;
    uint32_t commandSets;
} IdeInterface;

void ideWriteControlRegister(IdeChannels *channel, uint8_t reg, uint8_t data) {
    outb(channel->control + reg, data);
}

void ideWriteRegister(IdeChannels *channel, uint8_t reg, uint16_t data) {
    outb(channel->base + reg, data);
    if (data > 0xFF) {
        // set higher half bit in the control register
        outb(channel->control, 0x0A);
        outb(channel->base + reg, data >> 8);
        outb(channel->control, 0x02);
    }
}

uint8_t ideReadControlRegister(IdeChannels *channel, uint8_t reg) {
    return inb(channel->control + reg);
}

uint8_t ideReadRegister(IdeChannels *channel, uint16_t reg) {
    return inw(channel->base + reg);
}

void ideReadBuffer(IdeChannels *channel, uint32_t *data) {
    for (uint16_t i = 0; i < 128; i++) {
        data[i] = ini(channel->base + IDE_DATA_REGISTER);
    }
}

void printError(uint16_t status) {
    if (status & 0x01) {
        printf("No address mark\n");
    }
    if (status & 0x02) {
        printf("Track 0 not found\n");
    }
    if (status & 0x04) {
        printf("command aborted\n");
    }
    if (status & 0x08) {
        printf("Media change request\n");
    }
    if (status & 0x10) {
        printf("ID mark not found\n");
    }
    if (status & 0x20) {
        printf("Media changed\n");
    }
    if (status & 0x40) {
        printf("uncorrectable data\n");
    }
    if (status & 0x80) {
        printf("bad block\n");
    }
}

void ideWaitUntilNotBusy(IdeChannels *channel, bool checkErrors) {
    for (uint8_t i = 0; i < 4; i++) {
        ideReadControlRegister(channel, 0);
    }
    uint16_t status = ideReadRegister(channel, IDE_STATUS_REGISTER);
    while (status & IDE_STATUS_BUSY) {
        status = ideReadRegister(channel, IDE_STATUS_REGISTER);
        yields();
    }
    if (checkErrors && status & IDE_STATUS_ERROR) {
        printf("encountered an error\n");
        printError(ideReadRegister(channel, IDE_ERROR_REGISTER));
    }
}

void ideCommand(IdeChannels *channel, uint8_t command) {
    ideWaitUntilNotBusy(channel, false);
    ideWriteRegister(channel, IDE_COMMAND_REGISTER, command);
}

#define CHANNEL channels[channel]

void ideDriveAccess(HardDrive *hardDrive, uint64_t lbaIn, void *buffer,
                    uint16_t sectorCount, uint8_t direction) {
    uint32_t lba = lbaIn;
    uint8_t accessMode = 0, head = 0, sector = 0, cylinder = 0;
    uint32_t lbaIo[6] = {0, 0, 0, 0, 0, 0};
    IdeInterface *interface = hardDrive->interface;
    if (lba >= 0x10000000) {
        accessMode = 2;
        lbaIo[0] = (uint8_t)lba;
        lbaIo[1] = (lba & 0x0000FF00) >> 8;
        lbaIo[2] = (lba & 0x00FF0000) >> 16;
        lbaIo[3] = (lba & 0xFF000000) >> 24;
    } else {
        accessMode = 1;
        lbaIo[0] = lba & 0xFF;
        lbaIo[1] = (lba & 0x000FF00) >> 8;
        lbaIo[2] = (lba & 0x0FF0000) >> 16;
    }
    ideWaitUntilNotBusy(interface->channels, true);
    if (accessMode == 0) {
        // chs
        ideWriteRegister(interface->channels, IDE_SDH_REGISTER,
                         0xA0 | (interface->drive << 4));
    } else {
        // lba
        ideWriteRegister(interface->channels, IDE_SDH_REGISTER,
                         0xE0 | (interface->drive << 4) | head);
    }
    ideWriteRegister(interface->channels, IDE_SECTOR_COUNT_REGISTER,
                     sectorCount);
    ideWriteRegister(interface->channels, IDE_SECTOR_NUMBER_REGISTER,
                     lbaIo[3] << 8 | lbaIo[0]);
    ideWriteRegister(interface->channels, IDE_CYLINDER_LOW_REGISTER,
                     lbaIo[4] << 8 | lbaIo[1]);
    ideWriteRegister(interface->channels, IDE_CYLINDER_HIGH_REGISTER,
                     lbaIo[5] << 8 | lbaIo[2]);
    uint8_t command = 0;
    if (direction == 0) {
        // read
        if (accessMode == 2) {
            command = 0x24;
        } else {
            command = 0x20;
        }
    } else {
        // write
        if (accessMode == 2) {
            command = 0x34;
        } else {
            command = 0x30;
        }
    }
    ideCommand(interface->channels, command);
    ideWaitUntilNotBusy(interface->channels, true);
    if (direction == 0) {
        uint16_t status;
        void *current = buffer;
        for (uint8_t i = 0; i < sectorCount; i++) {
            ideWaitUntilNotBusy(interface->channels, true);
            ideReadBuffer(interface->channels, (void *)buffer + 512 * i);
        }
    }
}

void setupIDEChannels(IdeChannels **channels, PciDevice *pciDevice) {
    if (pciDevice->programmingInterface & 0x01) {
        channels[0]->base = pciDevice->bar0;
        channels[0]->control = pciDevice->bar1;
    } else {
        channels[0]->base = 0x1F0;
        channels[0]->control = 0x3F6;
    }
    if (pciDevice->programmingInterface & 0x04) {
        channels[1]->base = pciDevice->bar2;
        channels[1]->control = pciDevice->bar3;
    } else {
        channels[1]->base = 0x170;
        channels[1]->control = 0x376;
    }
}

void initializeIdeController(PciDevice *pciDevice, ListElement **drives) {
    uint8_t *buffer = malloc(2048);
    IdeChannels *channels[] = {malloc(sizeof(IdeChannels)),
                               malloc(sizeof(IdeChannels))};
    uint8_t usedChannels = 0; // to be able to clean up later
    setupIDEChannels(channels, pciDevice);
    for (uint8_t channel = 0; channel < 2; channel++) {
        if (ideReadRegister(CHANNEL, IDE_STATUS_REGISTER) == 0) {
            // no device
            continue;
        }
        ideWriteControlRegister(CHANNEL, IDE_ALTERNATE_STATUS_CONTROL_REGISTER,
                                2);
        for (uint8_t drive = 0; drive < 2; drive++) {
            ideWaitUntilNotBusy(CHANNEL, true);
            // set the drive
            ideWriteRegister(CHANNEL, IDE_SDH_REGISTER, 0xA0 | (drive << 4));
            ideWriteRegister(CHANNEL, IDE_SECTOR_COUNT_REGISTER, 0);
            ideWriteRegister(CHANNEL, IDE_SECTOR_NUMBER_REGISTER, 0);
            ideWriteRegister(CHANNEL, IDE_CYLINDER_LOW_REGISTER, 0);
            ideWriteRegister(CHANNEL, IDE_CYLINDER_HIGH_REGISTER, 0);
            ideCommand(CHANNEL, IDE_IDENTIFY_COMMAND);
            ideWaitUntilNotBusy(CHANNEL, true);
            ideReadBuffer(CHANNEL, (void *)buffer);
            if (*((uint16_t *)buffer) == 0) {
                // any bit has to be set if it is a real device
                continue;
            }
            usedChannels |= 1 << channel;
            HardDrive *hardDrive = malloc(sizeof(HardDrive));
            IdeInterface *interface = malloc(sizeof(IdeInterface));
            hardDrive->interface = interface;
            addIdentifyData(hardDrive, (void *)buffer);
            hardDrive->access = ideDriveAccess;
            interface->channels = CHANNEL;
            interface->channel = channel;
            interface->drive = drive;
            interface->type = *(uint16_t *)buffer;
            interface->cylinders = *(uint16_t *)buffer + 2;
            interface->heads = *(uint16_t *)buffer + 6;
            interface->sectors = *(uint16_t *)buffer + 12;
            interface->capabilities = *(uint16_t *)buffer + 98;
            interface->commandSets = *(uint32_t *)(buffer + 164);
            listAdd(drives, hardDrive);
        }
    }
    if (!(usedChannels & 1)) {
        free(channels[0]);
    }
    if (!(usedChannels & 2)) {
        free(channels[1]);
    }
    free(buffer);
}