#ifndef _UDP_SOCKET_H
#define _UDP_SOCKET_H

/*

UDP (UNIVERSAL DATAGRAM PROTOCOL) C++ CLASS...

Author: Matthew W. Coan
Date: Sat Jan  7 21:35:49 EST 2017

*/

#include <cstdio>
#include <cstdlib> 
#include <cstring>

#ifdef __WIN32__
#include <windows.h>
#include <winsock2.h>
#include <Ws2tcpip.h>
int inet_pton(int af, const char *src, void *dst)
{
  struct sockaddr_storage ss;
  int size = sizeof(ss);
  char src_copy[INET6_ADDRSTRLEN+1];

  ZeroMemory(&ss, sizeof(ss));
  /* stupid non-const API */
  strncpy (src_copy, src, INET6_ADDRSTRLEN+1);
  src_copy[INET6_ADDRSTRLEN] = 0;

  if (WSAStringToAddress(src_copy, af, NULL, (struct sockaddr *)&ss, &size) == 0) {
    switch(af) {
      case AF_INET:
    *(struct in_addr *)dst = ((struct sockaddr_in *)&ss)->sin_addr;
    return 1;
      case AF_INET6:
    *(struct in6_addr *)dst = ((struct sockaddr_in6 *)&ss)->sin6_addr;
    return 1;
    }
  }
  return 0;
}
#else
#include <netdb.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <arpa/inet.h>
#include <unistd.h>

#include <sys/types.h>
#include <unistd.h>

#include <netinet/in.h>
#include <arpa/inet.h>

#include <netdb.h>
#define closesocket(fd) close(fd)
#define SOCKET int
#endif



#include <string>
#include <iostream>

#include "bsstream.h"

namespace net_tools {

using namespace std;

class udp_socket {
public:
   enum udp_type { UDP_SERVER, UDP_CLIENT };
   struct udp_exception { };

private:
   string host; 
   int port;
   struct hostent  *ptrh;
   struct protoent *ptrp;
   struct sockaddr_in sad;
   SOCKET fd;
   struct sockaddr fskt;
   int fskt_len;
   int n;

public: 
   udp_socket(const string & host, const int port, udp_socket::udp_type the_type) throw(udp_exception) {
      this->host = host;
      this->port = port;
      n = 0;

      fskt_len = sizeof(fskt);

      memset((char *)&sad,0,sizeof(sad));
      sad.sin_family = AF_INET;
      sad.sin_port = htons((u_short)port);

      ptrh = gethostbyname(host.c_str());

      if ( ((char *)ptrh) == NULL ) {
         fprintf(stderr,"invalid host: %s\n", host.c_str());
         throw udp_exception();
      }

      memcpy(&sad.sin_addr, ptrh->h_addr, ptrh->h_length);

      if ( ((size_t)(ptrp = getprotobyname("udp"))) == 0) {
         fprintf(stderr, "cannot map \"udp\" to protocol number");
         throw udp_exception();
      }

      fd = socket(AF_INET, SOCK_DGRAM, ptrp->p_proto);

      if (fd < 0) {
         fprintf(stderr, "socket creation failed\n");
         throw udp_exception();
      }

      switch(the_type) {
      case UDP_SERVER:
         sad.sin_addr.s_addr = htonl(INADDR_ANY);
         if (bind(fd, (struct sockaddr *)&sad, sizeof(sad)) < 0) {
            cerr << "unable to bind socket..." << endl;
            throw udp_exception();
         }
      break;

      case UDP_CLIENT:
#ifdef __WIN32__
         if (inet_pton(AF_INET, host.c_str(), (void*)&sad.sin_addr)==0) {
#else
         if (inet_aton(host.c_str(), &sad.sin_addr)==0) {
#endif
            cerr << "unable to inet_aton..." << endl;
            throw udp_exception();
         }
      break;
      }
   }

   ~udp_socket() {
      closesocket(fd);
   }

   int send(const char * buf, const size_t size) {
      n = sendto(fd, buf, size, 0, (struct sockaddr *)&sad, sizeof(sad));
      return n;
   }
   
   int recive(char * buf, const size_t size) {
#ifdef __WIN32__
      n = recvfrom(fd, buf, size, 0, (struct sockaddr *)&sad, (ssize_t*)&fskt_len); 
#else
      n = recvfrom(fd, buf, size, 0, (struct sockaddr *)&sad, (socklen_t*)&fskt_len); 
#endif

      return n;
   }

   operator void*() {
      return ((void*)(n != -1));
   }
};

}

#endif /* _UDP_SOCKET_H */
