#include <_stdio.h>
#include <alloc.h>
#include <hardDrive.h>

// information source:
// https://www.intel.com/content/dam/www/public/us/en/documents/technical-specifications/serial-ata-ahci-spec-rev1_3.pdf

#define SATA_ATA 0x00000101
#define SATA_ATAPI 0xEB140101
#define SATA_ENCLOSURE_MANAGEMENT_BRIDGE 0xC33C0101
#define SATA_SIG_PM 0x96690101

#define AHCI_DEV_NULL 0
#define AHCI_SATA_DEVICE 1
#define AHCI_SEMB_DEVICE 2
#define AHCI_DEV_PM 3
#define AHCI_DEV_SATAPI 4

#define POWER_ACTIVE 1
#define DEVICE_PRESENT 3

#define ATA_DEV_BUSY 0x80
#define ATA_DEV_DRQ 0x08

typedef struct {
    uint32_t commandListLower, commandListUpper;
    uint32_t fisLower, fisUpper;
    uint32_t interruptStatus, interruptEnable;
    uint32_t command;
    uint32_t reserved;
    uint32_t taskFileData;
    uint32_t signature;
    uint32_t sataStatus, sataControl, sataError, sataActive;
    uint32_t commandIssue;
    uint32_t sataNotification;
    uint32_t fisSwitchingControl;
    uint32_t deviceSleep;
    uint8_t reserved_[50];
    uint32_t vendorSpecific;
} PortControlRegisters;

typedef struct {
    uint32_t hostCapabilities, globalHostControl, interruptStatus,
        implementedPorts, version, cccCtl, cccPts, enclosureManagementLocation,
        enclosureManagementControl, hostCapabilitiesExtended, handoffStatus;
    uint8_t reserved[116];
    uint8_t vendor[96];
    PortControlRegisters ports[32];
} HbaMemoryRegisters;

typedef struct {
    uint8_t fisLength : 5, ATAPI : 1, write : 1, prefetchable : 1;

    uint8_t reset : 1, bist : 1, clearBusy : 1, reserved : 1,
        portMultiplierPort : 4;

    uint16_t regionDescriptorCount;

    volatile uint32_t transferedBytes;

    uint32_t commandTableLower, commandTableUpper;

    uint32_t reserved_[4];
} CommandHeader;

typedef struct {
    CommandHeader commandHeaders[32];
} CommandListStructure;

typedef struct {
    uint32_t dataBaseLower, dataBaseUpper;
    uint32_t reserved;
    uint32_t byteCount : 22, reserved_ : 9, interrupt : 1;
} PhysicalRegionDescriptor;

typedef struct {
    uint8_t fis[64];
    uint8_t atapi[16];
    uint8_t reserved[48];
    PhysicalRegionDescriptor regionDescriptor[8];
} CommandTable;

typedef struct {
    HbaMemoryRegisters *registers;
    PortControlRegisters *port;
} SataInterface;

#define FIS_TO_DEVICE_TYPE 0x27

typedef struct {
    uint8_t type;
    uint8_t portMultiplierPort : 4, reserved : 3, commandControl : 1;

    uint8_t command, featureLow;

    uint8_t lba0, lba1, lba2, device;

    uint8_t lba3, lba4, lba5, featureHigh;

    uint8_t countLow, countHigh, isochronousCommandCompletion, control;

    uint8_t reserved_[4];
} FisToDevice;

void *getPointer(uint32_t lower, uint32_t upper) {
    return (void *)(lower | (uint64_t)upper << 32);
}

uint32_t getType(PortControlRegisters *registers) {
    uint32_t status = registers->sataStatus;
    uint8_t detectionStatus = status & 0xF;
    uint8_t interfacePowerManagement = (status >> 8) & 0xF;
    if (detectionStatus != DEVICE_PRESENT ||
        interfacePowerManagement != POWER_ACTIVE) {
        return AHCI_DEV_NULL;
    }
    switch (registers->signature) {
    case SATA_ATAPI:
        return AHCI_DEV_SATAPI;
    case SATA_ENCLOSURE_MANAGEMENT_BRIDGE:
        return AHCI_SEMB_DEVICE;
    case SATA_SIG_PM:
        return AHCI_DEV_PM;
    default:
        return AHCI_SATA_DEVICE;
    }
}

int8_t findFreeCommandHeader(PortControlRegisters *port) {
    uint32_t filledSlots = port->sataControl | port->commandIssue;
    for (uint8_t i = 0; i < 32; i++) {
        if (!(filledSlots & 1 >> i)) {
            return i;
        }
    }
    return -1;
}

void insertLba(FisToDevice *fis, uint64_t lba) {
    fis->lba0 = (uint8_t)lba;
    fis->lba1 = (uint8_t)(lba >> 8);
    fis->lba2 = (uint8_t)(lba >> 16);
    fis->lba3 = (uint8_t)(lba >> 24);
    fis->lba4 = (uint8_t)(lba >> 32);
    fis->lba5 = (uint8_t)(lba >> 40);
}

void setupFisToRead(FisToDevice *fis, uint16_t sectorCount, uint64_t lba,
                    uint8_t command) {
    fis->type = FIS_TO_DEVICE_TYPE;
    fis->commandControl = 1;
    fis->command = command;
    fis->device = 1 << 6;
    fis->countLow = (uint8_t)sectorCount;
    fis->countHigh = (uint8_t)(sectorCount >> 8);
    insertLba(fis, lba);
}

void sataIssueCommand(PortControlRegisters *port) {
    while (port->taskFileData & (ATA_DEV_BUSY | ATA_DEV_DRQ)) {
        yields();
    }
    port->commandIssue = 1;
    while (port->commandIssue == 1) {
        yields();
    }
}

void sataCommand(PortControlRegisters *port, uint8_t command, void *buffer,
                 uint64_t lba, uint16_t sectorCount) {
    port->interruptStatus = -1;
    CommandListStructure *commandList =
        getPointer(port->commandListLower, port->commandListUpper);
    int8_t slot = findFreeCommandHeader(port);
    if (slot == -1) {
        printf("sata command failed: no free command header slot found!\n");
        return;
    }
    CommandHeader *commandHeader = &(commandList->commandHeaders[slot]);
    commandHeader->fisLength = sizeof(FisToDevice) / sizeof(uint32_t);
    commandHeader->write = 0;
    commandHeader->regionDescriptorCount = 1;
    CommandTable *commandTable = getPointer(commandHeader->commandTableLower,
                                            commandHeader->commandTableUpper);
    memset(commandTable, 0,
           sizeof(CommandTable) + (commandHeader->regionDescriptorCount - 1) *
                                      sizeof(PhysicalRegionDescriptor));
    commandTable->regionDescriptor[0].dataBaseLower = (uint64_t)buffer;
    commandTable->regionDescriptor[0].dataBaseUpper = (uint64_t)buffer >> 32;
    commandTable->regionDescriptor[0].byteCount = 511;
    commandTable->regionDescriptor[0].interrupt = 1;
    FisToDevice *fis = (void *)&(commandTable->fis);
    setupFisToRead(fis, sectorCount, lba, command);
    sataIssueCommand(port);
}

void sataDriveAccess(HardDrive *hardDrive, uint64_t lba, void *buffer,
                     uint16_t sectorCount, uint8_t direction) {
    SataInterface *interface = hardDrive->interface;
    sataCommand(interface->port, ATA_CMD_READ_DMA_EXT, buffer, lba,
                sectorCount);
}

#define HBA_PxCMD_ST 0x0001
#define HBA_PxCMD_FRE 0x0010
#define HBA_PxCMD_FR 0x4000
#define HBA_PxCMD_CR 0x8000

void disableCommands(PortControlRegisters *port) {
    port->command &= ~(HBA_PxCMD_ST | HBA_PxCMD_FRE);
    while (port->command & HBA_PxCMD_FR && port->command + HBA_PxCMD_CR) {
    }
}

void enableCommands(PortControlRegisters *port) {
    while (port->command & HBA_PxCMD_CR) {
    }
    port->command |= HBA_PxCMD_ST | HBA_PxCMD_FRE;
}

void relocatePort(PortControlRegisters *port) {
    // call only once for each port, nothing is being freed here
    disableCommands(port);
    CommandListStructure *commandList =
        mallocAligned(sizeof(CommandListStructure), 10);
    port->commandListLower = commandList;
    port->commandListUpper = (uint64_t)commandList >> 32;
    uint64_t fis = (uint64_t)mallocAligned(sizeof(CommandListStructure), 10);
    port->fisLower = fis;
    port->fisUpper = (uint64_t)commandList >> 32;
    memset(commandList, 0, sizeof(CommandListStructure));
    for (uint8_t i = 0; i < 32; i++) {
        uint16_t size = 256; // sizeof(CommandTable)
        commandList->commandHeaders[i].regionDescriptorCount = 8;
        uint64_t commandTable = (uint64_t)mallocAligned(size, 8);
        commandList->commandHeaders[i].commandTableLower = commandTable;
        commandList->commandHeaders[i].commandTableUpper = commandTable >> 32;
        memset((void *)commandTable, 0, size);
    }
    enableCommands(port);
}

void sataDeviceInitialize(PortControlRegisters *port, ListElement **drives,
                          HbaMemoryRegisters *registers) {
    HardDrive *hardDrive = malloc(sizeof(HardDrive));
    SataInterface *interface = malloc(sizeof(SataInterface));
    hardDrive->access = sataDriveAccess;
    hardDrive->interface = interface;
    interface->registers = registers;
    interface->port = port;
    relocatePort(interface->port);
    uint8_t *buffer = malloc(512);
    sataCommand(port, ATA_CMD_IDENTIFY, buffer, 0, 1);
    char *model = malloc(41);
    for (int i = 0; i < 40; i += 2) {
        *((uint16_t *)&model[i]) = buffer[54 + i + 1];
        *((uint16_t *)&model[i + 1]) = buffer[54 + i];
    }
    model[41] = 0;
    printf("model: %s\n", model);
    hardDrive->model = model;
    free(buffer);
    listAdd(drives, hardDrive);
}

void initializeSataController(PciDevice *pciDevice, ListElement **drives) {
    if (pciDevice->programmingInterface != 1) {
        // todo
        return;
    }
    HbaMemoryRegisters *registers = (HbaMemoryRegisters *)pciDevice->bar5;
    for (uint8_t i = 0; i < 32; i++) {
        if (!(registers->implementedPorts & (1 << i))) {
            continue;
        }
        uint32_t type = getType(&registers->ports[i]);
        if (type == AHCI_DEV_NULL) {
            continue;
        }
        if (type != AHCI_SATA_DEVICE) {
            printf("type %i not implemented yet...\n", type);
            // todo
            continue;
        }
        sataDeviceInitialize(&registers->ports[i], drives, registers);
    }
}