#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdio.h>
#include <unistd.h>
#include <netdb.h>
#include <stdlib.h>
#include <sys/time.h>
#include <time.h>
#include <pthread.h>

#include "socket.h"

#define mem_alloc(a,b) malloc(a)
#define mem_free(a) free(a)
#define mem_realloc(a,b,c) realloc(a,b)

#define max(a,b) ((a) > (b)) ? (a) : (b)
#define min(a,b) ((a) < (b)) ? (a) : (b)

struct asocket
   {
   int sockfd;
   struct sockaddr_in addr;
   };

/* Connects to a host  on a given port
 */

asocket *socket_connect (char *host, int port)
   {
   asocket *retval;
   struct hostent *he;
   
   retval = mem_alloc (sizeof (asocket), "socket_connect");
      
   if ((he = gethostbyname(host)) == NULL)
      {
      perror("gethostbyname");
      mem_free (retval);
      return NULL;
      }
   
   if ((retval->sockfd = socket(AF_INET, SOCK_STREAM, 0)) == -1) 
      {
      perror("socket");
      mem_free (retval);
      return NULL;
      }
   
   retval->addr.sin_family = AF_INET;      /* host byte order */
   retval->addr.sin_port = htons(port);    /* short, network byte order */
   retval->addr.sin_addr = *((struct in_addr *)he->h_addr);
   memset (retval->addr.sin_zero, 0, 8);
   
   if (connect (retval->sockfd, (struct sockaddr *)&retval->addr, 
		sizeof(struct sockaddr)) == -1)
      {
      perror("connect");
      mem_free (retval);
      return NULL;
      }
         
   return retval;
   }

/* Sets up a socket for listening on port port
 * backport is set to 10
 */

asocket *socket_listen(int port)
	{
	asocket *retval;

	retval = mem_alloc (sizeof(asocket), "socket_listen");
	if (!retval)
		return NULL;

	if ((retval->sockfd = socket(PF_INET, SOCK_STREAM, 0)) == -1)
		{
		mem_free(retval);
		perror("socket error");
		return NULL;
		}

	/* tracer_log("Got sockfd: %d", retval->sockfd); */
	retval->addr.sin_family = AF_INET;
	retval->addr.sin_port = htons(port);
	retval->addr.sin_addr.s_addr = INADDR_ANY;
	memset(&retval->addr.sin_zero, 0, 8);

	if (bind(retval->sockfd, (struct sockaddr *)&retval->addr,
	         sizeof(struct sockaddr)) == -1)
		{
		mem_free(retval);
		perror("bind error");
		return NULL;
		}

	if (listen(retval->sockfd, 10) == -1)
		{
		perror("listen error");
		mem_free(retval);
		return NULL;
		}

	return retval;
	}

/* waits for a new connection to an already listening socket
 * mallocs memory for the return value (which should later be free'd
 * using socket_destroy)
 */
asocket *socket_accept(asocket * socket)
	{
	socklen_t sin_size;
	asocket *retval;

	sin_size = sizeof(struct sockaddr_in);

	retval = mem_alloc(sizeof(asocket), "socket_accept");

	if ((retval->sockfd = accept(socket->sockfd,
	                             &retval->addr, &sin_size)) == -1)
		{
		mem_free(retval);
		perror("accept error");
		return NULL;
		}

	return retval;
	}

/* Starts up a server on an already defined socket. When a connection 
 * is made, connect_cb is called, if this returns true, then the connection
 * is allowed, otherwise it is destroyed. The user_data pointer can be set
 * to an arbitrary value, this will be passed to the read_cb callback. Each
 * socket will have its own user_data pointer.
 * Whenever a socket has data waiting, read_cb is called with that socket's
 * user_data pointer. This function should read all available data from
 * the socket
 */ 
int socket_server_start_callback (asocket *socket,
				  int (*connect_cb)(asocket *socket,
						    void **user_data),
				  int (*read_cb)(asocket *socket,
						 void *user_data))
   {
   fd_set fds;
   int i, retval, done = 0, n, n_clients = 0;
   asocket **clients = NULL;
   asocket *new_client;
   void **user_data = NULL;   
   socklen_t sin_size;

   sin_size = sizeof(struct sockaddr_in);
   
   while (!done)
      {
      FD_ZERO (&fds);
      
      FD_SET (socket->sockfd, &fds);
      n = socket->sockfd + 1;
      for (i = 0; i < n_clients; i++)
	 if (clients[i])
	    {
	    FD_SET (clients[i]->sockfd, &fds);
	    n = max (n, clients[i]->sockfd + 1);
	    }
      
      retval = select (n, &fds, NULL, NULL, NULL);

      if (retval)
	 {
	 if (FD_ISSET (socket->sockfd, &fds)) /* new connection */
	    {
	    n_clients++;
	    clients = mem_realloc (clients, n_clients * sizeof (asocket *),
				   "socket_server_start_callback");
	    user_data = mem_realloc (user_data, n_clients * sizeof (void *),
				     "socket_server_start_callback2");
	    clients[n_clients - 1] = socket_accept (socket);
	    new_client = clients[n_clients - 1];
	    if (connect_cb)
	       connect_cb (new_client, &user_data[n_clients - 1]);
	    else
	       user_data[n_clients - 1] = NULL;
	    FD_SET (new_client->sockfd, &fds);
	    n = max (n, new_client->sockfd + 1);
	    }
	 else /* existing socket connection */
	    {
	    int removed = 0;
	    
	    for (i = 0; i < n_clients; i++)
	       {
	       if (clients[i] && FD_ISSET (clients[i]->sockfd, &fds))
		  {
		  if (read_cb)
		     if (read_cb (clients[i], user_data[i]) == -1)
			{
			FD_CLR (clients[i]->sockfd, &fds);
			socket_destroy (clients[i]);
			clients[i] = NULL;
			removed = 1;
			}
		  break;
		  }
	       }
	    
	    if (removed)
	       {
	       n = socket->sockfd + 1;
	       for (i = 0; i < n_clients; i++)
		  {
		  if (clients[i])
		     {
		     n = max (n, clients[i]->sockfd + 1);
		     }
		  }
	       }
	    }
	 }
      }
   
   return 0;
   }

typedef struct
   {
   asocket *socket;
   int (*cb)(asocket *socket, void *user_data);
   void *user_data;
   } socket_thread;

static void *start_thread (void *socket)
   {
   socket_thread *data = socket;
   
   data->cb (data->socket, data->user_data);
   
   return NULL;
   }

int socket_server_start_threaded (asocket *socket,
				  int (*connect_cb)(asocket *socket,
						    void **user_data),
				  int (*start_cb)(asocket *socket,
						  void *user_data))
   {
   asocket *new_client;
   int done = 0;
   socket_thread data;
   void *user_data;
   pthread_t thread;
   
   while (!done)
      {
      user_data = NULL;
      new_client = socket_accept (socket);
      if (new_client)
	 {
	 if (!connect_cb || connect_cb (new_client, &user_data))
	    {
	    if (start_cb)
	       {
	       data.socket = new_client;
	       data.cb = start_cb;
	       data.user_data = user_data;	    
	       pthread_create (&thread, NULL, start_thread, &data);
	       }
	    else
	       socket_destroy (new_client);
	    }
	 }
      }

   
   return 0;
   }



/* sends a block of data down a socket. Every MAXPACKETLEN bytes, it waits
 * to receive a single byte containing the number 1 (0x01), if this is not
 * received then there is an error
 */
int socket_send(asocket * socket, void *buffer, int len)
	{
	int retval;
	   
	if (!socket)
	   return -1;
	
	retval = send(socket->sockfd, buffer, len, 0);
	return retval;
	}

/* sends a line of characters, appending \n\r to the end of them
 * Useful for text protocols like HTTP, POP etc...
 * returns the length of the line that was successfully sent (does
 * not include the length of \n\r)
 */
int socket_sendline (asocket *socket, char *line)
   {
   int retval, len;
   
   if (!socket)
      return -1;
   
   if (line)
      len = strlen (line);
   else
      len = 0;
   
   retval = send (socket->sockfd, line, len, 0);
   if (retval != len)
      return -1;
   
   if (send (socket->sockfd, "\n\r", 2, 0) != 2)
      return -1;
   
   return retval;
   }

/* receives data from an existing socket. every MAXPACKETLEN bytes it will
 * send a single bytes containing '1' (0x01)
 */
int socket_recv(asocket * socket, void *buffer, int len)
	{
	int retval;

	if (!socket)
		return -1;
	retval = recv(socket->sockfd, buffer, len, 0);
	
	return retval;
	}

/* Reads from the socket until maxlen bytes read, or \n is read
 * Null terminates the string
 */ 
int socket_recvline(asocket *socket, char *buffer, int maxlen)
   {
   int done = 0, tmp, len = 0;
   char ch;

   if (!socket)
      {
      buffer[0] = '\0';
      return -1;
      }

   while (!done)
      {
      tmp = recv (socket->sockfd, &ch, 1, 0);

      if (tmp != 1)
	 {
	 buffer[len] = '\0';
	 len = -1;
	 break;
	 }

      if (ch == '\n')
	 {
	 buffer[len] = '\0';
	 done = 1;
	 }
      else
	 {
	 buffer[len] = ch;
	 len++;
	 }

      if (len == maxlen)
	 {
	 buffer[maxlen] = '\0';
	 done = 1;
	 }
      }

   return len;   
   }

/* free's all resources associated with a socket */
int
socket_destroy(asocket * socket)
	{
	if (!socket)
		return 0;
	close(socket->sockfd);
	
	mem_free(socket);
	return 0;
	}

/* returns some info about a socket, atm this is its IP address in string
 * format
 */
char *
socket_info(asocket * socket)
	{
	if (!socket)
		return "Invalid socket";
	return inet_ntoa(socket->addr.sin_addr);
	}