#include <ahci.h>
#include <alloc.h>
#include <ata.h>
#include <hardDrive.h>
#include <paging.h>
#include <stdio.h>

#define LOWER(ptr) (uint32_t)(uintptr_t) ptr
#define UPPER(ptr) (uint32_t)((uint64_t)(uintptr_t)ptr >> 32)

uint8_t CheckPortType(HBAPort *Port) {
    uint32_t SataStatus = Port->SataStatus;

    uint8_t InterfacePowerManagement = (SataStatus >> 8) & 0b111;
    uint8_t DeviceDetection = SataStatus & 0b111;

    if (DeviceDetection != HBA_PORT_DEV_PRESENT) {
        return None;
    }
    if (InterfacePowerManagement != HBA_PORT_IPM_ACTIVE) {
        return None;
    }

    switch (Port->Signature) {
    case SATA_SIG_ATAPI:
        return SATAPI;
    case SATA_SIG_ATA:
        return SATA;
    case SATA_SIG_PM:
        return PM;
    case SATA_SIG_SEMB:
        return SEMB;
    default:
        return None;
    }
}

void configurePort(Port *);
bool ahciCommand(HBAPort *HBAPortPtr, uint64_t Sector, void *Buffer,
                 uint32_t SectorCount, uint8_t command);
void StartCMD(HBAPort *HBAPortPtr);
void StopCMD(HBAPort *HBAPortPtr);

AtaIdentifyData ident;

void accessAHCIDrive(HardDrive *drive, uint64_t lba, void *buffer,
                     uint16_t count, uint8_t direction) {
    Port *port = (void *)drive->interface;
    printf("start ");
    yields();
    ahciCommand(port->HBAPortPtr, lba, buffer, count, ATA_CMD_READ_DMA_EX);
    printf("end %x\n", *(uint64_t *)buffer);
    yields();
}

void ahciIdentify(HBAPort *hbaPort, void *buffer);

void preparePorts(HBAMemory *hba, ListElement **drives) {
    for (uint8_t i = 0; i < 32; i++) {
        if (!(hba->PortsImplemented & (uint32_t)(1 << i))) {
            continue;
        }
        uint8_t portType = CheckPortType(&hba->Ports[i]);
        if (portType != SATA) {
            continue;
        }
        Port *port = malloc(sizeof(Port));
        port->AHCIPortType = portType;
        port->HBAPortPtr = &hba->Ports[i];
        configurePort(port);
        printf("reconfigured the port\n");
        yields();
        ahciCommand(&hba->Ports[i], 0, &ident, 1, ATA_CMD_IDENTIFY);
        HardDrive *hardDrive = malloc(sizeof(HardDrive));
        hardDrive->interface = port;
        hardDrive->access = accessAHCIDrive;
        addIdentifyData(hardDrive, &ident);
        listAdd(drives, hardDrive);
        yields();
    }
}

void configurePort(Port *port) {
    HBAPort *HBAPortPtr = port->HBAPortPtr;
    StopCMD(HBAPortPtr);

    HBACommandHeader *CommandHeader = mallocAligned(1024, 10);
    HBAPortPtr->CommandListBase = LOWER(CommandHeader);
    HBAPortPtr->CommandListBaseUpper = UPPER(CommandHeader);
    memset(CommandHeader, 0, 1024);

    void *fis = mallocAligned(1024, 10);
    HBAPortPtr->FISBaseAddress = LOWER(fis);
    HBAPortPtr->FISBaseAddressUpper = UPPER(fis);
    memset(fis, 0, 256);
    for (int i = 0; i < 32; i++) {
        CommandHeader[i].PRDTLength = 8;
        void *CommandTable = mallocAligned(sizeof(HBACommandTable), 8);
        CommandHeader[i].CommandTableBaseAddress = LOWER(CommandTable);
        CommandHeader[i].CommandTableBaseAddressUpper = UPPER(CommandTable);
        memset(CommandTable, 0, 256);
    }
    StartCMD(HBAPortPtr);
}

void StopCMD(HBAPort *HBAPortPtr) {
    HBAPortPtr->CommandStatus &= ~HBA_PxCMD_ST;
    HBAPortPtr->CommandStatus &= ~HBA_PxCMD_FRE;

    while (true) {
        if (HBAPortPtr->CommandStatus & HBA_PxCMD_FR) {
            continue;
        }
        if (HBAPortPtr->CommandStatus & HBA_PxCMD_CR) {
            continue;
        }
        break;
    }
}

void StartCMD(HBAPort *HBAPortPtr) {
    while (HBAPortPtr->CommandStatus & HBA_PxCMD_CR) {
    }
    HBAPortPtr->CommandStatus |= HBA_PxCMD_FRE;
    HBAPortPtr->CommandStatus |= HBA_PxCMD_ST;
}

#define HBA_PxIS_TFES (1 << 30)

bool ahciCommand(HBAPort *hbaPort, uint64_t Sector, void *Buffer,
                 uint32_t SectorCount, uint8_t command) {
    hbaPort->InterruptStatus = (uint32_t)-1;
    HBACommandHeader *commandHeader =
        (void *)(uintptr_t)hbaPort->CommandListBase;
    commandHeader->CommandFISLength = sizeof(FisToDevice) / sizeof(uint32_t);
    commandHeader->Write = 0;
    commandHeader->PRDTLength = 1;

    HBACommandTable *CommandTable =
        (void *)(uintptr_t)(commandHeader->CommandTableBaseAddress);
    memset(CommandTable, 0, sizeof(HBACommandTable));

    CommandTable->PRDTEntry[0].DataBaseAddress = LOWER(Buffer);
    CommandTable->PRDTEntry[0].DataBaseAddressUpper = UPPER(Buffer);
    CommandTable->PRDTEntry[0].ByteCount = 512 * SectorCount - 1;
    CommandTable->PRDTEntry[0].InterruptOnCompletion = 1;

    FisToDevice *CommandFIS = (FisToDevice *)(&CommandTable->CommandFIS);
    memset(CommandFIS, 0, sizeof(FisToDevice));
    CommandFIS->FISType = FIS_TYPE_REG_H2D;
    CommandFIS->commandControl = 1;
    CommandFIS->command = command;
    CommandFIS->LBA0 = (uint8_t)(Sector);
    CommandFIS->LBA1 = (uint8_t)(Sector >> 8);
    CommandFIS->LBA2 = (uint8_t)(Sector >> 16);
    CommandFIS->LBA3 = (uint8_t)(Sector >> 24);
    CommandFIS->LBA4 = (uint8_t)(Sector >> 32);
    CommandFIS->LBA5 = (uint8_t)(Sector >> 40);
    CommandFIS->DeviceRegister = 1 << 6;
    CommandFIS->countLow = SectorCount & 0xFF;
    CommandFIS->countHigh = (SectorCount >> 8) & 0xFF;
    uint64_t Spin = 0;
    while ((hbaPort->TaskFileData & (ATA_DEV_BUSY | ATA_DEV_DRQ)) &&
           Spin < 1000000) {
        Spin++;
    }
    if (Spin == 1000000) {
        printf("ahci command timeout\n");
        return false;
    }
    hbaPort->CommandIssue = 1;
    while (true) {
        if (hbaPort->CommandIssue == 0) {
            break;
        }
        if (hbaPort->InterruptStatus & HBA_PxIS_TFES) {
            return false;
        }
    }
    printf("count: %i ", commandHeader->PRDBCount);
    return true;
}

#define OsOwnership (1 << 1)
#define BiosOwnership (1 << 0)

void initializeSataController(PciDevice *pciDevice, ListElement **drives) {
    if (pciDevice->programmingInterface != 1) {
        return;
    }
    HBAMemory *hba = (void *)(uintptr_t)(pciDevice->bar5 & 0xFFFFFFF0);
    markMMIO(hba);
    if (!(hba->GlobalHostControl & (1 << 31))) {
        printf("controller not in ahci mode!\n");
        return;
    }
    if (hba->HostCapabilitiesExtended & 1) {
        printf("the os does not own the hba, getting it now...\n");
        yields();
        hba->BIOSHandoffControlStatus |= OsOwnership;
        while (hba->BIOSHandoffControlStatus & BiosOwnership ||
               !(hba->BIOSHandoffControlStatus & OsOwnership)) {
        }
    }
    preparePorts(hba, drives);
}