/****************************************************************************
	 *   Copyright (C) 2006-2010 by Jason Ansel, Kapil Arya, and Gene Cooperman *
	 *   jansel@csail.mit.edu, kapil@ccs.neu.edu, gene@ccs.neu.edu              *
	 *                                                                          *
	 *   This file is part of the dmtcp/src module of DMTCP (DMTCP:dmtcp/src).  *
	 *                                                                          *
	 *  DMTCP:dmtcp/src 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 3 of the      *
	 *  License, or (at your option) any later version.                         *
	 *                                                                          *
	 *  DMTCP:dmtcp/src 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 DMTCP:dmtcp/src.  If not, see                        *
	 *  <http://www.gnu.org/licenses/>.                                         *
	 ****************************************************************************/
	
	#include "connectionrewirer.h"
	#include "dmtcpmessagetypes.h"
	#include "syscallwrappers.h"
	
	
	
	void dmtcp::ConnectionRewirer::onData ( jalib::JReaderInterface* sock )
	{
	  JASSERT ( sock->bytesRead() == sizeof ( DmtcpMessage ) ) ( sock->bytesRead() ) ( sizeof ( DmtcpMessage ) );
	  DmtcpMessage& msg = * ( DmtcpMessage* ) sock->buffer();
	  msg.assertValid();
	
	  if ( msg.type == DMT_FORCE_RESTART )
	  {
	    JTRACE ( "got DMT_FORCE_RESTART, exiting ConnectionRewirer" ) ( _pendingOutgoing.size() ) ( _pendingIncoming.size() );
	    _pendingIncoming.clear();
	    _pendingOutgoing.clear();
	    finishup();
	    return;
	  }
	
	  JASSERT ( msg.type==DMT_RESTORE_WAITING ) ( msg.type ).Text ( "unexpected message" );
	
	  // Find returns iterator 'i' w/ 0 or more elts, with first elt matching key.
	  iterator i = _pendingOutgoing.find ( msg.restorePid );
	
	  if ( i == _pendingOutgoing.end() )
	  {
	    // 'i' is an iterator over 0 elements.
	    JTRACE ( "got RESTORE_WAITING MESSAGE [not used]" ) ( msg.restorePid ) ( _pendingOutgoing.size() ) ( _pendingIncoming.size() );
	  }
	  else
	  {
	    // 'i' is an iterator over 1 element.
	    JTRACE ( "got RESTORE_WAITING MESSAGE, reconnecting..." )
	    ( msg.restorePid ) ( msg.restorePort ) ( msg.restoreAddrlen ) ( _pendingOutgoing.size() ) ( _pendingIncoming.size() );
	    const dmtcp::vector<int>& fds = i->second;
	    JASSERT ( fds.size() > 0 );
	    int fd0 = fds[0];
	
	    jalib::JSocket remote = jalib::JSocket::Create();
	    remote.changeFd ( fd0 );
	    errno = 0;
	    JASSERT ( remote.connect ( ( sockaddr* ) &msg.restoreAddr,msg.restoreAddrlen,msg.restorePort ) )
	    ( msg.restorePid ) ( msg.restorePort ) ( JASSERT_ERRNO )
	    .Text ( "failed to restore connection" );
	
	    {
	      DmtcpMessage peerMsg;
	      peerMsg.type = DMT_RESTORE_RECONNECTED;
	      peerMsg.restorePid = msg.restorePid;
	      addWrite ( new jalib::JChunkWriter ( remote, ( char* ) &peerMsg,sizeof ( peerMsg ) ) );
	    }
	
	    for ( size_t n = 1; n<fds.size(); ++n )
	    {
	      JTRACE ( "restoring extra fd" ) ( fd0 ) ( fds[n] );
	      JASSERT ( _real_dup2 ( fd0,fds[n] ) == fds[n] ) ( fd0 ) ( fds[n] ) ( msg.restorePid )
	      .Text ( "dup2() failed" );
	    }
	    _pendingOutgoing.erase ( i );
	  }
	
	  if ( pendingCount() ==0 ) finishup();
	#ifdef DEBUG
	  else debugPrint();
	#endif
	}
	
	void dmtcp::ConnectionRewirer::onConnect ( const jalib::JSocket& sock,  const struct sockaddr* /*remoteAddr*/,socklen_t /*remoteLen*/ )
	{
	  jalib::JSocket remote = sock;
	  DmtcpMessage msg;
	  msg.poison();
	  remote >> msg;
	  msg.assertValid();
	  JASSERT ( msg.type == DMT_RESTORE_RECONNECTED ) ( msg.type ).Text ( "unexpected message" );
	
	  iterator i = _pendingIncoming.find ( msg.restorePid );
	

	  JASSERT ( i != _pendingIncoming.end() ) ( msg.restorePid )
	  .Text ( "got unexpected incoming restore request" );
	

	  const dmtcp::vector<int>& fds = i->second;
	  JASSERT ( fds.size() > 0 );
	  int fd0 = fds[0];
	
	  remote.changeFd ( fd0 );
	
	  JTRACE ( "restoring incoming connection" ) ( msg.restorePid ) ( fd0 ) ( fds.size() );
	
	  for ( size_t i = 1; i<fds.size(); ++i )
	  {
	    JTRACE ( "restoring extra fd" ) ( fd0 ) ( fds[i] );
	    JASSERT ( _real_dup2 ( fd0,fds[i] ) == fds[i] ) ( fd0 ) ( fds[i] ) ( msg.restorePid )
	    .Text ( "dup2() failed" );
	  }
	
	  _pendingIncoming.erase ( i );
	
	
	  if ( pendingCount() ==0 ) finishup();
	#ifdef DEBUG
	  else debugPrint();
	#endif
	}
	
	void dmtcp::ConnectionRewirer::finishup()
	{
	  JTRACE ( "finishup begin" ) ( _listenSockets.size() ) ( _dataSockets.size() );
	  //i expect both sizes above to be 1
	  //close the restoreSocket
	  for ( size_t i=0; i<_listenSockets.size(); ++i )
	    _listenSockets[i].close();
	  //poison the coordinator socket listener
	  for ( size_t i=0; i<_dataSockets.size(); ++i )
	    _dataSockets[i]->socket() = -1;
	
	//     JTRACE("finishup end");
	}
	
	void dmtcp::ConnectionRewirer::onDisconnect ( jalib::JReaderInterface* sock )
	{
	  JASSERT ( sock->socket().sockfd() < 0 )
	  .Text ( "dmtcp_coordinator disconnected" );
	}
	
	int dmtcp::ConnectionRewirer::coordinatorFd() const
	{
	  return _coordinatorFd;
	}
	
	
	void dmtcp::ConnectionRewirer::setCoordinatorFd ( const int& theValue )
	{
	  _coordinatorFd = theValue;
	}
	
	void dmtcp::ConnectionRewirer::doReconnect()
	{
	  if ( pendingCount() > 0 ) monitorSockets();
	}
	
	void dmtcp::ConnectionRewirer::registerIncoming ( const ConnectionIdentifier& local
	        , const dmtcp::vector<int>& fds )
	{
	  _pendingIncoming[local] = fds;
	
	  JTRACE ( "announcing pending incoming" ) ( local );
	  DmtcpMessage msg;
	  msg.type = DMT_RESTORE_WAITING;
	  msg.restorePid = local;
	
	  JASSERT ( _coordinatorFd > 0 );
	  addWrite ( new jalib::JChunkWriter ( _coordinatorFd , ( char* ) &msg, sizeof ( DmtcpMessage ) ) );
	}
	
	void dmtcp::ConnectionRewirer::registerOutgoing ( const ConnectionIdentifier& remote
	        , const dmtcp::vector<int>& fds )
	{
	  _pendingOutgoing[remote] = fds;
	}
	
	void dmtcp::ConnectionRewirer::debugPrint() const
	{
	  JASSERT_STDERR << "Pending Incoming:\n";
	  for ( const_iterator i = _pendingIncoming.begin(); i!=_pendingIncoming.end(); ++i )
	  {
	    JASSERT_STDERR << i->first << " numFds=" << i->second.size() << " firstFd=" << i->second[0] << '\n';
	  }
	  JASSERT_STDERR << "Pending Outgoing:\n";
	  for ( const_iterator i = _pendingOutgoing.begin(); i!=_pendingOutgoing.end(); ++i )
	  {
	    JASSERT_STDERR << i->first << " numFds=" << i->second.size() << " firstFd=" << i->second[0] << '\n';
	  }
	}