/*
 * Copyright (C) 2020-2021 Aloysius Indrayanto
 *
 * This file is part of the eCxx library, see LICENSE file for the license details.
 */

/*
 * ----------------------------------------------------------------------------------
 * The code in this file is based on:
 *     TWI/I2C library for Wiring & Arduino
 * ----------------------------------------------------------------------------------
 * Copyright (C) 2006 Nicholas Zambetti. All right reserved.
 *
 * This library is free software; you can redistribute it and/or modify it under the
 * terms of the GNU Lesser General Public License as published by the Free Software
 * Foundation; either version 2.1 of the License, or (at your option) any later
 * version.
 *
 * This library is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
 * PARTICULAR PURPOSE. See the GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License along
 * with this library; if not, write to the Free Software Foundation, Inc., 51
 * Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
 * ----------------------------------------------------------------------------------
 * In addition information from these sources may also have been used:
 *
 *     http://ww1.microchip.com/downloads/en/appnotes/00002480a.pdf
 *     http://ww1.microchip.com/downloads/en/appnotes/atmel-2565-using-the-twi-module-as-i2c-slave_applicationnote_avr311.pdf
 *
 *     https://www.exploreembedded.com/wiki/Basics_of_I2C_with_AVR
 *     https://www.engineersgarage.com/avr-microcontroller/how-to-use-i2c-twi-two-wire-interface-in-avr-atmega32-part-36-46
 *     http://www.chrisherring.net/all/tutorial-interrupt-driven-twi-interface-for-avr-part1
 * ----------------------------------------------------------------------------------
 */


#include <util/twi.h>

#include "../AVRCommon/Internal.h"
#include "TWIDriver.h"

#include <eCxx/Platform/AVR/PinMapping.h>


// eCxx only supports one TWI controller in ATmega328PB
#ifdef __AVR_ATmega328PB__
    #define TWAR     TWAR0
    #define TWBR     TWBR0
    #define TWCR     TWCR0
    #define TWDR     TWDR0
    #define TWSR     TWSR0
    #define TWI_vect TWI0_vect
    // According to the datasheets, the position of TWI1's configuration bits are the same as TWI0's
    #if (TWEN != TWEN1 || TWEA != TWEA1 || TWIE != TWIE1 || TWINT != TWINT || TWSTA != TWSTA1 || TWPS1 != TWPS11 || TWPS0 != TWPS10)
    #error "TWI1's configuration bits position are not the same as TWI0's"
    #endif
#endif


// Workaround for missing TWI macros in some versions of AVR GCC
#ifndef TWSR
    #if defined(__AVR_ATmega16U4__ )
        #define TWBR  _SFR_MEM8(0xB8)
        #define TWSR  _SFR_MEM8(0xB9)
        #define TWPS0 0
        #define TWPS1 1
        #define TWS3  3
        #define TWS4  4
        #define TWS5  5
        #define TWS6  6
        #define TWS7  7
        #define TWAR  _SFR_MEM8(0xBA)
        #define TWDR  _SFR_MEM8(0xBB)
        #define TWCR  _SFR_MEM8(0xBC)
        #define TWIE  0
        #define TWEN  2
        #define TWWC  3
        #define TWSTO 4
        #define TWSTA 5
        #define TWEA  6
        #define TWINT 7
    #endif
#endif


namespace eCxx {


//
// Allow the ISR code to not be linked when not explictly used
//
static void (*twi_isr_handler)() = nullptr;


//
// A standard AVR TWI driver
//
void TWIDriver::begin(uint8_t address, SlaveRxCallbackF rxucbSlave, SlaveTxCallbackF txucbSlave, uint32_t frequency)
{
    // Safe a pointer to the real ISR handler
    twi_isr_handler = &TWIDriver::__twi_isr_handler;

    // Initialize variables
    _state         = TWIState::Ready;
    _sendStop      = true;
    _repeatedStart = false;

    _resultCode    = TWIResultCode::NoError;

    _rxucbSlave    = rxucbSlave;
    _txucbSlave    = txucbSlave;

    // Activate internal pullups
    Application::setPinMode(twiPinSDA, Application::PinMode::InputPullUp);
    Application::setPinMode(twiPinSCL, Application::PinMode::InputPullUp);

    // Ensure the frequency is within range
         if(frequency <   1000UL) frequency =   1000UL;
    else if(frequency > 400000UL) frequency = 400000UL;

    // TWI clock frequency calculation
    //
    // TWPS1   TWPS0   Prescaler
    // 0       0        1
    // 0       1        4
    // 1       0       16
    // 1       1       64
    //
    // SCL frequency = F_CPU / ( 16 + ( 2 * TWBR * Prescaler ) )

    // Determine the divider and prescaler
    uint32_t cpuFreq = Application::getCPUFrequency();
    uint16_t divider = (cpuFreq / frequency - 16) / 2;
    uint8_t  twps    = 0;
    while(divider > 255) {
        ++twps;
        cpuFreq /= 4;
        divider  = (cpuFreq / frequency - 16) / 2;
    }
    if(twps > 3) {
        twps    =   3;
        divider = 255;
    }

    // Set the prescaler
    TWSR |= (twps << TWPS0);

    // Set the divider
    TWBR = static_cast<uint8_t>(divider);

    // Enable TWI, acknowledge, and interrupt
    TWCR = _BV(TWEN) | _BV(TWEA) | _BV(TWIE);

    // Set (slave) address
    if(address) TWAR = address << 1;
}


void TWIDriver::end()
{
    // Disable TWI, acknowledge, and interrupt
    TWCR &= ~( _BV(TWEN) | _BV(TWEA) | _BV(TWIE) );

    // Deactivate internal pullups
    Application::setPinMode(twiPinSDA, Application::PinMode::Input);
    Application::setPinMode(twiPinSCL, Application::PinMode::Input);
}


TWIDriver::TWIResultCode TWIDriver::writeTo(uint8_t address, const uint8_t* data, uint8_t len, bool wait, bool sendStop)
{
    // Check if there will be enough space in the buffer
    if(_msBuff.capacity() < len) return TWIResultCode::DataTooLong;

    // Wait until ready
    uint32_t timeoutDelay = Application::CommunicationTimeoutCycles;

    while(_state != TWIState::Ready) {
        Application::yield();
        if(--timeoutDelay == 0) return TWIResultCode::TimeOut;
    }

    // Make it a master transmitter
    _state       = TWIState::MTx;
    _sendStop    = sendStop;

    _resultCode  = TWIResultCode::NoError;

    // Copy data to buffer
    _msBuff.clear();
    for(uint8_t i = 0; i < len; ++i) _msBuff.put(*data++);

    // Build SLA+W
    _slaveAddressRWBit = TW_WRITE | (address << 1);

    // Do not send start condition again if it has been sent
    if(_repeatedStart) {
        // Reset flag
        _repeatedStart = false;
        // Send the SLA+W
        timeoutDelay = Application::CommunicationTimeoutCycles;
        do {
            TWDR = _slaveAddressRWBit;
            Application::yield();
            if(--timeoutDelay == 0) {
                sendStopCondition();
                return TWIResultCode::TimeOut;
            }
        } while( TWCR & _BV(TWWC) );
        // Enable interrupt
        TWCR = _BV(TWEN) | _BV(TWEA)  | _BV(TWIE) | _BV(TWINT);
    }
    // Enable interrupt and send start condition
    else {
       TWCR = _BV(TWEN) | _BV(TWEA)  | _BV(TWIE) | _BV(TWINT) | _BV(TWSTA);
    }

    // Wait for completion
    timeoutDelay = Application::CommunicationTimeoutCycles;

    while( wait && (_state == TWIState::MTx) ) {
        Application::yield();
        if(--timeoutDelay == 0) {
            sendStopCondition();
            return TWIResultCode::TimeOut;
        }
    }

    // Done
    return _resultCode;
}


uint8_t TWIDriver::readFrom(uint8_t address, uint8_t* data, uint8_t len, bool sendStop)
{
    // Check if there will be enough space in the buffer
    if(_msBuff.capacity() < len) return 0;

    // Wait until ready
    uint32_t timeoutDelay = Application::CommunicationTimeoutCycles;

    while(_state != TWIState::Ready) {
        Application::yield();
        if(--timeoutDelay == 0) {
            _resultCode = TWIResultCode::TimeOut;
            return 0;
        }
    }

    // Make it a master receiver
    _state             = TWIState::MRx;
    _sendStop          = sendStop;
    _expectedReadBytes = len;

    _resultCode        = TWIResultCode::NoError;

    // Clear buffer
    _msBuff.clear();

    // Build SLA+R
    _slaveAddressRWBit = TW_READ | (address << 1);

    // Do not send start condition again if it has been sent
    if(_repeatedStart) {
        // Reset flag
        _repeatedStart = false;
        // Send the SLA+R
        timeoutDelay = Application::CommunicationTimeoutCycles;
        do {
            TWDR = _slaveAddressRWBit;
            Application::yield();
            if(--timeoutDelay == 0) {
                sendStopCondition();
                _resultCode = TWIResultCode::TimeOut;
                return 0;
            }
        } while( TWCR & _BV(TWWC) );
        // Enable interrupt
        TWCR = _BV(TWEN) | _BV(TWEA)  | _BV(TWIE) | _BV(TWINT);
    }
    // Enable interrupt and send start condition
    else {
       TWCR = _BV(TWEN) | _BV(TWEA)  | _BV(TWIE) | _BV(TWINT) | _BV(TWSTA);
    }

    // Wait for completion
    timeoutDelay = Application::CommunicationTimeoutCycles;

    while(_state == TWIState::MRx) {
        Application::yield();
        if(--timeoutDelay == 0) {
            sendStopCondition();
           _resultCode = TWIResultCode::TimeOut;
            return 0;
        }
    }

    // Copy data from buffer
    const uint8_t got = _msBuff.size();
    for(uint8_t i = 0; i < got; ++i) _msBuff.get(*data++);

    // Done
    return got;
}


TWIDriver::TWIResultCode TWIDriver::transmitToMaster(const uint8_t* data, uint8_t len)
{
    // Check if it is in a slave transmitter mode
    if(_state != TWIState::STx) return TWIResultCode::NotSlaveTx;

    // Check if there is enough space in the buffer
    if(_txBuff.capacity() < _txBuff.size() + len) return TWIResultCode::DataTooLong;

    // Copy data to buffer
    for(uint8_t i = 0; i < len; ++i) _txBuff.put(*data++);

    // Done
    return TWIResultCode::NoError;
}


void TWIDriver::ack(bool ack)
{
    if(ack) TWCR = _BV(TWEN) | _BV(TWIE) | _BV(TWINT) | _BV(TWEA); // Send ACK
    else    TWCR = _BV(TWEN) | _BV(TWIE) | _BV(TWINT);             // Send NACK
}


void TWIDriver::sendStopCondition()
{
    // Send stop condition
    TWCR = _BV(TWEN) | _BV(TWEA) | _BV(TWIE) | _BV(TWINT) | _BV(TWSTO);

    // No need to wait if it is already timeout
    if(_resultCode == TWIResultCode::TimeOut) return;

    // Wait for stop condition completion
    uint32_t timeoutDelay = Application::CommunicationTimeoutCycles;

    while( TWCR & _BV(TWSTO) ) {
        Application::yield();
        if(--timeoutDelay == 0) {
            _resultCode = TWIResultCode::TimeOut;
            return;
        }
    }

    // Set state
    _state = TWIState::Ready;
}


void TWIDriver::releaseBusControl()
{
    // Release bus
    TWCR = _BV(TWEN) | _BV(TWEA) | _BV(TWIE) | _BV(TWINT);

    // Set state
    _state = TWIState::Ready;
}


void TWIDriver::__twi_isr_handler()
{
    // For convenience
    using TWIState = TWIDriver::TWIState;

    TWIDriver& twiDriver = TWIDriver::instance();

    // TWI state machine
    uint8_t c;

    switch(TW_STATUS) {

        // ### ### ### Both master and slave ### ### ###

        // No state information available
        case TW_NO_INFO:
            break;

        // Illegal start or stop condition
        case TW_BUS_ERROR:
            twiDriver._resultCode = TWIResultCode::OtherError;
            twiDriver.sendStopCondition();
            break;

        // ### ### ### Master transmitter and receiver ### ### ###

        // Start condition transmitted
        case TW_START:
            /* FALLTHROUGH */

        // Repeated start condition transmitted
        case TW_REP_START:
            // Send slave address and ACK
            TWDR = twiDriver._slaveAddressRWBit;
            twiDriver.ack(true);
            break;

        // Arbitration lost in SLA+R/W or data
        case TW_MR_ARB_LOST: /* Also handle TW_MT_ARB_LOST (it is the same value) */
            twiDriver._resultCode = TWIResultCode::OtherError;
            twiDriver.releaseBusControl();
            break;

        // ### ### ### Master transmitter ### ### ###

        // SLA+W transmitted, ACK received
        case TW_MT_SLA_ACK:
            /* FALLTHROUGH */

        // Data transmitted, ACK received
        case TW_MT_DATA_ACK:
            // If there are more byte(s), sent one byte and ACK
            if( !twiDriver._msBuff.empty() ) {
                twiDriver._msBuff.get(c);
                TWDR = c;
                twiDriver.ack(true);
            }
            // Otherwise, stop
            else{
                if(twiDriver._sendStop) {
                    twiDriver.sendStopCondition();
                }
                else {
                    twiDriver._repeatedStart = true;
                    twiDriver._state         = TWIState::Ready;
                    TWCR = _BV(TWEN) | _BV(TWINT) | _BV(TWSTA);
                }
            }
            break;

        // SLA+W transmitted, NACK received
        case TW_MT_SLA_NACK:
            twiDriver._resultCode = TWIResultCode::AddressNACK;
            twiDriver.sendStopCondition();
            break;

        // Data transmitted, NACK received
        case TW_MT_DATA_NACK:
            twiDriver._resultCode = TWIResultCode::DataNACK;
            twiDriver.sendStopCondition();
            break;

        // ### ### ### Master receiver ### ### ###

        // SLA+R transmitted, ACK received
        case TW_MR_SLA_ACK:
            // ACK there is any byte expected, NACK otherwise
            if(twiDriver._expectedReadBytes) twiDriver.ack(true );
            else                             twiDriver.ack(false);
            break;

        // Data received, ACK returned
        case TW_MR_DATA_ACK:
            // Put one byte into buffer
            c = TWDR;
            twiDriver._msBuff.put(c);
            // ACK if more bytes are expected, NACK otherwise
            if(twiDriver._msBuff.size() < twiDriver._expectedReadBytes - 1) twiDriver.ack(true );
            else                                                            twiDriver.ack(false);
            break;

        // SLA+R transmitted, NACK received
        case TW_MR_SLA_NACK:
            twiDriver.sendStopCondition();
            break;

        // Data received, NACK returned
        case TW_MR_DATA_NACK:
            // Put final byte into buffer
            c = TWDR;
            twiDriver._msBuff.put(c);
            // Done
            if(twiDriver._sendStop) {
                twiDriver.sendStopCondition();
            }
            else {
                twiDriver._repeatedStart = true;
                twiDriver._state         = TWIState::Ready;
                TWCR = _BV(TWEN) | _BV(TWINT) | _BV(TWSTA);
            }
            break;

        // ### ### ### Slave transmitter ### ### ###

        // SLA+R received, ACK returned
        case TW_ST_SLA_ACK:
            /* FALLTHROUGH */

        // Arbitration lost in SLA+RW, SLA+R received, ACK returned
        case TW_ST_ARB_LOST_SLA_ACK:
            // Enter slave transmitter mode
            twiDriver._state = TWIState::STx;
            // Clear the buffer
            twiDriver._txBuff.clear();
            // Call user callback from which the user is expected to call transmitToMaster()
            if(twiDriver._txucbSlave) (*twiDriver._txucbSlave)();
            // If the user does not put any data, NACK and break, otherwise fall through
            if( twiDriver._txBuff.empty() ) {
                twiDriver.ack(false);
                twiDriver._state = TWIState::Ready;
                break;
            }
            /* FALLTHROUGH */

        // Data transmitted, ACK received
        case TW_ST_DATA_ACK:
            // Send one byte
            twiDriver._txBuff.get(c);
            TWDR = c;
            // ACK if there is more byte(s), NACK otherwise
            if( !twiDriver._txBuff.empty() ) twiDriver.ack(true );
            else                             twiDriver.ack(false);
            break;

        // Data transmitted, NACK received
        case TW_ST_DATA_NACK:
            /* FALLTHROUGH */

        // Last data byte transmitted, ACK received
        case TW_ST_LAST_DATA:
            twiDriver.ack(true);
            twiDriver._state = TWIState::Ready;
            break;

        // ### ### ### Slave receiver ### ### ###

        // SLA+W received, ACK returned
        case TW_SR_SLA_ACK:
            /* FALLTHROUGH */

        // General call received, ACK returned
        case TW_SR_GCALL_ACK:
            /* FALLTHROUGH */

        // Arbitration lost in SLA+RW, SLA+W received, ACK returned
        case TW_SR_ARB_LOST_SLA_ACK:
            /* FALLTHROUGH */

        // Arbitration lost in SLA+RW, general call received, ACK returned
        case TW_SR_ARB_LOST_GCALL_ACK:
            // Enter slave receiver mode
            twiDriver._state = TWIState::SRx;
            // Clear the Rx buffer and ACK
            twiDriver._rxBuff.clear();
            twiDriver.ack(true);
            break;

        // Data received, ACK returned
        case TW_SR_DATA_ACK:
            /* FALLTHROUGH */

        // General call data received, ACK returned
        case TW_SR_GCALL_DATA_ACK:
            // ACK if there is still available space, ...
            if( !twiDriver._rxBuff.full() ){
                // Put one byte into buffer and ACK
                c = TWDR;
                twiDriver._rxBuff.put(c);
                twiDriver.ack(true);
            }
            // ... NACK otherwise
            else{
                twiDriver.ack(false);
            }
            break;

        // Stop or repeated start condition received while selected
        case TW_SR_STOP:
            // Release bus control
            twiDriver.releaseBusControl();
            // Call user callback to send the byte(s) to the user
            if(twiDriver._rxucbSlave) (*twiDriver._rxucbSlave)( twiDriver._rxBuff.data(), twiDriver._rxBuff.size() );
            twiDriver._rxBuff.clear();
            break;

        // Data received, NACK returned
        case TW_SR_DATA_NACK:
            /* FALLTHROUGH */

        // General call data received, NACK returned
        case TW_SR_GCALL_DATA_NACK:
            twiDriver.ack(false);
            break;

    } // switch TW_STATUS
}


} // namespace eCxx


//
// TWI interrupt vector
//
ISR(TWI_vect)
{
    if(!eCxx::twi_isr_handler) return;

    (*eCxx::twi_isr_handler)();
}

