#include <string.h>
#include <vector>

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


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

    uint8_t W25Qxx::init()
    {
        lock_guard<Mutex>lock(*mutex_);

    	powerDown();
        delay_(5);

    	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::getChunkBytes(ChunkSize size)
    {
        switch (size)
        {
        case ChunkSize::PAGE:
            return W25QXX_PAGE_SIZE;

        case ChunkSize::SECTOR:
            return W25QXX_SECTOR_SIZE;

        case ChunkSize::BLOCK_32KB:
            return W25QXX_BLOCK_32_KB_SIZE;

        case ChunkSize::BLOCK_64KB:
            return W25QXX_BLOCK_64_KB_SIZE;
        
        default:
            STA_UNREACHABLE();
            break;
        }
    }

    uint32_t W25Qxx::findLast(std::function<bool(uint8_t*)> criterion, ChunkSize size, uint32_t startAddr /*= 0 */, uint32_t endAddr /* = W25QXX_32B_MEM_SIZE */)
    {
        STA_ASSERT(startAddr <= endAddr);

        lock_guard<Mutex>lock(*mutex_);
        uint32_t bytes = getChunkBytes(size);

        uint32_t left   = startAddr / bytes;
        uint32_t right  = endAddr / bytes; // (addrMode_ == AddressMode::_32BIT ? W25QXX_32B_MEM_SIZE : W25QXX_24B_MEM_SIZE) / bytes;
        uint32_t middle;

        uint8_t * buffer = new uint8_t[bytes];

        while (left < right)
        {
            middle = (left + right) / 2;
            STA_DEBUG_PRINTF("left=%d, middle=%d, right=%d", left, middle, right);
            readData(middle * bytes, buffer, bytes);

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

        middle = (left + right) / 2;
        STA_DEBUG_PRINTLN(middle);

        readData(middle * bytes, buffer, bytes);
        if (criterion(buffer))
        {
            middle += 1;
            STA_DEBUG_PRINTLN("PLUS ONE");
        }

        delete[] buffer;

        return middle * bytes;
    }

    uint8_t W25Qxx::setAddressMode(AddressMode addrMode)
    {
        lock_guard<Mutex>lock(*mutex_);
        busWrite(W25QXX_4_BYTE_ADDR_ENABLE);

        while (isBusy()) {}

        return getAddressMode() == addrMode;
    }

    AddressMode W25Qxx::getAddressMode()
    {   
        lock_guard<Mutex>lock(*mutex_);

        uint8_t status;
        readStatusRegister(3, &status);

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

    uint8_t W25Qxx::getChipID()
    {
        lock_guard<Mutex>lock(*mutex_);

        uint8_t buffer[4];
        busRead(W25QXX_RELEASE_POWER_DOWN, buffer, 3);

        return buffer[3];
    }

    uint8_t W25Qxx::getManufacturerID()
    {
        lock_guard<Mutex>lock(*mutex_);

        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()
    {
        lock_guard<Mutex>lock(*mutex_);

        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()
    {
        lock_guard<Mutex>lock(*mutex_);

        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 */)
    {
        lock_guard<Mutex>lock(*mutex_);

        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
        {
            if (fast)
            {
                uint8_t addrBuffer[4] = {
                    (uint8_t) (address >> 16),
                    (uint8_t) (address >> 8),
                    (uint8_t) (address),
                    0x00 // Dummy byte for fast mode
                };

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

    uint8_t W25Qxx::pageProgram(uint32_t address, uint8_t * buffer, size_t length)
    {
        lock_guard<Mutex>lock(*mutex_);

        STA_ASSERT(length <= W25QXX_PAGE_SIZE);
    	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, 3);
        }
    }

    uint8_t W25Qxx::sectorProgram(uint32_t address, uint8_t * buffer, size_t length)
    {
        lock_guard<Mutex>lock(*mutex_);

        STA_ASSERT(length <= W25QXX_SECTOR_SIZE);

        uint32_t nPages = length / W25QXX_PAGE_SIZE;
        uint32_t remainder = length % W25QXX_PAGE_SIZE;

        uint8_t rslt = 1;

        for (uint8_t i = 0; i < nPages; i++)
            rslt &= pageProgram(address + i * W25QXX_PAGE_SIZE, buffer + i * W25QXX_PAGE_SIZE, W25QXX_PAGE_SIZE);

        rslt &= pageProgram(address + nPages * W25QXX_PAGE_SIZE, buffer + nPages * W25QXX_PAGE_SIZE, remainder);

        return rslt;
    }

    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