// This file is part of noxos and licensed under the MIT open source license

#include "drivers/ahci.h"
#include "drivers/fs/mbr.h"
#include "utils/memory.h"
#include "utils/logger.h"
#include "mm/page_map.h"
#include "mm/page_frame.h"

string_t g_ahci_device_type_strings[AHCI_DEVICE_TYPE_END_OF_ENUM] = {
    "None",
    "SATA",
    "SEMB",
    "Port Multiplier",
    "SATAPI"
};

#include "utils/math.h"

ahci_controller_T* ahci_controller_alloc(pci_device_T* pci_device) {
    ahci_controller_T* controller = memory_allocate(sizeof(ahci_controller_T));

    controller->pci_device = pci_device;
    controller->abar       = (ahci_hba_memory_T*)((pci_header_0_T*)pci_device->header)->bar5;
    controller->num_ports  = 0;

    page_map_map_memory(g_kernel_page_map, controller->abar, controller->abar, PM_FLAG_READ_WRITE);

    // probe ports
    for (int i = 0; i < 32; i++) {
        if ((controller->abar->ports_implemented & (1 << i)) == 0) continue;

        ahci_device_type_E type = ahci_port_get_type(&controller->abar->ports[i]);

        switch (type) {
            case AHCI_DEVICE_SATAPI:
            case AHCI_DEVICE_SATA: {
                ahci_port_T* port = &controller->ports[controller->num_ports];
                port->type        = type;
                port->hba_port    = &controller->abar->ports[i];
                port->id          = controller->num_ports;

                controller->num_ports += 1;
                break;
            }
        }
    }

    // configure ports
    log(LOG_INFO, "AHCI devices:");
    for (int i = 0; i < controller->num_ports; i++) {
        log(LOG_INFO, "    %s", g_ahci_device_type_strings[controller->ports[i].type]);
        ahci_port_init(&controller->ports[i]);

        uint8_t* addr = memory_allocate(512 + PFRAME_SIZE);
        void* buffer = (void*)CEIL_TO((uint64_t)addr, PFRAME_SIZE);

        if (ahci_port_read(&controller->ports[i], 0, 1, buffer)) {
            mbr_header_T* mbr = (mbr_header_T*)buffer;

            log(LOG_INFO, "      MBR:");
            for (int j = 0; j < 4; j++) {
                mbr_partition_types_E type = mbr->partitions[j].partition_type;
                if (type == MBR_PARTITION_TYPE_UNUSED) continue;
                log(LOG_INFO, "        p%d: 0x%xb - %s", j, type, mbr_partition_type_to_string(type));
            }
        }

        memory_free(addr);
    }

    return controller;
}

void ahci_controller_destruct(ahci_controller_T* controller) {
    for (int i = 0; i < controller->num_ports; i++) {
        ahci_port_destruct(&controller->ports[i]);
    }

    page_map_unmap_memory(g_kernel_page_map, controller->abar);

    memory_free(controller);
}

void ahci_port_init(ahci_port_T* port) {
    ahci_port_command_stop(port);

    port->cmd_list_base                     = pframe_request();
    port->fis_base                          = pframe_request();
    port->hba_port->command_list_base       = (uint32_t)(uint64_t)port->cmd_list_base;
    port->hba_port->command_list_base_upper = (uint32_t)((uint64_t)port->cmd_list_base >> 32);
    port->hba_port->fis_base_address        = (uint32_t)(uint64_t)port->fis_base;
    port->hba_port->fis_base_address_upper  = (uint32_t)((uint64_t)port->fis_base >> 32);
    memory_set(port->cmd_list_base, 0, 0x400);
    memory_set(port->fis_base, 0, 0x100);

    ahci_hba_command_T* command_header = (ahci_hba_command_T*)port->cmd_list_base;

    for (uint32_t i = 0; i < 32; i++) {
        port->command_tables[i] = pframe_request();

        uint64_t address = (uint64_t)port->command_tables[i] + (i << 8);
        command_header[i].prdt_length                      = 8;
        command_header[i].command_table_base_address       = (uint32_t)address;
        command_header[i].command_table_base_address_upper = (uint32_t)(address >> 32);

        memory_set((void*)address, 0, 0x100);
    }

    ahci_port_command_start(port);
}

void ahci_port_destruct(ahci_port_T* port) {

    for (uint32_t i = 0; i < 32; i++) {
        pframe_free(port->command_tables[i]);
    }

    pframe_free(port->fis_base);
    pframe_free(port->cmd_list_base);
}

ahci_device_type_E ahci_port_get_type(ahci_hba_port_T* port) {
    if ((port->sata_status & 0b111) != HBA_PORT_DEVICE_PRESENT)    return AHCI_DEVICE_NONE;
    if (((port->sata_status >> 8) & 0b111) != HBA_PORT_IPM_ACTIVE) return AHCI_DEVICE_NONE;

    switch (port->signature) {
        case SATA_SIG_ATA:              return AHCI_DEVICE_SATA;
        case SATA_SIG_ATAPI:            return AHCI_DEVICE_SATAPI;
        case SATA_SIG_SEMB:             return AHCI_DEVICE_SEMB;
        case SATA_SIG_PORT_MULTIPLIER:  return AHCI_DEVICE_PORT_MULTIPLIER;
        default:                        return AHCI_DEVICE_NONE;
    }
}

void ahci_port_command_stop(ahci_port_T* port) {
    port->hba_port->command_status &= ~HBA_PX_COMMAND_START;
    port->hba_port->command_status &= ~HBA_PX_COMMAND_FIS_RECV_ENABLE;

    while ((port->hba_port->command_status & HBA_PX_COMMAND_CMD_LIST_RUNNING) &&
           (port->hba_port->command_status & HBA_PX_COMMAND_FIS_RECV_RUNNING)) {
        asm("hlt");
    }
}

void ahci_port_command_start(ahci_port_T* port) {
    while (port->hba_port->command_status & HBA_PX_COMMAND_CMD_LIST_RUNNING) asm("hlt");

    port->hba_port->command_status |= HBA_PX_COMMAND_FIS_RECV_ENABLE;
    port->hba_port->command_status |= HBA_PX_COMMAND_START;
}

// reads num_sector*512 bytes from offset (sector*512) into buffer_out, that needs to be page aligned
bool ahci_port_read(ahci_port_T* port, uint64_t sector, uint32_t num_sectors, void* buffer_out) {
    if (port->type == AHCI_DEVICE_SATAPI) return false;

    uint64_t blocked = 0;
    while ((port->hba_port->task_file_data & (ATA_DEVICE_BUSY | ATA_DEVICE_DRQ))) {
        if (blocked > 1000000) return false;
        blocked++;
    }

    uint64_t buffer                  = (uint64_t) page_map_get_physical_address(g_kernel_page_map, buffer_out);
    uint32_t sector_low              = (uint32_t)sector;
    uint32_t sector_high             = (uint32_t)(sector >> 32);
    port->hba_port->interrupt_status = (uint32_t)-1;


    ahci_hba_command_T* command_header = (ahci_hba_command_T*)(uint64_t)port->hba_port->command_list_base;
    command_header->command_fis_length = sizeof(AHCI_FIS_REGISTER_HOST_TO_DEVICE) / sizeof(uint32_t);
    command_header->write              = 0;
    command_header->prdt_length        = 1;

    ahci_hba_command_table_T* command_table = (ahci_hba_command_table_T*)((uint64_t)command_header->command_table_base_address);
    memory_set(command_table, 0, sizeof(ahci_hba_command_table_T) + (command_header->prdt_length - 1) * sizeof(ahci_hba_prdt_entry_T));

    command_table->prdt_entry[0].data_base_address       = (uint32_t)(uint64_t)buffer;
    command_table->prdt_entry[0].data_base_address_upper = (uint32_t)((uint64_t)buffer >> 32);
    command_table->prdt_entry[0].byte_count              = (num_sectors << 9) - 1;
    command_table->prdt_entry[0].interrupt_on_completion = 1;


    ahci_fis_register_host_to_device_T* command_fis = (ahci_fis_register_host_to_device_T*)(&command_table->command_fis);
    command_fis->type                   = AHCI_FIS_REGISTER_HOST_TO_DEVICE;
    command_fis->command_control_select = 1;
    command_fis->command                = ATA_CMD_READ_DMA_EX;
    command_fis->lba_0                  = (uint8_t)sector_low;
    command_fis->lba_1                  = (uint8_t)(sector_low >> 8);
    command_fis->lba_2                  = (uint8_t)(sector_low >> 16);
    command_fis->lba_3                  = (uint8_t)sector_high;
    command_fis->lba_4                  = (uint8_t)(sector_high >> 8);
    command_fis->lba_5                  = (uint8_t)(sector_high >> 16);
    command_fis->device                 = 1 << 6;
    command_fis->count_low              = (uint8_t)(num_sectors & 0xFF);
    command_fis->count_high             = (uint8_t)((num_sectors >> 8) & 0xFF);

    port->hba_port->command_issue       = 1;

    while (port->hba_port->command_issue) {
        if (port->hba_port->interrupt_status & HBA_PX_IS_TFES) return false;
    }

    return true;
}

bool ahci_port_write(ahci_port_T* port, uint64_t sector, uint32_t num_sectors, void* buffer_in) {
    DEBUG("AHCI write N/A");
    return false;
}