#include <windows.h>
#include <WinSock2.h>
#include <iostream>

using namespace std;

const int STRLEN = 256; char message[STRLEN] = {0};
unsigned int port = 666;
char recMessage[STRLEN] = {0};
HANDLE h = GetStdHandle(STD_OUTPUT_HANDLE);

class Socket
{
    protected:
        WSADATA     wsaData;
        SOCKET      mySocket;
        SOCKET      myBackup;
        SOCKET      acceptSocket;
        sockaddr_in sockad;
    public:
        Socket();
        ~Socket();
        bool SendData( char* );
        bool RecvData( char*, int );
        void CloseConnection();
        void GetAndSendMessage();
};

class ClientSocket : public Socket  //Inheritance
{
    public:
        void ConnectToServer( const char *ipAddress, int port );
}; ClientSocket sockClient;

void ClientSocket::ConnectToServer( const char *ipAddress, int port )
{
    sockad.sin_family = AF_INET;
    sockad.sin_addr.s_addr = inet_addr( ipAddress );
    sockad.sin_port = htons( port );

    if ( connect( mySocket, (SOCKADDR*) &sockad, sizeof( sockad ) ) == SOCKET_ERROR )
    {
        cerr << "ClientSocket: Failed to connect\n";
        system("pause");
        WSACleanup();
        exit(13);
    }
    cout << "\nCONNECTED!\n\n";
}

void Socket::GetAndSendMessage()   //Member functions
{
    cin.ignore();
    message[0] = '\0';

    cout << "You    > ";
    cin.get( message, STRLEN );
    SendData( message );
}

bool Socket::SendData( char *buffer )
{
     send( mySocket, buffer, strlen( buffer ), 0 );
     return true;
}

bool Socket::RecvData( char *buffer, int size )
{
     int i = recv( mySocket, buffer, size, 0 );

     cout<<WSAGetLastError()<<"\r\n";
     cout <<i<< "\n";

     buffer[i] = '\0';
     return true;
}

Socket::Socket()
{
    if( WSAStartup( MAKEWORD(2, 2), &wsaData ) != NO_ERROR )
    {
        cerr << "Socket Initialization: Error with WSAStartup\n";
        system("pause");
        WSACleanup();
        exit(10);
    }

    //Create a socket
    mySocket = socket( AF_INET, SOCK_STREAM, IPPROTO_TCP );

    if ( mySocket == INVALID_SOCKET )
    {
        cerr << "Socket Initialization: Error creating socket" << endl;
        system("pause");
        WSACleanup();
        exit(11);
    }

    myBackup = mySocket;
}

Socket::~Socket()
{
    WSACleanup();
}
/////////////////////////////////////////////////////////////////////////
class ServerSocket : public Socket  //Inheritance
{
    public:
        void Bind( int port );
        void Listen();
        void StartHosting( int port );
};    ServerSocket sockServer;

void ServerSocket::StartHosting( int port )
{
     Bind( port );
     Listen();
}

void ServerSocket::Bind( int port )
{
     sockad.sin_family = AF_INET;
     sockad.sin_addr.s_addr = inet_addr( "0.0.0.0" );
     sockad.sin_port = htons( port );

     if ( bind ( mySocket, (SOCKADDR*) &sockad, sizeof( sockad) ) == SOCKET_ERROR )
     {
         cerr << "ServerSocket: Failed to connect\n";
         system("pause"); WSACleanup(); exit(14);
     }
}

void ServerSocket::Listen()
{
     if ( listen ( mySocket, 1 ) == SOCKET_ERROR )
     {
         cerr << "ServerSocket: Error listening on socket\n";
         system("pause");
         WSACleanup();
         exit(15);
     }

     acceptSocket = accept( myBackup, NULL, NULL );
     while ( acceptSocket == SOCKET_ERROR )
     {
         acceptSocket = accept( myBackup, NULL, NULL );
     }
     mySocket = acceptSocket;
     cout << "\nCONNECTED!\n\n";
}

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

int main()
{
    string choice;
    string ipAddress;
    while(1){
    cout << "1) Client" << endl;
    cout << "2) Server" << endl;
    cin  >> choice;

    /***************************************CLIENT***************************************/
    if ( choice == "1" )
    {
         cout << "ATTEMPTING TO CONNECT..." << endl;
         sockClient.ConnectToServer( "127.0.0.1", port );
        while (1)
        {
           sockClient.GetAndSendMessage();
        }
        sockClient.CloseConnection();
    }
    /****************************************SERVER***************************************/
    else if ( choice == "2" )
    {
       //(HANDLE)_beginthreadex(NULL, 0, LISTENandACCEPT , NULL, 0, &Th2);
        cout << "\t--------------\n" << endl;
        cout << "\t****SERVER****\n" << endl;
        cout << "\t--------------\n" << endl;
        cout << "WAITING FOR CLIENT..." << endl;
        sockServer.StartHosting( port );

        while(1)
        {
              sockClient.RecvData( recMessage, STRLEN );
              cout << "Friend > " << recMessage << endl; Sleep(1999);
        }
    }
    else  cout << "\nType 1 or 2 Please!\n\n";
    }
}