/**********************************************************************************************************************************
  WiFiSocket.cpp - Library for Arduino Wifi shield.

  Based on and modified from

  1) WiFi101 Library                 https://github.com/arduino-libraries/WiFi101

  Forked then modified by Khoi Hoang https://github.com/khoih-prog/WiFi101

  Copyright (c) 2011-2014 Arduino.  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

  Version: 1.0.0

  Version Modified By   Date      Comments
  ------- -----------  ---------- -----------
  1.0.0   K Hoang      14/11/2022 Fix severe limitation to permit sending much larger data than total 4K
 ***********************************************************************************************************************************/

extern "C"
{
#include "driver/include/m2m_wifi.h"
#include "socket/include/m2m_socket_host_if.h"
#include "driver/source/m2m_hif.h"
#include "driver/include/m2m_periph.h"
}

#include "WiFiSocket.h"

////////////////////////////////////////

#ifdef LIMITED_RAM_DEVICE
  #define SOCKET_BUFFER_SIZE 64
#else
  #define SOCKET_BUFFER_SIZE 1472
#endif

////////////////////////////////////////

extern uint8 hif_receive_blocked;

////////////////////////////////////////

enum
{
  SOCKET_STATE_INVALID,
  SOCKET_STATE_IDLE,
  SOCKET_STATE_CONNECTING,
  SOCKET_STATE_CONNECTED,
  SOCKET_STATE_BINDING,
  SOCKET_STATE_BOUND,
  SOCKET_STATE_LISTEN,
  SOCKET_STATE_LISTENING,
  SOCKET_STATE_ACCEPTED
};

////////////////////////////////////////

WiFiSocketClass::WiFiSocketClass()
{
  for (int i = 0; i < MAX_SOCKET; i++)
  {
    _info[i].state = SOCKET_STATE_INVALID;
    _info[i].parent = -1;
    _info[i].recvMsg.s16BufferSize = 0;
    _info[i].buffer.data = NULL;
    _info[i].buffer.head = NULL;
    _info[i].buffer.length = 0;
    memset(&_info[i]._lastSendtoAddr, 0x00, sizeof(_info[i]._lastSendtoAddr));
  }
}

////////////////////////////////////////

WiFiSocketClass::~WiFiSocketClass()
{
}

////////////////////////////////////////

SOCKET WiFiSocketClass::create(uint16 u16Domain, uint8 u8Type, uint8 u8Flags)
{
  SOCKET sock = socket(u16Domain, u8Type, u8Flags);

  if (sock >= 0)
  {
    _info[sock].state = SOCKET_STATE_IDLE;
    _info[sock].parent = -1;
  }

  return sock;
}

////////////////////////////////////////

sint8 WiFiSocketClass::bind(SOCKET sock, struct sockaddr *pstrAddr, uint8 u8AddrLen)
{
  if (::bind(sock, pstrAddr, u8AddrLen) < 0)
  {
    return 0;
  }

  _info[sock].state = SOCKET_STATE_BINDING;

  unsigned long start = millis();

  while (_info[sock].state == SOCKET_STATE_BINDING && millis() - start < 2000)
  {
    m2m_wifi_handle_events(NULL);
  }

  if (_info[sock].state != SOCKET_STATE_BOUND)
  {
    _info[sock].state = SOCKET_STATE_IDLE;
    return 0;
  }

  _info[sock].recvMsg.s16BufferSize = 0;

  if (sock < TCP_SOCK_MAX)
  {
    // TCP
  }
  else
  {
    // UDP
    recvfrom(sock, NULL, 0, 0);
  }

  return 1;
}

////////////////////////////////////////

sint8 WiFiSocketClass::listen(SOCKET sock, uint8 backlog)
{
  if (::listen(sock, backlog) < 0)
  {
    return 0;
  }

  _info[sock].state = SOCKET_STATE_LISTEN;

  unsigned long start = millis();

  while (_info[sock].state == SOCKET_STATE_LISTEN && millis() - start < 2000)
  {
    m2m_wifi_handle_events(NULL);
  }

  if (_info[sock].state != SOCKET_STATE_LISTENING)
  {
    _info[sock].state = SOCKET_STATE_IDLE;
    return 0;
  }

  return 1;
}

////////////////////////////////////////

sint8 WiFiSocketClass::setopt(SOCKET socket, uint8 u8Level, uint8 option_name, const void *option_value,
                              uint16 u16OptionLen)
{
  return setsockopt(socket, u8Level, option_name, option_value, u16OptionLen);
}

////////////////////////////////////////

sint8 WiFiSocketClass::connect(SOCKET sock, struct sockaddr *pstrAddr, uint8 u8AddrLen)
{
  if (::connect(sock, pstrAddr, u8AddrLen) < 0)
  {
    return 0;
  }

  _info[sock].state = SOCKET_STATE_CONNECTING;

  unsigned long start = millis();

  while (_info[sock].state == SOCKET_STATE_CONNECTING && millis() - start < 20000)
  {
    m2m_wifi_handle_events(NULL);
  }

  if (_info[sock].state != SOCKET_STATE_CONNECTED)
  {
    _info[sock].state = SOCKET_STATE_IDLE;
    return 0;
  }

  _info[sock].recvMsg.s16BufferSize = 0;
  _info[sock].recvMsg.strRemoteAddr.sin_port = ((struct sockaddr_in*)pstrAddr)->sin_port;
  _info[sock].recvMsg.strRemoteAddr.sin_addr.s_addr = ((struct sockaddr_in*)pstrAddr)->sin_addr.s_addr;
  recv(sock, NULL, 0, 0);

  return 1;
}

////////////////////////////////////////

uint8 WiFiSocketClass::connected(SOCKET sock)
{
  m2m_wifi_handle_events(NULL);

  return (_info[sock].state == SOCKET_STATE_CONNECTED);
}

////////////////////////////////////////

uint8 WiFiSocketClass::listening(SOCKET sock)
{
  m2m_wifi_handle_events(NULL);

  return (_info[sock].state == SOCKET_STATE_LISTENING);
}

////////////////////////////////////////

uint8 WiFiSocketClass::bound(SOCKET sock)
{
  m2m_wifi_handle_events(NULL);

  return (_info[sock].state == SOCKET_STATE_BOUND);
}

////////////////////////////////////////

int WiFiSocketClass::available(SOCKET sock)
{
  m2m_wifi_handle_events(NULL);

  if (_info[sock].state != SOCKET_STATE_CONNECTED && _info[sock].state != SOCKET_STATE_BOUND)
  {
    return 0;
  }

  return (_info[sock].buffer.length + _info[sock].recvMsg.s16BufferSize);
}

////////////////////////////////////////

int WiFiSocketClass::peek(SOCKET sock)
{
  m2m_wifi_handle_events(NULL);

  if (_info[sock].state != SOCKET_STATE_CONNECTED && _info[sock].state != SOCKET_STATE_BOUND)
  {
    return -1;
  }

  if (available(sock) == 0)
  {
    return -1;
  }

  if (_info[sock].buffer.length == 0 && _info[sock].recvMsg.s16BufferSize)
  {
    if (!fillRecvBuffer(sock))
    {
      return -1;
    }
  }

  return *_info[sock].buffer.head;
}

////////////////////////////////////////

int WiFiSocketClass::read(SOCKET sock, uint8_t* buf, size_t size)
{
  m2m_wifi_handle_events(NULL);

  if (_info[sock].state != SOCKET_STATE_CONNECTED && _info[sock].state != SOCKET_STATE_BOUND)
  {
    return 0;
  }

  int avail = available(sock);

  if (avail <= 0)
  {
    return 0;
  }

  if ((int)size > avail)
  {
    size = avail;
  }

  int bytesRead = 0;

  while (size)
  {
    if (_info[sock].buffer.length == 0 && _info[sock].recvMsg.s16BufferSize)
    {
      if (!fillRecvBuffer(sock))
      {
        break;
      }
    }

    int toCopy = size;

    if (toCopy > _info[sock].buffer.length)
    {
      toCopy = _info[sock].buffer.length;
    }

    memcpy(buf, _info[sock].buffer.head, toCopy);
    _info[sock].buffer.head += toCopy;
    _info[sock].buffer.length -= toCopy;

    buf += toCopy;
    size -= toCopy;
    bytesRead += toCopy;
  }

  if (_info[sock].buffer.length == 0 && _info[sock].recvMsg.s16BufferSize == 0)
  {
    if (sock < TCP_SOCK_MAX)
    {
      // TCP
      recv(sock, NULL, 0, 0);
    }
    else
    {
      // UDP
      recvfrom(sock, NULL, 0, 0);
    }

    m2m_wifi_handle_events(NULL);
  }

  return bytesRead;
}

////////////////////////////////////////

IPAddress WiFiSocketClass::remoteIP(SOCKET sock)
{
  return _info[sock].recvMsg.strRemoteAddr.sin_addr.s_addr;
}

////////////////////////////////////////

uint16_t WiFiSocketClass::remotePort(SOCKET sock)
{
  return _info[sock].recvMsg.strRemoteAddr.sin_port;
}

////////////////////////////////////////

size_t WiFiSocketClass::write(SOCKET sock, const uint8_t *buf, size_t size)
{
  m2m_wifi_handle_events(NULL);

  if (_info[sock].state != SOCKET_STATE_CONNECTED)
  {
    return 0;
  }

#ifdef CONF_PERIPH
  // Network led ON (rev A then rev B).
  m2m_periph_gpio_set_val(M2M_PERIPH_GPIO16, 0);
  m2m_periph_gpio_set_val(M2M_PERIPH_GPIO5, 0);
#endif

  sint16 err;

  while ((err = send(sock, (void *)buf, size, 0)) < 0)
  {
    // Exit on fatal error, retry if buffer not ready.
    if (err != SOCK_ERR_BUFFER_FULL)
    {
      size = 0;
      break;
    }
    else if (hif_receive_blocked)
    {
      size = 0;
      break;
    }

    m2m_wifi_handle_events(NULL);
  }

#ifdef CONF_PERIPH
  // Network led OFF (rev A then rev B).
  m2m_periph_gpio_set_val(M2M_PERIPH_GPIO16, 1);
  m2m_periph_gpio_set_val(M2M_PERIPH_GPIO5, 1);
#endif

  return size;
}

////////////////////////////////////////

sint16 WiFiSocketClass::sendto(SOCKET sock, void *pvSendBuffer, uint16 u16SendLength, uint16 flags,
                               struct sockaddr *pstrDestAddr, uint8 u8AddrLen)
{
  m2m_wifi_handle_events(NULL);

  if (_info[sock].state != SOCKET_STATE_BOUND)
  {
    return -1;
  }

  if (memcmp(&_info[sock]._lastSendtoAddr, pstrDestAddr, sizeof(_info[sock]._lastSendtoAddr)) != 0)
  {
    memcpy(&_info[sock]._lastSendtoAddr, pstrDestAddr, sizeof(_info[sock]._lastSendtoAddr));

    return ::sendto(sock, pvSendBuffer, u16SendLength, flags, pstrDestAddr, u8AddrLen);
  }
  else
  {
    return ::send(sock, pvSendBuffer, u16SendLength, 0);
  }
}

////////////////////////////////////////

sint8 WiFiSocketClass::close(SOCKET sock)
{
  m2m_wifi_handle_events(NULL);

  if (_info[sock].state == SOCKET_STATE_CONNECTED || _info[sock].state == SOCKET_STATE_BOUND)
  {
    if (_info[sock].recvMsg.s16BufferSize > 0)
    {
      _info[sock].recvMsg.s16BufferSize = 0;

      // flush any data not processed
      hif_receive(0, NULL, 0, 1);
    }
  }

  _info[sock].state = SOCKET_STATE_INVALID;
  _info[sock].parent = -1;

  if (_info[sock].buffer.data != NULL)
  {
    free(_info[sock].buffer.data);
  }

  _info[sock].buffer.data = NULL;
  _info[sock].buffer.head = NULL;
  _info[sock].buffer.length = 0;
  _info[sock].recvMsg.s16BufferSize = 0;
  memset(&_info[sock]._lastSendtoAddr, 0x00, sizeof(_info[sock]._lastSendtoAddr));

  return ::close(sock);
}

////////////////////////////////////////

int WiFiSocketClass::hasParent(SOCKET sock, SOCKET child)
{
  if (_info[child].parent != sock)
  {
    return 0;
  }

  return 1;
}

////////////////////////////////////////

SOCKET WiFiSocketClass::accepted(SOCKET sock)
{
  m2m_wifi_handle_events(NULL);

  for (SOCKET s = 0; s < TCP_SOCK_MAX; s++)
  {
    if (_info[s].parent == sock && _info[s].state == SOCKET_STATE_ACCEPTED)
    {
      _info[s].state = SOCKET_STATE_CONNECTED;

      _info[s].recvMsg.s16BufferSize = 0;
      recv(s, NULL, 0, 0);

      return s;
    }
  }

  return -1;
}

////////////////////////////////////////

void WiFiSocketClass::eventCallback(SOCKET sock, uint8 u8Msg, void *pvMsg)
{
  WiFiSocket.handleEvent(sock, u8Msg, pvMsg);
}

////////////////////////////////////////

void WiFiSocketClass::handleEvent(SOCKET sock, uint8 u8Msg, void *pvMsg)
{
  switch (u8Msg)
  {
    /* Socket bind. */
    case SOCKET_MSG_BIND:
    {
      tstrSocketBindMsg *pstrBind = (tstrSocketBindMsg *)pvMsg;

      if (pstrBind && pstrBind->status == 0)
      {
        _info[sock].state = SOCKET_STATE_BOUND;
      }
      else
      {
        _info[sock].state = SOCKET_STATE_IDLE;
      }
    }
    break;

    /* Socket listen. */
    case SOCKET_MSG_LISTEN:
    {
      tstrSocketListenMsg *pstrListen = (tstrSocketListenMsg *)pvMsg;

      if (pstrListen && pstrListen->status == 0)
      {
        _info[sock].state = SOCKET_STATE_LISTENING;
      }
      else
      {
        _info[sock].state = SOCKET_STATE_IDLE;
      }
    }
    break;

    /* Socket accept. */
    case SOCKET_MSG_ACCEPT:
    {
      tstrSocketAcceptMsg *pstrAccept = (tstrSocketAcceptMsg*)pvMsg;

      if (pstrAccept && pstrAccept->sock > -1)
      {
        _info[pstrAccept->sock].state = SOCKET_STATE_ACCEPTED;
        _info[pstrAccept->sock].parent = sock;
        _info[pstrAccept->sock].recvMsg.strRemoteAddr = pstrAccept->strAddr;
      }
    }
    break;

    /* Socket connected. */
    case SOCKET_MSG_CONNECT:
    {
      tstrSocketConnectMsg *pstrConnect = (tstrSocketConnectMsg *)pvMsg;

      if (pstrConnect && pstrConnect->s8Error >= 0)
      {
        _info[sock].state = SOCKET_STATE_CONNECTED;
      }
      else
      {
        _info[sock].state = SOCKET_STATE_IDLE;
      }
    }
    break;

    /* Socket data received. */
    case SOCKET_MSG_RECV:
    case SOCKET_MSG_RECVFROM:
    {
      tstrSocketRecvMsg *pstrRecvMsg = (tstrSocketRecvMsg *)pvMsg;

#ifdef CONF_PERIPH
      // Network led ON (rev A then rev B).
      m2m_periph_gpio_set_val(M2M_PERIPH_GPIO16, 0);
      m2m_periph_gpio_set_val(M2M_PERIPH_GPIO5, 0);
#endif

      if (pstrRecvMsg->s16BufferSize <= 0)
      {
        close(sock);
      }
      else if (_info[sock].state == SOCKET_STATE_CONNECTED || _info[sock].state == SOCKET_STATE_BOUND)
      {
        _info[sock].recvMsg.pu8Buffer = pstrRecvMsg->pu8Buffer;
        _info[sock].recvMsg.s16BufferSize = pstrRecvMsg->s16BufferSize;

        if (sock < TCP_SOCK_MAX)
        {
          // TCP
        }
        else
        {
          // UDP
          _info[sock].recvMsg.strRemoteAddr = pstrRecvMsg->strRemoteAddr;
        }

        fillRecvBuffer(sock);
      }
      else
      {
        // not connected or bound, discard data
        hif_receive(0, NULL, 0, 1);
      }

#ifdef CONF_PERIPH
      // Network led OFF (rev A then rev B).
      m2m_periph_gpio_set_val(M2M_PERIPH_GPIO16, 1);
      m2m_periph_gpio_set_val(M2M_PERIPH_GPIO5, 1);
#endif
    }
    break;

    /* Socket data sent. */
    case SOCKET_MSG_SEND:
    {
      // sint16 *s16Sent = (sint16 *)pvMsg;
    }
    break;

    default:
      break;
  }
}

////////////////////////////////////////

int WiFiSocketClass::fillRecvBuffer(SOCKET sock)
{
  if (_info[sock].buffer.data == NULL)
  {
    _info[sock].buffer.data = (uint8_t*)malloc(SOCKET_BUFFER_SIZE);
    _info[sock].buffer.head = _info[sock].buffer.data;
    _info[sock].buffer.length = 0;
  }

  int size = _info[sock].recvMsg.s16BufferSize;

  if (size > SOCKET_BUFFER_SIZE)
  {
    size = SOCKET_BUFFER_SIZE;
  }

  uint8 lastTransfer = ((sint16)size == _info[sock].recvMsg.s16BufferSize);

  if (hif_receive(_info[sock].recvMsg.pu8Buffer, _info[sock].buffer.data, (sint16)size, lastTransfer) != M2M_SUCCESS)
  {
    return 0;
  }

  _info[sock].buffer.head = _info[sock].buffer.data;
  _info[sock].buffer.length = size;
  _info[sock].recvMsg.pu8Buffer += size;
  _info[sock].recvMsg.s16BufferSize -= size;

  return 1;
}

////////////////////////////////////////

WiFiSocketClass WiFiSocket;