/*
 * Connection.cpp
 *
 * Copyright (C) 2008 Matthias Schiffer <matthias@gamezock.de>
 *
 * This program is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the
 * Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program. If not, see <http://www.gnu.org/licenses/>.
 */

#include "Connection.h"
#include "FdManager.h"
#include "IPAddress.h"
#include "ThreadManager.h"

#include <cstring>
#include <sys/socket.h>

#include <boost/bind.hpp>

namespace Mad {
namespace Net {


Connection::StaticInit Connection::staticInit;


Connection::~Connection() {
  if(_isConnected())
    doDisconnect();

  if(transR.data)
    delete [] transR.data;

  while(!_sendQueueEmpty()) {
    delete [] transS.front().data;
    transS.pop();
  }

  gnutls_certificate_free_credentials(x509_cred);

  if(peer)
    delete peer;
}

void Connection::handshake() {
  boost::unique_lock<boost::shared_mutex> lock(stateLock);
  if(state != CONNECT)
    return;

  state = HANDSHAKE;
  lock.unlock();

  doHandshake();
}

void Connection::bye() {
  boost::unique_lock<boost::shared_mutex> lock(stateLock);
  if(state != DISCONNECT)
    return;

  state = BYE;
  lock.unlock();

  doBye();
}

void Connection::doHandshake() {
  boost::shared_lock<boost::shared_mutex> lock(stateLock);
  if(state != HANDSHAKE)
    return;

  int ret = gnutls_handshake(session);
  if(ret < 0) {
    lock.unlock();

    if(ret == GNUTLS_E_INTERRUPTED || ret == GNUTLS_E_AGAIN) {
      updateEvents();
      return;
    }

    // TODO: Error
    doDisconnect();
    return;
  }

  state = CONNECTION_HEADER;
  lock.unlock();

  connectionHeader();
}

void Connection::doBye() {
  if(state != BYE)
    return;

  int ret = gnutls_bye(session, GNUTLS_SHUT_RDWR);
  if(ret < 0) {
    if(ret == GNUTLS_E_INTERRUPTED || ret == GNUTLS_E_AGAIN) {
      updateEvents();
      return;
    }

    // TODO: Error
    doDisconnect();
    return;
  }

  doDisconnect();
}

void Connection::enterReceiveLoop() {
  boost::unique_lock<boost::shared_mutex> lock(stateLock);

  if(!_isConnected() || _isDisconnecting())
    return;

  if(_isConnecting())
    ThreadManager::get()->pushWork(boost::bind((void (boost::signal0<void>::*)())&boost::signal0<void>::operator(), &connectedSignal));

  state = PACKET_HEADER;
  lock.unlock();

  rawReceive(sizeof(Packet::Data), boost::bind(&Connection::packetHeaderReceiveHandler, this, _1, _2));
}

void Connection::packetHeaderReceiveHandler(const void *data, unsigned long length) {
  if(state != PACKET_HEADER)
      return;

  if(length != sizeof(Packet::Data)) {
    // TODO: Error
    doDisconnect();
    return;
  }

  header = *(const Packet::Data*)data;

  if(header.length == 0) {
    ThreadManager::get()->pushWork(boost::bind((void (boost::signal1<void, const Packet&>::*)(const Packet&))&boost::signal1<void, const Packet&>::operator(), &receiveSignal, Packet(ntohs(header.requestId))));

    enterReceiveLoop();
  }
  else {
    state = PACKET_DATA;
    rawReceive(ntohs(header.length), boost::bind(&Connection::packetDataReceiveHandler, this, _1, _2));
  }
}

void Connection::packetDataReceiveHandler(const void *data, unsigned long length) {
  if(state != PACKET_DATA)
    return;

  if(length != ntohs(header.length)) {
    // TODO: Error
    doDisconnect();
    return;
  }

  ThreadManager::get()->pushWork(boost::bind((void (boost::signal1<void, const Packet&>::*)(const Packet&))&boost::signal1<void, const Packet&>::operator(), &receiveSignal, Packet(ntohs(header.requestId), data, length)));

  enterReceiveLoop();
}

void Connection::doReceive() {
  if(!isConnected())
    return;

  boost::unique_lock<boost::mutex> lock(receiveLock);

  if(_receiveComplete())
    return;

  ssize_t ret = gnutls_record_recv(session, transR.data+transR.transmitted, transR.length-transR.transmitted);

  if(ret < 0) {
    lock.unlock();
    if(ret == GNUTLS_E_INTERRUPTED || ret == GNUTLS_E_AGAIN)
      return;

    // TODO: Error
    doDisconnect();
    return;
  }

  transR.transmitted += ret;

  if(_receiveComplete()) {
    // Save data pointer, as transR.notify might start a new reception
    uint8_t *data = transR.data;
    transR.data = 0;

    lock.unlock();

    transR.notify(data, transR.length);

    delete [] data;
  }
  else {
    lock.unlock();
  }

  updateEvents();
}

bool Connection::rawReceive(unsigned long length,
    const boost::function2<void,const void*,unsigned long> &notify)
{
  if(!isConnected())
    return false;

  boost::unique_lock<boost::mutex> lock(receiveLock);
  if(!_receiveComplete())
    return false;

  transR.data = new uint8_t[length];
  transR.length = length;
  transR.transmitted = 0;
  transR.notify = notify;

  lock.unlock();

  updateEvents();

  return true;
}

void Connection::doSend() {
  if(!isConnected())
    return;

  boost::unique_lock<boost::mutex> lock(sendLock);
  while(!_sendQueueEmpty()) {
    ssize_t ret = gnutls_record_send(session, transS.front().data+transS.front().transmitted,
        transS.front().length-transS.front().transmitted);

    if(ret < 0) {
      lock.unlock();

      if(ret == GNUTLS_E_INTERRUPTED || ret == GNUTLS_E_AGAIN)
        return;

      // TODO: Error
      doDisconnect();
      return;
    }

    transS.front().transmitted += ret;

    if(transS.front().transmitted == transS.front().length) {
      delete [] transS.front().data;
      transS.pop();
    }
  }

  lock.unlock();

  updateEvents();
}

bool Connection::rawSend(const uint8_t *data, unsigned long length) {
  if(!isConnected())
    return false;

  Transmission trans = {length, 0, new uint8_t[length], boost::function2<void,const void*,unsigned long>()};
  std::memcpy(trans.data, data, length);

  sendLock.lock();
  transS.push(trans);
  sendLock.unlock();

  updateEvents();

  return true;
}

void Connection::sendReceive(short events) {
  if(events & POLLHUP || events & POLLERR) {
    doDisconnect();
    return;
  }

  switch(state) {
    case CONNECT:
      handshake();
      return;
    case HANDSHAKE:
      doHandshake();
      return;
    case DISCONNECT:
      if(!_sendQueueEmpty())
        break;

      bye();
      return;
    case BYE:
      doBye();
      return;
    default:
      break;
  }

  if(events & POLLIN)
    doReceive();

  if(events & POLLOUT)
    doSend();
}

bool Connection::send(const Packet &packet) {
  stateLock.lock_shared();
  bool err = (!_isConnected() || _isConnecting() || _isDisconnecting());
  stateLock.unlock_shared();

  if(err)
    return false;

  return rawSend((const uint8_t*)packet.getRawData(), packet.getRawDataLength());
}

void Connection::disconnect() {
  boost::unique_lock<boost::shared_mutex> lock(stateLock);
  if(!_isConnected() || _isDisconnecting())
    return;

  state = DISCONNECT;

  lock.unlock();

  updateEvents();
}

void Connection::doDisconnect() {
  boost::unique_lock<boost::shared_mutex> lock(stateLock);

  if(_isConnected()) {
    FdManager::get()->unregisterFd(sock);

    shutdown(sock, SHUT_RDWR);
    close(sock);

    gnutls_deinit(session);

    ThreadManager::get()->pushWork(boost::bind((void (boost::signal0<void>::*)())&boost::signal0<void>::operator(), &disconnectedSignal));

    state = DISCONNECTED;
  }
}

void Connection::updateEvents() {
  receiveLock.lock();
  short events = (_receiveComplete() ? 0 : POLLIN);
  receiveLock.unlock();

  sendLock.lock();
  events |= (_sendQueueEmpty() ? 0 : POLLOUT);
  sendLock.unlock();

  stateLock.lock_shared();
  if(state == HANDSHAKE || state == BYE)
    events = ((gnutls_record_get_direction(session) == 0) ? POLLIN : POLLOUT);
  else if(state == CONNECT || state == DISCONNECT)
    events |= POLLOUT;

  FdManager::get()->setFdEvents(sock, events);
  stateLock.unlock_shared();
}

}
}