#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 fisBaseLower, fisBaseHigher;
    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, void *buffer) {
    uint64_t bufferPosition = (uint64_t)buffer;
    fis->lba0 = (uint8_t)bufferPosition;
    fis->lba1 = (uint8_t)(bufferPosition >> 8);
    fis->lba2 = (uint8_t)(bufferPosition >> 16);
    fis->lba3 = (uint8_t)(bufferPosition >> 24);
    fis->lba4 = (uint8_t)(bufferPosition >> 32);
    fis->lba5 = (uint8_t)(bufferPosition >> 40);
}

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

void sataDriveAccess(HardDrive *hardDrive, uint32_t lba, void *buffer,
                     uint8_t sectorCount, uint8_t direction) {
    SataInterface *interface = hardDrive->interface;
    CommandListStructure *commandList = getPointer(
        interface->port->commandListLower, interface->port->commandListUpper);
    interface->port->interruptStatus = -1; // clear pending interrupts
    int8_t slot = findFreeCommandHeader(interface->port);
    if (slot == -1) {
        printf("sata access failed: no free command header spot found!\n");
        return;
    }
    CommandHeader *commandHeader = &(commandList->commandHeaders[slot]);
    printf("interface: %x\ncommand header: %x\ncommandList: %x\n", interface,
           commandHeader, commandList);
    yields();
    commandHeader->fisLength = sizeof(FisToDevice) / sizeof(uint32_t);
    printf("size: %i\n", commandHeader->fisLength);
    commandHeader->write = 0; // read
    commandHeader->regionDescriptorCount = ((sectorCount - 1) >> 4) + 1;

    CommandTable *commandTable = getPointer(commandHeader->commandTableLower,
                                            commandHeader->commandTableUpper);
    memset(commandTable, 0,
           sizeof(CommandTable) + (commandHeader->regionDescriptorCount - 1) *
                                      sizeof(PhysicalRegionDescriptor));
    uint32_t i;
    void *currentPosition = buffer;
    for (i = 0; i < commandHeader->regionDescriptorCount - 1; i++) {
        commandTable->regionDescriptor[i].dataBaseLower = currentPosition;
        commandTable->regionDescriptor[i].dataBaseUpper =
            ((uint64_t)currentPosition) >> 32;
        commandTable->regionDescriptor[i].byteCount = 8 * 1024 - 1;
        commandTable->regionDescriptor[i].interrupt = 1;
        currentPosition += 8 * 1024;
        sectorCount -= 16;
    }
    commandTable->regionDescriptor[i].dataBaseLower = (uint64_t)currentPosition;
    commandTable->regionDescriptor[i].dataBaseUpper =
        (uint64_t)currentPosition >> 32;
    commandTable->regionDescriptor[i].byteCount = (sectorCount << 9) - 1;
    commandTable->regionDescriptor[i].interrupt = 1;
    FisToDevice *fis = (void *)&(commandTable->fis);
    setupFisToRead(fis, buffer, sectorCount);
    printf("waiting.");
    while (interface->port->taskFileData & (ATA_DEV_BUSY | ATA_DEV_DRQ)) {
        yields();
    }
    printf("..");
    yields();

    uint32_t slotMask = 1 << slot;
    interface->port->commandIssue |= slotMask;
    while (1) {
        if (!(interface->port->commandIssue & slotMask)) {
            break;
        }
    }
    printf("\n");
}

#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
    CommandListStructure *commandList =
        mallocAligned(sizeof(CommandListStructure), 10);
    disableCommands(port);
    port->commandListLower = commandList;
    port->commandListUpper = (uint64_t)commandList >> 32;
    memset(commandList, 0, sizeof(CommandListStructure));
    for (uint8_t i = 0; i < 2; i++) {
        uint16_t size = 256; // sizeof(CommandTable)
        commandList->commandHeaders[i].regionDescriptorCount = 8;
        uint64_t commandTable = (uint64_t)mallocAligned(size, 7);
        commandList->commandHeaders[i].commandTableLower = commandTable;
        commandList->commandHeaders[i].commandTableUpper = commandTable >> 32;
        memset((void *)commandTable, 0, size);
    }
    enableCommands(port);
}

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;
        }
        HardDrive *hardDrive = malloc(sizeof(HardDrive));
        SataInterface *interface = malloc(sizeof(SataInterface));
        hardDrive->access = sataDriveAccess;
        hardDrive->interface = interface;
        interface->registers = registers;
        interface->port = &registers->ports[i];
        relocatePort(interface->port);
        listAdd(drives, hardDrive);
        printf("device on port %i has type %x\n", i, type);
    }
}