#pragma once

#include <Windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fstream>

#pragma comment( lib, "ws2_32.lib" )

#include <vector>

/* TEST */
#include <Psapi.h>
#include <TlHelp32.h>
/* TEST */

#include "err.hpp"
#include "util.hpp"

/*
protocol rules:
every msg must be xored
first byte is the xorkey

*/

namespace client
{
	constexpr auto PORT_NUM = 6969;
	constexpr auto BUFFER_SIZE = 255;

	class c_connect {
	public:
		c_connect( const char* ip ) : 
			m_ip( inet_addr( ip ) ) { }

		~c_connect( ) {
			if( m_socket )
				closesocket( m_socket );

			WSACleanup( );
		}

		bool setup( ) {
			int code{ };

			if( WSAStartup( MAKEWORD( 2, 2 ), &m_wsdata ) )
				code = err::ERR_WSA;
			else {
				m_socket = socket( AF_INET, SOCK_STREAM, 0 );
				if( m_socket == INVALID_SOCKET )
					code = err::ERR_WSA;
			}

			if( code != err::ERR_NONE ) {
				MessageBoxA( nullptr, err::translate_err( code ), "", MB_OK );
				return false;
			}

			return true;
		}

		bool connect( ) {
			sockaddr_in server_address{ };
			int code{ };

			server_address.sin_addr.s_addr = m_ip;
			server_address.sin_port = htons( PORT_NUM );
			server_address.sin_family = AF_INET;
			
			code = ::connect( m_socket, ( sockaddr* )( &server_address ), 
				sizeof( server_address ) );

			if( code == -1 ) {
				MessageBoxA( nullptr, err::translate_err( err::ERR_CONNECT ), "", MB_OK );
				return false;
			}

			return true;
		}

		void decode_buffer( uint8_t* buf, size_t length ) {
			auto key = buf[ 0 ];
			for( size_t i{ 1 }; i < length; ++i )
				buf[ i ] ^= key;
		}

		std::string get_string( ) {
			std::string ret{ };
			char buffer[ BUFFER_SIZE ];


			while ( true ) {
				int received = recv( m_socket, buffer, BUFFER_SIZE, 0 );
				if ( received < 0 )
					break;

				for ( int i{ }; i < received; ++i )
					ret.push_back( buffer[ i ] );

				if ( received < BUFFER_SIZE )
					break;
			}

			if ( ret.size( ) ) {
				decode_buffer( ( uint8_t* )ret.data( ), ret.size( ) );
				ret.erase( ret.begin( ) );
			}

			return ret;
		}

		std::vector< uint8_t > get_msg( ) {
			std::vector< uint8_t > ret;
			char buffer[ BUFFER_SIZE ];
			int received = 0;
			
			while( true ) {
				received = recv( m_socket, buffer, BUFFER_SIZE, 0 );
				if( received < 0 )
					break;

				for( int i{ }; i < received; ++i )
					ret.push_back( buffer[ i ] );

				if( received < BUFFER_SIZE )
					break;
			}

			if( ret.size( ) ) {
				decode_buffer( ret.data( ), ret.size( ) );
				ret.erase( ret.begin( ) );
			}
			return ret;
		}

		void send_msg( const uint8_t* msg, int length ) {
			auto buffer = std::make_unique< uint8_t[ ] >( length + 1 );
			uint8_t key = util::random_number( 0, 255 ) & 0xff;

			buffer[ 0 ] = key;
			memcpy( buffer.get( ) + 1,
				msg,
				length );

			for( int i = 1; i <= length; ++i ) {
				buffer[ i ] ^= key;
			}

			int ret = send( m_socket, ( char* )buffer.get( ), length + 1, 0 );
			if ( ret == SOCKET_ERROR ) {
				printf( xors( "error sending message error code: %d" ), WSAGetLastError( ) );
			}
		}

		void send_msg( const char* msg ) {
			auto length = strlen( msg );
			auto buffer = std::make_unique< uint8_t[ ] >( length + 1 );
			uint8_t key = util::random_number( 0, 255 ) & 0xff;

			buffer[ 0 ] = key;
			memcpy( buffer.get( ) + 1,
				msg,
				length );

			for ( size_t i = 1; i <= length; ++i ) {
				buffer[ i ] ^= key;
			}

			int ret = send( m_socket, ( char* )buffer.get( ), (int) length + 1, 0 );
			if ( ret == SOCKET_ERROR ) {
				printf( xors( "error sending message error code: %d" ), WSAGetLastError( ) );
			}
		}


		void send_msg( const char msg ) {
			auto buffer = std::make_unique< uint8_t[ ] >( 2 );
			uint8_t key = util::random_number( 0, 255 ) & 0xff;

			buffer[ 0 ] = key;
			buffer[ 1 ] = msg;
			buffer[ 1 ] ^= buffer[ 0 ];

			int ret = send( m_socket, ( char* )buffer.get( ), 2, 0 );
			if ( ret == SOCKET_ERROR ) {
				printf( xors( "error sending message error code: %d" ), WSAGetLastError( ) );
			}
		}

		void handle( ) {
			auto msg = get_string( );
			if ( msg != xors( "hello" ) ) {
				std::cout << "connection failed." << std::endl;
				//return 0;
			}

			send_msg( "hello" );
		}

	private:
		SOCKET  m_socket;
		WSADATA m_wsdata;
		int     m_ip;
	};
}