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

#include "drivers/ahci.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"
};

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]);
    }

    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);
}

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;
}

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;
    }
}