#include "NudpSocket.h"

#ifdef _DEBUG
#define DEB( x ) x
#else
#define DEB( x )
#endif

NudpSocket::NudpSocket( SocketHandler& h, int ibufsz, int obufsz ) : UdpSocket( h, ibufsz ),
_state( NUDP_CLOSED ),
_last_top( 0 ),
_data_buf( NULL ),
_got_packet( false ),
_last_data_length( 0 ),
_obufsz( obufsz ),
_log( NULL ),
_err_log( NULL )
{
	_in_packet = new char[_obufsz+NUDP_HEADER_LENGTH];
	_out_packet = new char[_obufsz+NUDP_HEADER_LENGTH];
	_out_packet[0] = 0xFF;
	_out_packet[1] = 0x00;

	_data_frag = new bool[NUDP_DATA_PACKETS];
	memset( _data_frag, 0, NUDP_DATA_PACKETS );

	DEB( printf( "NudpSocket created. State is CLOSED.\n\n" ); )
}

NudpSocket::~NudpSocket( void )
{
	delete _in_packet;
	delete _out_packet;
	delete _data_frag;
	Log( "NudpSocket deleted.\n\n" );
}

int NudpSocket::Bind( ipaddr_t iface, port_t port )
{
	char p_str[6];
	sprintf( p_str, "%u", port );
	l2ip( iface, _my_address );
	_my_address += ":";
	_my_address += p_str;

	return UdpSocket::Bind( iface, port );
}

bool NudpSocket::Open( ipaddr_t host, port_t port )
{
	Log( "NudpSocket::Open():\n" );

	if ( _state != NUDP_CLOSED )
	{
		LogErr( "Socket still in use.\n\n" );
		return false;
	}
	if ( !this->UdpSocket::Open( host, port ) )
	{
		LogErr( "UdpSocket::Open() failed.\n\n" );
		return false;
	}

	_remote_host = host;
	_remote_port = port;
	_state = NUDP_READY;

	Log( "CLOSED->READY.\n\n" );
	return true;
}

bool NudpSocket::Suspend( void )
{
	Log( "NudpSocket::Suspend():\n" );

	if ( !(_state & NUDP_READY) )
	{
		LogErr( "Socket state is not READY.\n\n" );
		return false;
	}

	_state = NUDP_CLOSED;
	Log( "READY->CLOSED.\n\n" );
	return true;
}

char NudpSocket::ChkSum( const char* buf )
{
	char chk = 0;
	for ( int i = 0; i < NUDP_HEADER_LENGTH; i++ )
	{
		if ( i == 3 ) continue;
		chk += buf[i];
	}
	return ~chk;
}

bool NudpSocket::SendCmd( const char number[4], const char* data, const size_t data_length, const char top )
{
	unsigned int n = (unsigned char)number[3];
	n = (n << 8) + (unsigned char)number[2];
	n = (n << 8) + (unsigned char)number[1];
	n = (n << 8) + (unsigned char)number[0];

	return SendCmd( n, data, data_length, top );
}

bool NudpSocket::SendCmd( unsigned int number, const char* data, const size_t data_length, const char top )
{
	Log( "NudpSocket::SendCmd():\n" );

	if ( _state & NUDP_CMD_ERR )
	{
		Log( "NUDP_CMD_ERR flag cleared.\n" );
		_state &= ~NUDP_CMD_ERR;
	}
	if ( (_state & 0x0F) != NUDP_READY )
	{
		LogErr( "Socket state != READY, packet not sent.\n\n" );
		return false;
	}
	if ( (number & 0xFF) == 0x08 )
	{
		if ( !(_state & NUDP_DATA) )
		{
			Log( "Command to send is 0x08: start RAW data dumping.\n" );
			// prepare for data taking, even if command will fail
			memset( _data_frag, 0, NUDP_DATA_PACKETS );
			_data_packet_count = 0;
			_last_raw_num = 0;
		}
		else
		{
			LogErr( "Cannot send command 0x08 while receiving RAW data.\n\n" );
			return false;
		}
	}
	_state = (_state & 0xF0) | NUDP_WAITING;

	_out_packet[2] = top & 0x0F;
	memcpy( _out_packet+4, &number, 4 );
	_out_packet[3] = ChkSum( _out_packet );

	size_t packet_length = NUDP_HEADER_LENGTH;
	if ( data && data_length )
	{
		if ( data_length > _obufsz )
		{
			LogErr( "Data field too long, packet not sent.\n\n" );
			_state = (_state & 0xF0) | NUDP_READY;
			return false;
		}
		memcpy( _out_packet+NUDP_HEADER_LENGTH, data, data_length );
		packet_length += data_length;
	}

	memcpy( _chk_header, _out_packet, NUDP_HEADER_LENGTH );
	_chk_header[2] |= 0x80;  // set ACK=1
	_chk_header[3] = ChkSum( _chk_header );

	// set NUDP_DATA flag just before sending 0x08 command
	if ( (number & 0xFF) == 0x08 ) _state |= NUDP_DATA;

	SendBuf( _out_packet, packet_length );

	char log_buf[128];
	std::string dest_ip_str;
	l2ip( _remote_host, dest_ip_str );
	sprintf( log_buf, "To: %s:%u\nNumber field: %.8X\nData length: %u\nREADY->WAITING\n\n",
		dest_ip_str.c_str(), _remote_port,
		htonl(number), packet_length-NUDP_HEADER_LENGTH );
	Log( log_buf );

	return true;
}

bool NudpSocket::GetResponse( unsigned int* number, char* data, size_t* data_length, char* top )
{
	if ( (_state & 0x0F) != NUDP_READY )
	{
		LogErr( "NudpSocket::GetResponse():\nState is not READY.\n\n" );
		return false;
	}

	if ( top ) *top = _in_packet[2] & 0x0F;
	if ( number ) *number = *((unsigned int*)_in_packet+1);
	if ( data && data_length && _last_data_length )
	{
		*data_length = _last_data_length;
		memcpy( data, _in_packet+NUDP_HEADER_LENGTH, _last_data_length );
	}

	return true;
}

bool NudpSocket::PrepareForData( char* buf_addr )
{
	if ( _state & NUDP_DATA )
	{
		LogErr( "NudpSocket::PrepareForData():\nCannot change setup while receiving RAW data.\n\n" );
		return false;
	}
	memset( _data_frag, 0, NUDP_DATA_PACKETS );
	_data_packet_count = 0;
	_last_raw_num = 0;
	_data_buf = buf_addr;

	return true;
}

bool NudpSocket::FragmentReceivedByOffset( int offset )
{
	if ( offset % NUDP_DATA_LENGTH ) return false;
	return  FragmentReceivedByIndex( offset / NUDP_DATA_LENGTH );
}

bool NudpSocket::FragmentReceivedByIndex( int index )
{
	if ( (index >= 0) && (index < NUDP_DATA_PACKETS ) )
		return _data_frag[index];
	else return false;
}

bool NudpSocket::AllDataReceived( void )
{
	if ( _data_packet_count == NUDP_DATA_PACKETS ) return true;
	else return false;
}

void NudpSocket::EmergencyDataStandBy( void )
{
	_state &= ~NUDP_DATA;
}

void NudpSocket::StopData( void )
{
	_state &= ~NUDP_DATA;
	_data_buf = NULL;

	char log_buf[64];
	sprintf( log_buf, "NudpSocket::StopData():\nNUDP state: %X.\n\n", (unsigned char)_state );
	Log( log_buf );
}

void NudpSocket::StopWaiting( void )
{
	Log( "NudpSocket::StopWaiting():\n" );
	if ( (_state & 0x0F) == NUDP_WAITING )
	{
		_state = (_state & 0xF0) | NUDP_READY;
		Log( "WAITING->READY\n\n" );
	}
	else { Log( "State is not WAITING.\n\n" ); }
}

bool NudpSocket::GotPacket( void )
{
	if ( _got_packet )
	{
		_got_packet = false;
		return true;
	}
	else return false;
}

void NudpSocket::Log( const char* msg )
{
	if ( _log )
	{
		fprintf( _log, msg );
		fflush( _log );
	}
	DEB( printf( msg ); )
}

void NudpSocket::LogErr( const char* msg )
{
	if ( _log )
	{
		fprintf( _log, msg );
		fflush( _log );
	}
	if ( _err_log )
	{
		fprintf( _err_log, msg );
		fflush( _err_log );
	}
	DEB( printf( msg ); )
}

bool NudpSocket::SetLogFile( FILE* log )
{
	if ( _log ) return false;
	_log = log;
	return true;
}

bool NudpSocket::SetErrLogFile( FILE* err_log )
{
	if ( _err_log ) return false;
	_err_log = err_log;
	return true;
}

void NudpSocket::OnRawData( const char* buf, size_t len, struct sockaddr* sender_addr, socklen_t sa_len )
{
	char log_buf[128];

	if ( _log )
	{
		sprintf( log_buf, "NudpSocket::OnRawData() (at %s):\n", _my_address.c_str() );
		Log( log_buf );
	}

	if ( sa_len != sizeof(struct sockaddr_in) ) return; // unexpected, IPv6 - do nothing

	struct sockaddr_in sa;
	memcpy( &sa, sender_addr, sa_len);
	ipaddr_t sender_ip;
	memcpy( &sender_ip, &sa.sin_addr, 4 );

	if ( _log )
	{
		std::string sender_ip_str;
		l2ip( sender_ip, sender_ip_str );
		sprintf( log_buf, "From: %s:%u\n", sender_ip_str.c_str(), ntohs(sa.sin_port) );
		Log( log_buf );
	}

	if ( (sender_ip != _remote_host) || ( ntohs(sa.sin_port) != _remote_port) )
	{
		LogErr( "Unexpected sender.\n\n" );
		return;
	}

	if ( len < NUDP_HEADER_LENGTH )
	{
		sprintf( log_buf, "UDP packet length < %u.\n\n", NUDP_HEADER_LENGTH );
		LogErr( log_buf );
		return;
	}

	if ( ((unsigned char)buf[0] != 0xFF) || ((unsigned char)buf[1] != 0x00) )
	{
		LogErr( "No NUDP signature.\n\n" );
		return;
	}

	_got_packet = true;
	_last_top = buf[2] & 0x0F;
	if ( buf[2] == 0x07 ) // 0x07: ack=0, ver=0, top=7 -> RAW data
	{
		Log( "RAW data packet.\n" );
		if ( buf[3] != ChkSum( buf ) )
		{
			LogErr( "Checksum is not correct.\n\n" );
			return;
		}

		if ( !(_state & NUDP_DATA) )
		{
			LogErr( "Socket is not expecting RAW data.\n" );
			return;
		}

		int offset = *((int*)(buf+4)); offset = offset << 1;
		int raw_num = offset / NUDP_DATA_LENGTH;
		if ( !_data_frag[raw_num] )
		{
			if ( _data_buf ) memcpy( _data_buf+offset, buf+NUDP_HEADER_LENGTH, len-NUDP_HEADER_LENGTH );
			_data_frag[raw_num] = true;
			_data_packet_count++;
			_last_raw_num = raw_num;
			if ( raw_num == NUDP_DATA_PACKETS-1 )
			{
				Log( "Last RAW data packet in this image.\n" );
				StopData();
			}
		}
		else
		{
			sprintf( log_buf, "RAW data packet #%d (offset=%X) duplicated.\n", raw_num, offset );
			LogErr( log_buf );
		}

		if ( _log )
		{
			sprintf( log_buf, "RAW data packet #%d (offset=%X, length=%d) parsed.\n\n", raw_num, offset, len-NUDP_HEADER_LENGTH );
			Log( log_buf );
		}
	}
	else
	{
		Log( "Camera response packet.\n" );
		switch ( _state & 0x0F )
		{
		case NUDP_WAITING:
			if ( (buf[2] & 0x0F) == 5 )
			{
				_chk_header[6] = buf[6];
				_chk_header[7] = buf[7];
				_chk_header[3] = ChkSum( _chk_header );
			}

			if ( memcmp( buf, _chk_header, NUDP_HEADER_LENGTH ) == 0 )
			{
				memcpy( _in_packet, buf, len );
				_last_data_length = len - NUDP_HEADER_LENGTH;
				_state = (_state & 0xF0) | NUDP_READY;
				Log( "WAITING->READY\n" );
			}
			else
			{
				// buf[3] - checksum calculated in camera
				if ( buf[3] != ChkSum( buf ) )
				{
					_state = (_state & 0xF0) | NUDP_READY | NUDP_CMD_ERR;
					LogErr( "Response damaged (bad checksum).\nWAITING->READY|CMD_ERR\n" );
				}
				else
				{
					// good checksum - but on something unexpected 
					LogErr( "Unexpected camera response.\nStill WAITING\n" );
				}
			}
			break;

		case NUDP_READY:
			LogErr( "Unexpected packet (not a problem for a driver, but should not happen).\n" );
			break;
		case NUDP_CLOSED:
			LogErr( "Socket is suspended!\n" );
			break;
		}
		Log( "Camera response packet parsed.\n\n" );
	}
}