#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/mman.h>
#include "sd-spi.h"

#define MAP_SIZE            4096UL
// Bit Manipulation Macros
#define SET_LSBITS(len) ((1 << (len)) - 1) // set the first len bits from right to left
#define GET_BITS_RANGE(x, up_pos, low_pos) ((x & (SET_LSBITS(up_pos-low_pos+1) << low_pos)) >> low_pos)
// QSPI Register Offsets (>>2 since will be 4-byte addressed)
#define SOFTWARE_RESET_REG  (0x40>>2)   // Software Reset Register (SRR)
#define SPI_CONTROL_REG     (0x60>>2)   // SPI Control Register (SPICR)
#define SPI_STATUS_REG      (0x64>>2)   // SPI Status Register (SPISR)
#define SPI_DATA_TX_REG     (0x68>>2)   // SPI Data Transmit Register (Single or FIFO) (SPI DTR)
#define SPI_DATA_RX_REG     (0x6C>>2)   // SPI Data Receive Register (Single or FIFO) (SPI DRR)
#define SPI_SLV_SEL_REG     (0x70>>2)   // SPI Slave Select Register (SPISSR)
#define SPI_TX_FIFO_OCC_REG (0x74>>2)   // SPI Transmit FIFO Occupancy Register (SPI Transmit FIFO Occupancy Register)
#define SPI_RX_FIFO_OCC_REG (0x78>>2)   // SPI Receive FIFO Occupancy Register (SPI Receive FIFO Occupancy Register)
// QSPI Base Addresses
#define SPI0_BASE           0xA0050000  // SPI0 Base Address
#define SPI1_BASE           0xA0060000  // SPI1 Base Address
#define PL1_BASE            0xFF5E0000
// QSPI Status Register Masks
#define SPI_COMMAND_ERR     (1<<10)     // 1 == core configured in dual/quad SPI and has SPI DTR FIFO mismatch
#define SPI_LOOPBACK_ERR    (1<<9)      // 1 == spi command, address, and data are non-standard spi and loopback enabled
#define SPI_MSB_ERR         (1<<8)      // 1 == core configured in dual/quad SPI and MSB mode is disabled
#define SPI_SLV_MODE_ERR    (1<<7)      // 1 == core configured in dual/quad SPI and set to slave mdoe
#define SPI_CPOL_CPHA_ERR   (1<<6)      // 1 == CPOL and CPHA are 01 or 10
#define SPI_SLV_MODE_SEL    (1<<5)      // 0 == core is configured in slave mode and selected by external SPI master (1 == standard mode)
#define SPI_MODF_ERR        (1<<4)      // 1 == SS signal goes active while SPI device is configured as master (mode-fault error)
#define SPI_TX_FULL         (1<<3)      // 1 == TX FIFO is full
#define SPI_TX_EMPTY        (1<<2)      // 1 == TX FIFO is empty
#define SPI_RX_FULL         (1<<1)      // 1 == RX FIFO is full
#define SPI_RX_EMPTY        (1<<0)      // 1 == RX FIFO is empty
// QSPI Peripheral Info
#define NUM_SLAVES          2
#define FIFO_SIZE           256
#define TIMEOUT_US          1000        // Timeout time in uSeconds
// SD Card Commands
#define CMD0                0           // Software reset
#define CMD1                1           // Initiate initialization process
#define CMD8                8           // For SDC V2, check voltage range
#define CMD16               16          // Change R/W block length
#define CMD17               17          // Read a block
#define CMD24               24          // Write a block
#define ACMD41              41          // For SDC, initiate initialization process
#define CMD55               55          // Leading command to ACMDs
#define CMD58               58          // Read OCR

uint32_t GetQSPIPeripheralStatus(volatile uint32_t* spix_regs) {
    volatile uint32_t* spisr = spix_regs + SPI_STATUS_REG;
    uint32_t status = *spisr;
    return status;
}

// Reset QSPI peripheral
void ResetQSPIPeripheral(volatile uint32_t* spix_regs) {
    volatile uint32_t* srr = spix_regs + SOFTWARE_RESET_REG;
    // Writing 0x0000000A to the SRR resets the core register for four AXI clock cycles
    *srr = 0x0000000A;
    usleep(5);
}

// Initialize QSPI peripheral
void InitQSPIPeripheral(volatile uint32_t* spix_regs) {
    volatile uint32_t* spicr = spix_regs + SPI_CONTROL_REG;
    uint32_t init = 0;
    init = init & ~(0 << 9);    // MSB file transfer format (0)
    init = init & ~(0 << 8);    // Enable master transactions
    init = init | (1 << 7);     // Manual slave select assertion
    init = init | (1 << 6);     // RX FIFO reset (automatically set back to 0 one clock cycle later)
    init = init | (1 << 5);     // TX FIFO reset (automatically set back to 0 one clock cycle later)
    init = init & ~(0 << 4);    // Clock phase (data is valid on first SCK edge) (CPHA=0 for MMC/SDC)
    init = init & ~(0 << 3);    // Clock polarity (active high clock) (CPOL=0 for MMC/SDC)
    init = init | (1 << 2);     // Master configuration
    init = init & ~(0 << 1);    // Disable SPI system (will immediately transmit if data placed in fifo when enabled)
    init = init & ~(0 << 0);    // Normal operation (not loop-back mode)
    *spicr = init;
    usleep(5);
}

// Disable QSPI peripheral
void DisableQSPIPeripheral(volatile uint32_t* spix_regs) {
    volatile uint32_t* spicr = spix_regs + SPI_CONTROL_REG;
    *spicr = *spicr & ~(0 << 1);
    usleep(20);
}

// Enable QSPI peripheral
void EnableQSPIPeripheral(volatile uint32_t* spix_regs) {
    volatile uint32_t* spicr = spix_regs + SPI_CONTROL_REG;
    *spicr = *spicr | (1 << 1);
    usleep(20);
}

// Read/Write a single byte to/from the QSPI peripheral fifo
// Since SPI is full-duplex, every transaction is a read/write
// Ignore either the read/write depending on the transaction
// Returns 1 on error, 0 on success
uint32_t SPIReadWriteByte(volatile uint32_t* spix_regs, uint8_t sendbyte, uint8_t* readbyte) {
    uint32_t status = GetQSPIPeripheralStatus(spix_regs);
    if (status & SPI_TX_FULL) return 1; // TX is full, cannot transmit
    
    volatile uint32_t* datatx = spix_regs + SPI_DATA_TX_REG;
    volatile uint32_t* datarx = spix_regs + SPI_DATA_RX_REG;

    *datatx = sendbyte;

    // Wait until the response reaches the FIFO
    status = GetQSPIPeripheralStatus(spix_regs);
    // if (status & SPI_RX_EMPTY) *readbyte = 0xFF;
    // else *readbyte = (*datarx) & 0xFF;
    while (status & SPI_RX_EMPTY) {
        status = GetQSPIPeripheralStatus(spix_regs);
    };

    *readbyte = (*datarx) & 0xFF;

    return 0;
}

// Select the slave device
// Returns 1 on error (slv number is invalid)
uint32_t SPISelectSlave(volatile uint32_t* spix_regs, uint8_t slv) {
    if (slv > NUM_SLAVES) return 1;
    // At most a single bit can be low and this is the bit the master communicates with
    uint32_t slv_sel = ~(1 << slv);

    volatile uint32_t* spissr = spix_regs + SPI_SLV_SEL_REG;
    *spissr = slv_sel;
    return 0;
}

// Get the number of elements in the transmit FIFO - 1 (occupancy - 1)
uint32_t GetNumTXFIFOElements(volatile uint32_t* spix_regs) {
    volatile uint32_t* tx_occ = spix_regs + SPI_TX_FIFO_OCC_REG;
    uint32_t fifo_mask = (FIFO_SIZE - 1);
    uint32_t tx_num = (*tx_occ) & fifo_mask;
    return tx_num;
}

// Get the number of elements in the receive FIFO - 1 (occupancy - 1)
uint32_t GetNumRXFIFOElements(volatile uint32_t* spix_regs) {
    volatile uint32_t* rx_occ = spix_regs + SPI_RX_FIFO_OCC_REG;
    uint32_t fifo_mask = (FIFO_SIZE - 1);
    uint32_t rx_num = (*rx_occ) & fifo_mask;
    return rx_num;
}

// Send a command to the SD Card
// Returns the response from the SD Card
uint8_t SDSPISendCommand(volatile uint32_t* spix_regs, uint8_t cmd, uint32_t args, uint8_t slv) {
    bool retry = 1;
    uint8_t byte_response;
    while (retry) {
        retry = 0;
        // Command frame consists of the following
        // 8-bit command ({0,1,cmd})
        // 32-bit args
        // 8-bit crc (not checked in SPI mode)
        if (cmd > 63) return 1;
        uint8_t crc = 0x01; // crc is irrelevant for all commands except CMD0/CMD8
        if (cmd == CMD0) {
            crc = 0x95; // valid CRC for CMD0(0)
        } else if (cmd == CMD8) {
            crc = 0x87; // valid CRC for CMD8(0x1AA)
        }

        // Disable the QSPI peripheral to select slave
        // DisableQSPIPeripheral(spix_regs);

        // Enable the QSPI peripheral for transmit/receive
        EnableQSPIPeripheral(spix_regs);

        // Disable the current slv device (set SS to high for 10ms)
        // Easy way to achieve this is to enable a different slv
        uint8_t otherslv;
        if (slv == 0) otherslv = 1;
        else otherslv = slv - 1;
        SPISelectSlave(spix_regs, otherslv);

        // Pulse the clock while the device is disabled
        // Send at least 74 clock pulses
        // Will send 80 with 10 8-bit messages
        uint8_t dummy;
        for (uint8_t i = 0; i < 10; i++) {
            SPIReadWriteByte(spix_regs, 0xFF, &dummy);
        }

        // Enable the slv device (set SS to low)
        SPISelectSlave(spix_regs, slv);

        // Send the command, ignore the data received while the command is being sent
        uint8_t cmd_data = (1 << 6) | cmd;
        SPIReadWriteByte(spix_regs, cmd_data, &byte_response);                      // {0, 1, cmd}
        SPIReadWriteByte(spix_regs, GET_BITS_RANGE(args, 31, 24), &byte_response);  // arg[31:24]
        SPIReadWriteByte(spix_regs, GET_BITS_RANGE(args, 23, 16), &byte_response);  // arg[23:16]
        SPIReadWriteByte(spix_regs, GET_BITS_RANGE(args, 15, 8), &byte_response);   // arg[15:8]
        SPIReadWriteByte(spix_regs, GET_BITS_RANGE(args, 7, 0), &byte_response);    // arg[7:0]
        SPIReadWriteByte(spix_regs, crc, &byte_response);                           // CRC, only checked or CMD0 and CMD8 in SPI mode

        // Wait for the response (transmit 0xFF) (will only take < 10 messages)
        byte_response = 0xFF;
        for (uint8_t i = 0; i < 10; i++) {
            SPIReadWriteByte(spix_regs, 0xFF, &byte_response);
            if (byte_response != 0xFF) break;
        }
        
        if (byte_response & 0x80) {
            printf("ERROR: Timed out, retrying (response: %X)\n", byte_response);
            retry = 1;
        } else {
            retry = 0;
        }
    }

    return byte_response;
}

// Send CMD0 to an SD card (Software reset)
uint8_t SDSPISoftwareReset(volatile uint32_t* spix_regs, uint8_t slv) {
    // Send CMD0 with CS LOW (IMPORTANT!!!! :D :3 meow)
    // Card samples CS signal on a CMD0 (if CS is low, the card enters SPI mode and responds with R1 with Idle state set)
    // This is the ONLY SPI command that MUST have a valid CRC (CRC is disabled once SPI mode is entered)
    return SDSPISendCommand(spix_regs, CMD0, 0x0, slv); // no args for CMD0
}

// Send CMD1 to an SD card (Begin initialization)
// Only used for MMC v3 devices
uint8_t SDSPIBeginInitialize(volatile uint32_t* spix_regs, uint8_t slv) {
    return SDSPISendCommand(spix_regs, CMD1, 0x0, slv);
}

// Send CMD8 to an SD card (check volatage range, only for SDC V2)
uint8_t SDSPICheckVoltageRange(volatile uint32_t* spix_regs, uint8_t slv) {
    uint32_t arg = 0x1AA; // 2.7 to 3.3V voltage range support
    return SDSPISendCommand(spix_regs, CMD8, arg, slv);
}

// Send CMD16 to an SD card (set block length, will ALWAYS be 512)
uint8_t SDSPISetBlockLength512(volatile uint32_t* spix_regs, uint8_t slv) {
    uint32_t arg = 0x200; // 512 bytes
    return SDSPISendCommand(spix_regs, CMD16, arg, slv);
}

// Send ACMD41 to an SD card (check for high capacity)
uint8_t SDSPIACMD41(volatile uint32_t* spix_regs, uint32_t arg, uint8_t slv) {
    // ACMD implies sending CMD55 prior to sending the next command
    SDSPISendCommand(spix_regs, CMD55, 0x0, slv);
    
    return SDSPISendCommand(spix_regs, ACMD41, arg, slv);
}

// Send CMD58 to an SD card
// This will read the OCR
uint32_t SDSPIReadOCR(volatile uint32_t* spix_regs, uint8_t slv) {
    uint8_t status = SDSPISendCommand(spix_regs, CMD58, 0x0, slv);

    uint8_t byte;
    uint32_t ocr = 0;
    SPIReadWriteByte(spix_regs, 0xFF, &byte);
    ocr = ocr | (byte << 24);
    SPIReadWriteByte(spix_regs, 0xFF, &byte);
    ocr = ocr | (byte << 16);
    SPIReadWriteByte(spix_regs, 0xFF, &byte);
    ocr = ocr | (byte << 8);
    SPIReadWriteByte(spix_regs, 0xFF, &byte);
    ocr = ocr | (byte);

    return ocr;
}

// Begin the initialization flow for an SD card
// Returns 1 on error, 0 on success
// http://elm-chan.org/docs/mmc/m/sdinit.png
uint8_t SDSPIInit(volatile uint32_t* spix_regs, uint8_t slv) {
    // This routine should never be called immediately after plugging in an sd card
    // Must wait at least 1ms after power on

    // Issue CMD0 (Software reset)
    // Read the response from the SD card (CMD0 responds in R1 form)
    uint8_t response;
    response = SDSPISoftwareReset(spix_regs, slv);

    // If the response is 0x1 (idle state) then issue CMD8
    // otherwise, the card is unknown and should not be accessed
    if (response != 0x1) {
        printf("ERROR: SD Card FAILED to respond successfully to CMD0\n");
        return 1;
    }
    
    printf("INFO: INITIALIZATION: CMD0 SUCCESS\n");

    // Issue CMD8 if the card successfully entered the idle state
    // Read the response from the SD card (CMD8 responds in R7 form)
    response = SDSPICheckVoltageRange(spix_regs, slv);

    if (response != 0x1) {
        printf("ERROR: SD Card FAILED to respond successfully to CMD8\n");
        return 1;
    }
    
    printf("INFO: INITIALIZATION: CMD8 SUCCESS\n");

    // Read the rest of the response (R7 contains 32-bit OCR)
    // R7 = R1 + OCR (32 bit)
    uint8_t byte;
    uint32_t ocr = 0;
    SPIReadWriteByte(spix_regs, 0xFF, &byte);
    ocr = ocr | (byte << 24);
    SPIReadWriteByte(spix_regs, 0xFF, &byte);
    ocr = ocr | (byte << 16);
    SPIReadWriteByte(spix_regs, 0xFF, &byte);
    ocr = ocr | (byte << 8);
    SPIReadWriteByte(spix_regs, 0xFF, &byte);
    ocr = ocr | (byte);

    bool sendCMD16 = false;
    if ((ocr & 0xFFF) == 0x1AA) {
        // Issue ACMD41(0x40000000)
        printf("INFO: INITIALIZATION: OCR == 0x1AA, send ACMD41(0x40000000)\n");
        response = 0x1;
        while (response == 0x1) {
            response = SDSPIACMD41(spix_regs, 0x40000000, slv); // set HCS flag, to check if the card is SDHC/SDXC
        }

        if (response == 0x0) {
            printf("INFO: INITIALIZATION: ACMD41(0x40000000) SUCCESS\n");
            
            // Send CMD58(0x00000000)
            uint32_t ocr = SDSPIReadOCR(spix_regs, slv);

            // Check CCS bit in OCR
            if (ocr & (1 << 30)) {
                // SD Ver 2+ (block address)
                printf("INFO: INITIALIZATION: SD Ver 2+ (block address)\n");
                printf("INFO: INITIALIZATION: Initialization completed successfully");
                return 0;
            } else {
                sendCMD16 = true;
                printf("INFO: INITIALIZATION: SD Ver 2+ (byte address), continuing...\n");
            }
        }
    } else {
        // Issue ACMD41(0x00000000)
        response = 0x1;
        while (response == 0x1) {
            response = SDSPIACMD41(spix_regs, 0x00000000, slv);
        }

        if (response == 0x0) {
            // SD Ver. 1
            printf("INFO: INITIALIZATION: SD Ver. 1 detected\n");
            sendCMD16 = true;
        } else {
            response = SDSPIBeginInitialize(spix_regs, slv);
            while (response == 0x1) {
                response = SDSPIACMD41(spix_regs, 0x00000000, slv);
            }

            if (response == 0x0) {
                // MMC Ver 3
                printf("INFO: INITIALIZATION: MMC Ver 3 detected\n");
                sendCMD16 = true;
            } else {
                printf("ERROR: SD Card FAILED to respond successfully to ACMD41\n");
                return 1;
            }
        }
    }

    if (sendCMD16) {
        // SD Ver 2+ (byte address), SD Ver 1, or MMC Ver 3
        // send CMD16(0x200)
        SDSPISetBlockLength512(spix_regs, slv);
        printf("INFO: INITIALIZATION: Initialization completed successfully");
        return 0;
    } else {
        printf("ERROR: Initialization should never reach here\n");
        return 1;
    }
}

// Send CMD17 to an SD card (read a single block)
// The buffer passed to the function MUST have a length of at least 512
// Return 1 on failure (0 on success)
uint8_t SDSPIReadOneBlock(volatile uint32_t* spix_regs, uint32_t addr, uint8_t* buff, uint8_t slv) {
    uint32_t arg = addr;
    uint8_t status = SDSPISendCommand(spix_regs, CMD17, arg, slv);

    if (status != 0x00) {
        printf("ERROR: SD Card FAILED to respond successfully to CMD17\n");
        return 1;
    }

    // Wait for the sd card to begin sending the packet
    uint8_t byteresponse = 0xFF;
    uint8_t expectedtoken = 0xFE; // data token for cmd 17
    while (byteresponse != expectedtoken) {
        SPIReadWriteByte(spix_regs, 0xFF, &byteresponse);
    }

    // Read the 512 bytes
    for (uint16_t i = 0; i < 512; i++) {
        SPIReadWriteByte(spix_regs, 0xFF, &byteresponse);
        buff[i] = byteresponse;
    }

    // Receive CRC (discard)
    SPIReadWriteByte(spix_regs, 0xFF, &byteresponse);
    SPIReadWriteByte(spix_regs, 0xFF, &byteresponse);

    printf("INFO: Done reading block\n");

    return 0;
}

// Send CMD24 to an SD card
// The buffer passed to the function MUST have a length of at least 512
// Return 1 on failure (0 on success)
uint8_t SDSPIWriteOneBlock(volatile uint32_t* spix_regs, uint32_t addr, uint8_t* buff, uint8_t slv) {
    uint32_t arg = addr;
    uint8_t status = SDSPISendCommand(spix_regs, CMD24, arg, slv);

    if (status != 0x00) {
        printf("ERROR: SD Card FAILED to respond successfully to CMD24\n");
        return 1;
    }

    usleep(5); // wait at least one byte of transfer before sending data packet

    uint8_t byteresponse;
    SPIReadWriteByte(spix_regs, 0xFE, &byteresponse); // data token for cmd 24
    for (uint16_t i = 0; i < 512; i++) {
        while (GetQSPIPeripheralStatus(spix_regs) & SPI_TX_FULL) {};
        SPIReadWriteByte(spix_regs, buff[i], &byteresponse);
    }

    // Send CRC (anything)
    SPIReadWriteByte(spix_regs, 0xFF, &byteresponse);
    SPIReadWriteByte(spix_regs, 0xFF, &byteresponse);

    // Get response
    SPIReadWriteByte(spix_regs, 0xFF, &byteresponse);

    if ((byteresponse & 0x1F) != 0x5) {
        printf("ERROR: SD Card FAILED to write block successfully\n");
        return 1;
    }

    // Wait until the device stops sending low (busy)
    uint8_t data = 0;
    while (data == 0x0) {
        SPIReadWriteByte(spix_regs, 0xFF, &data);
    }

    printf("INFO: Device is no longer busy, done writing block\n");

    return 0;
}

// Run the full initialization sequence as shown in the example application
// Sets PL1 frequency, resets QSPI peripheral, initializes QSPI peripheral, calls SDCard initialization
// Will return the addresses of the 2 QSPI device registers (in parameters)
// Return 1 on failure (0 on success)
uint8_t SDSPIFullInitialization(volatile uint32_t** spi0_regs_ret, volatile uint32_t** spi1_regs_ret) {
    int dh = open("/dev/mem", O_RDWR | O_SYNC);
    if(dh == -1) {
		printf("Unable to open /dev/mem.  Ensure it exists (major=1, minor=1)\n");
		printf("Must be root to run this routine.\n");
		return 1;
    }

    uint32_t spi0_status, spi1_status;
    volatile uint32_t* spi0_regs = mmap(NULL, MAP_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, dh, SPI0_BASE);
    volatile uint32_t* spi1_regs = mmap(NULL, MAP_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, dh, SPI1_BASE);

    // Set PL1 to 3.2MHz (SPI has 16 div, so should run at 200K)
    printf("INFO: Setting PL1 to ~3.2MHz, so QSPI peripheral runs ~200KHz\n");
    volatile uint32_t* pl_clk_regs = mmap(NULL, MAP_SIZE, PROT_READ|PROT_WRITE, MAP_SHARED, dh, PL1_BASE);
    volatile uint32_t* PL1 = pl_clk_regs + (0xC4 >> 2);

    *PL1 = (1 << 24) // enable
	| (15 << 16) // divisor 1
	| (30 << 8); // divisor 2

    //////////////////////////////////////////////////////////////////////
    // SD CARD BLOCK WRITE/READ TEST
    //////////////////////////////////////////////////////////////////////
    // Reset and initialize SPI peripheral
    printf("INFO: Resetting QSPI Peripheral\n");
    ResetQSPIPeripheral(spi0_regs);
    printf("INFO: Initializing QSPI Peripheral\n");
    InitQSPIPeripheral(spi0_regs); // initialize in disabled state;

    *spi0_regs_ret = spi0_regs;
    *spi1_regs_ret = spi1_regs;

    return 0;
}