#include <string.h>

#include <sta/debug/assert.hpp>
#include <sta/debug/debug.hpp>
#include <sta/drivers/w25qxx.hpp>


namespace sta
{
    W25Qxx::W25Qxx(SPIDevice * device, DelayUsFunc delay, AddressMode addrMode /* = AddressMode::_24BIT */)
        : device_{device},
          delay_{delay},
          state_{ChipState::POWERED_DOWN},
          addrMode_{addrMode}
    {
        STA_ASSERT(device != nullptr);
    }

    uint8_t W25Qxx::init()
    {
    	powerDown();
        delay_(3);

    	if (!releasePowerDown())
		{
			return 0;
		}

        // Check if the chip returns the correct device id.
        if (getManufacturerID() != W25QXX_DEVICE_ID_RESULT)
        {
            return 0;
        }

        // If requested, tell the flash chip to use 32 bit addresses.
        if (addrMode_ == AddressMode::_32BIT)
        {
            if (!setAddressMode(AddressMode::_32BIT))
            {
                return 0;
            }
        }

        return 1;
    }

    uint32_t W25Qxx::findLastPage(std::function<bool(uint8_t*)> criterion)
    {
        uint32_t left   = 0;
        uint32_t right  = 0xFFFFFFFF & (addrMode_ == AddressMode::_32BIT ? 0xFFFFFFFF : 0x00FFFFFF);
        uint32_t middle;
        uint8_t * page = new uint8_t[256];

        while (left < right)
        {
            middle = (left + right) / 2 + 1;
            readData(middle, page, 256);

            if (criterion(page))
            {
                left = middle;
            }
            else
            {
                right = middle-1;
            }
        }

        return middle;
    }

    uint8_t W25Qxx::setAddressMode(AddressMode addrMode)
    {
        busWrite(W25QXX_4_BYTE_ADDR_ENABLE);

        while (isBusy()) {}

        return getAddressMode() == addrMode;
    }

    AddressMode W25Qxx::getAddressMode()
    {   
        uint8_t status;
        readStatusRegister(3, &status);

        return status && 0x01 == 0x01 ? AddressMode::_32BIT : AddressMode::_24BIT; 
    }

    uint8_t W25Qxx::getChipID()
    {
        uint8_t buffer[4];
        busRead(W25QXX_RELEASE_POWER_DOWN, buffer, 3);

        return buffer[3];
    }

    uint8_t W25Qxx::getManufacturerID()
    {
        uint8_t dummy[3] = {0, 0, 0};
        uint8_t id[2] = {0, 0};

        busRead(W25QXX_DEVICE_ID, id, 2, dummy, 3);

        return id[0];
    }

    uint64_t W25Qxx::getUniqueID()
    {
        uint8_t dummy[4];
        uint8_t id[8];

        busRead(W25QXX_READ_UNIQUE_ID, id, 8, dummy, 4);

        uint64_t id_complete = 0;

        for (size_t i; i < 8; i++)
        {
            id_complete |= id[i] << (7-i) * 8;
        }
        
        return id_complete;
    }

    bool W25Qxx::isBusy()
    {
        uint8_t status = 0;
        readStatusRegister(1, &status);

        return (0x01 & status) == 0x01;
    }

    uint8_t W25Qxx::readData(uint32_t address, uint8_t * buffer, size_t length, bool fast /* = true */)
    {
        uint8_t instruction = fast ? W25QXX_FAST_READ : W25QXX_READ;

        // In fast mode we have to send a 8 dummy clock cycles first.
        if (fast)
        {
            // TODO
        }

        while (isBusy()) {}

        // Depending on address mode, send 3 bytes or 4 bytes.
        if (addrMode_ == AddressMode::_32BIT)
        {
            uint8_t addrBuffer[4] = {
                (uint8_t) (address << 24),
                (uint8_t) (address << 16),
                (uint8_t) (address << 8),
                (uint8_t) (address)
            };

            return busRead(instruction, buffer, length, addrBuffer, 4);
        }
        else
        {
            uint8_t addrBuffer[3] = {
                (uint8_t) (address << 16),
                (uint8_t) (address << 8),
                (uint8_t) (address)
            };

            return busRead(instruction, buffer, length, addrBuffer, 4);
        }
    }

    uint8_t W25Qxx::pageProgram(uint32_t address, uint8_t * buffer, size_t length)
    {
    	while (isBusy()) {}

        if (!writeEnable())
        {
            return 0;
        }

        while (!isWriteEnabled()) {}

        if (addrMode_ == AddressMode::_32BIT)
        {
            uint8_t addrBuffer[4] = {
                (uint8_t) (address << 24),
                (uint8_t) (address << 16),
                (uint8_t) (address << 8),
                (uint8_t) (address)
            };

            return busWrite(W25QXX_PAGE_PROGAM, buffer, length, addrBuffer, 4);
        }
        else
        {
            uint8_t addrBuffer[3] = {
                (uint8_t) (address << 16),
                (uint8_t) (address << 8),
                (uint8_t) (address)
            };

            return busWrite(W25QXX_PAGE_PROGAM, buffer, length, addrBuffer, 4);
        }
    }

    uint8_t W25Qxx::busWrite(uint8_t instruction, const uint8_t * data /* = nullptr */, size_t length /* = 0 */, uint8_t * arguments /* = nullptr */, size_t arg_length /* = 0 */, uint32_t delayCsHigh /*= 0 */)
    {
        device_->beginTransmission();

        // Send the instruction.
        device_->transfer(instruction);

        // If requested, send argument bytes before the actual data.
        if (arguments != nullptr && arg_length != 0)
        {
            device_->transfer(arguments, arg_length);
        }

        // If necessary, send the actual data bytes.
        if (data != nullptr && length != 0)
        {
            device_->transfer(data, length);
        }

        if (delayCsHigh != 0)
        {
            delay_(2);
        }

        device_->endTransmission();

        return 1;
    }

    uint8_t W25Qxx::busRead(uint8_t instruction, uint8_t * data, size_t length, uint8_t * arguments /* = nullptr */, size_t arg_length /* = 0 */)
    {
        device_->beginTransmission();

        // Send the instruction.
        device_->transfer(instruction);

        // If requested, send argument bytes before receiving the actual data.
        if (arguments != nullptr && arg_length != 0)
        {
            device_->transfer(arguments, arg_length);
        }

        // Send the actual data bytes.
        device_->receive(data, length);

        device_->endTransmission();

        return 1;
    }

    uint8_t W25Qxx::writeEnable()
    {
        return busWrite(W25QXX_WRITE_ENABLE);
    }

    uint8_t W25Qxx::writeVolatileEnable()
    {
        return busWrite(W25QXX_VOL_SR_WRITE_ENABLE);
    }

    uint8_t W25Qxx::writeDisable()
    {
        return busWrite(W25QXX_WRITE_DISABLE);
    }

    uint8_t W25Qxx::readStatusRegister(uint8_t regID, uint8_t * status_byte)
    {
        if (regID == 1)
        {
            return busRead(W25QXX_STATUS_REG_1_READ, status_byte, 1);
        }

        if (regID == 2)
        {
            return busRead(W25QXX_STATUS_REG_2_READ, status_byte, 1);
        }

        if (regID == 3)
        {
            return busRead(W25QXX_STATUS_REG_3_READ, status_byte, 1);
        }

        return 0;
    }

    uint8_t W25Qxx::writeStatusRegister(uint8_t regID, uint8_t * status_byte, bool nonvolatile /* = false */)
    {
        if (!writeEnable())
        {
            return 0;
        }

        while (!isWriteEnabled()) {}

        if (!nonvolatile)
        {
            uint8_t rslt = writeVolatileEnable();

            if (rslt == 0)
            {
                return 0;
            }
        }

        if (regID == 1)
        {
            return busRead(W25QXX_STATUS_REG_1_READ, status_byte, 1);
        }

        if (regID == 2)
        {
            return busRead(W25QXX_STATUS_REG_2_READ, status_byte, 1);
        }

        if (regID == 3)
        {
            return busRead(W25QXX_STATUS_REG_3_READ, status_byte, 1);
        }

        return 0;
    }

    bool W25Qxx::isWriteEnabled()
    {
    	uint8_t status = 0;
        readStatusRegister(1, &status);

        return (0x02 & status) == 0x02;
    }

    uint8_t W25Qxx::sectorErase(uint32_t address)
    {
        if (!writeEnable())
        {
            return 0;
        }

        while (!isWriteEnabled()) {}

        if (addrMode_ == AddressMode::_32BIT)
        {
            uint8_t addrBuffer[4] = {
                (uint8_t) (address << 24),
                (uint8_t) (address << 16),
                (uint8_t) (address << 8),
                (uint8_t) (address)
            };

            return busWrite(W25QXX_SECTOR_ERASE, addrBuffer, 4);
        }
        else
        {
            uint8_t addrBuffer[3] = {
                (uint8_t) (address << 16),
                (uint8_t) (address << 8),
                (uint8_t) (address)
            };

            return busWrite(W25QXX_SECTOR_ERASE, addrBuffer, 3);
        }

        return 1;
    }

    uint8_t W25Qxx::blockErase(uint32_t address, BlockSize blockSize)
    {
        if (!writeEnable())
        {
            return 0;
        }

        while (!isWriteEnabled()) {}

        uint8_t instruction = blockSize == BlockSize::_32KB ? W25QXX_BLOCK_ERASE_32_KB : W25QXX_BLOCK_ERASE_64_KB;

        if (addrMode_ == AddressMode::_32BIT)
        {
            uint8_t addrBuffer[4] = {
                (uint8_t) (address << 24),
                (uint8_t) (address << 16),
                (uint8_t) (address << 8),
                (uint8_t) (address)
            };

            return busWrite(instruction, addrBuffer, 4);
        }
        else
        {
            uint8_t addrBuffer[3] = {
                (uint8_t) (address << 16),
                (uint8_t) (address << 8),
                (uint8_t) (address)
            };

            return busWrite(instruction, addrBuffer, 3);
        }

        return 1;
    }

    uint8_t W25Qxx::chipErase()
    {
        if (!writeEnable())
        {
            return 0;
        }

        return busWrite(W25QXX_CHIP_ERASE);
    }

    uint8_t W25Qxx::suspendErase()
    {
        return busWrite(W25QXX_ERASE_SUSPEND_PROG);
    }

    uint8_t W25Qxx::resumeErase()
    {
        return busWrite(W25QXX_ERASE_RESUME_PROG);
    }

    uint8_t W25Qxx::powerDown()
    {
        return busWrite(W25QXX_POWER_DOWN);
    }

    uint8_t W25Qxx::releasePowerDown()
    {
        if (state_ == ChipState::POWERED_DOWN)
        {
            return busWrite(W25QXX_RELEASE_POWER_DOWN, nullptr, 0, nullptr, 0, 2);
        }

        return 0;
    }

    uint8_t W25Qxx::lockBlock(uint32_t address)
    {
        if (!writeEnable())
        {
            return 0;
        }

        while (!isWriteEnabled()) {}

        uint8_t addrBuffer[3] = {
            (uint8_t) (address << 16),
            (uint8_t) (address << 8),
            (uint8_t) (address)
        };

        return busWrite(W25QXX_INDIV_BLOCK_LOCK, addrBuffer, 3);
    }

    uint8_t W25Qxx::unlockBlock(uint32_t address)
    {
        if (!writeEnable())
        {
            return 0;
        }

        while (!isWriteEnabled()) {}

        uint8_t addrBuffer[3] = {
            (uint8_t) (address << 16),
            (uint8_t) (address << 8),
            (uint8_t) (address)
        };

        return busWrite(W25QXX_INDIV_BLOCK_UNLOCK, addrBuffer, 3);
    }
} // namespace sta