/*
  Copyright 2010 Nathan Phillip Brink

  This file is a part of DistRen.

  DistRen is free software: you can redistribute it and/or modify
  it under the terms of the GNU Affero General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  DistRen is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU Affero General Public License for more details.

  You should have received a copy of the GNU Affero General Public License
  along with DistRen.  If not, see <http://www.gnu.org/licenses/>.

*/

#include "common/config.h"

#include "common/libremoteio.h"
#include "common/execio.h"
#include "common/asprintf.h"

#include <list.h>

#include <errno.h>
#ifndef _WIN32
#include <netdb.h>
#endif
#include <poll.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/types.h>
#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#else
#include <sys/socket.h>
#endif
#include <unistd.h>
#include <queue.h>

#ifndef _WIN32
#include <sys/un.h>
#endif

/* local */

#define REMOTEIO_DEFAULT_PORT "4050"

int _remoteio_handle_write(multiio_context_t multiio,
			   int fd,
			   short revent,
			   struct remoteio_opts *opts,
			   struct remoteio *rem);
int _remoteio_handle_read(multiio_context_t multiio,
			  int fd,
			  short revent,
			  struct remoteio_opts *opts,
			  struct remoteio *rem);

int _remoteio_ssh_open(struct remoteio *rem, struct remoteio_server *server);
int _remoteio_ssh_read(struct remoteio *rem, void *buf, size_t len, size_t *bytesread);
int _remoteio_ssh_write(struct remoteio *rem, const void *buf, size_t len, size_t *byteswritten);
int _remoteio_ssh_close(struct remoteio *rem);

/**
   ``named pipes''
 */
#ifndef _WIN32
int _remoteio_sock_open(struct remoteio *rem, struct remoteio_server *server);
int _remoteio_sock_close(struct remoteio *rem);
#endif
int _remoteio_sock_read(struct remoteio *rem, void *buf, size_t len, size_t *bytesread);
int _remoteio_sock_write(struct remoteio *rem, const void *buf, size_t len, size_t *byteswritten);

/**
   These borrow from _remoteio_sock_read() and _remoteio_sock_write().
 */
int _remoteio_tcp_open(struct remoteio *rem, struct remoteio_server *server);
int _remoteio_tcp_close(struct remoteio *rem);

/**
  lookup table for different methods of remoteio:
  the enum remoteio_method is the index of the entry to use for that method. 
  Regardless, a NULL terminator is required because the configuration function
  searches through this table for the method specified in the config file.
*/
struct remoteio_method_funcmap funcmap[] = 
  {
    /* [REMOTEIO_METHOD_SSH] */
    {REMOTEIO_METHOD_SSH, &_remoteio_ssh_open, &_remoteio_ssh_read, &_remoteio_ssh_write, &_remoteio_ssh_close, "ssh"},
#ifndef _WIN32
    {REMOTEIO_METHOD_UNIX, &_remoteio_sock_open, &_remoteio_sock_read, &_remoteio_sock_write, &_remoteio_sock_close, "unix"},
#endif
    {REMOTEIO_METHOD_TCP, &_remoteio_tcp_open, &_remoteio_sock_read, &_remoteio_sock_write, &_remoteio_tcp_close, "tcp"},
    {REMOTEIO_METHOD_SOCKET, NULL, &_remoteio_sock_read, &_remoteio_sock_write, &_remoteio_sock_close},
    {REMOTEIO_METHOD_MAX, NULL, NULL, NULL, NULL, NULL}
  };

struct remoteio_server *remoteio_getserver(const struct remoteio_opts *opts, const char *servername);

void remoteio_packet_free(struct remoteio_packet *packet);

int remoteio_config(cfg_t *cfg, struct remoteio_opts *opts)
{
  size_t numservers;
  size_t counter, counter2;
  static int haslisted_methods = 0;
  
  struct remoteio_server aserver;

  multiio_socket_type_register(opts->multiio, &opts->socket_type);

  multiio_event_handler_register(opts->multiio,
				 opts->socket_type,
				 POLLOUT,
				 (multiio_event_handler_func_t)&_remoteio_handle_write,
				 opts);
  multiio_event_handler_register(opts->multiio,
				 opts->socket_type,
				 POLLIN,
				 (multiio_event_handler_func_t)&_remoteio_handle_read,
				 opts);

  opts->servers = list_init();
  if(!opts->servers)
    {
      fprintf(stderr, "@todo cleanup!\n");
      abort();
    }
  
  numservers = cfg_size(cfg, "server");
  for(counter = 0; counter < numservers; counter ++)
    {
      cfg_t *cfg_aserver;
      char *method;
      
      cfg_aserver = cfg_getnsec(cfg, "server", counter);
      
      aserver.name = strdup(cfg_title(cfg_aserver));
      aserver.hostname = strdup(cfg_getstr(cfg_aserver, "hostname"));
      aserver.username = strdup(cfg_getstr(cfg_aserver, "username"));

      aserver.method = REMOTEIO_METHOD_MAX;
      method = cfg_getstr(cfg_aserver, "method");
      for(counter2 = 0; funcmap[counter2].name; counter2 ++)
	if(strcmp(method, funcmap[counter2].name) == 0)
	  aserver.method = funcmap[counter2].method;
      if(aserver.method == REMOTEIO_METHOD_MAX)
	{
	  fprintf(stderr, "No such method as %s\n", method);
	  if(!haslisted_methods)
	    {
	      fprintf(stderr, "Available methods:\n");
	      for(counter2 = 0; funcmap[counter2].name; counter2 ++)
		fprintf(stderr, "\t%s\n", funcmap[counter2].name);
	      
	      haslisted_methods ++;
	    }
	  abort();
	}
      list_insert_after(opts->servers, &aserver, sizeof(struct remoteio_server));
    }
  
  return 0;
}

int remoteio_generic_data_set(struct remoteio_opts *opts, void *generic_data)
{
  opts->generic_handler_data = generic_data;

  return 0;
}

int remoteio_open_common(struct remoteio **remoteio,
			 enum remoteio_method method,
			 struct remoteio_opts *opts,
			 remoteio_read_handle_func_t read_handler,
			 void *read_handler_data,
			 remoteio_close_handle_func_t close_handler)
{
  struct remoteio *rem;

  rem = malloc(sizeof(struct remoteio));
  if(!rem)
    {
      fprintf(stderr, "OOM\n");
      return 2;
    }

  *remoteio = rem;

  rem->execio = NULL;
  rem->method = method;
  rem->opts = opts;
  rem->inbuf.data = NULL;
  rem->inbuf.len = 0;
  rem->outmsgs = q_init();
  rem->read_handler = read_handler;
  rem->read_handler_data = read_handler_data;
  rem->close_handler = close_handler;
  /*
   * the following initialization is very important... though the
   * others are too, I suppose :-p. See remoteio_close()
   */
  rem->careful_free = 0;

  return 0;
}

int remoteio_open_socket(struct remoteio **remoteio,
			 struct remoteio_opts *opts,
			 remoteio_read_handle_func_t read_handler,
			 void *read_handler_data,
			 remoteio_close_handle_func_t close_handler,
			 int fd)
{
  struct remoteio *rem;
  if(remoteio_open_common(remoteio, REMOTEIO_METHOD_SOCKET, opts, read_handler, read_handler_data, close_handler))
    return 1;
  rem = *remoteio;

  rem->sock = fd;
  multiio_socket_add(opts->multiio, rem->sock, opts->socket_type, rem, POLLIN);

  return 0;
}

int remoteio_open_server(struct remoteio **remoteio,
			 struct remoteio_opts *opts,
			 remoteio_read_handle_func_t read_handler,
			 void *read_handler_data,
			 remoteio_close_handle_func_t close_handler,
			 const char *servername)
{
  struct remoteio_server *theserver;
  struct remoteio *rem;

  int tmp;

  if(!opts)
    {
      fprintf(stderr, "%s:%d: no null opts!\n\tThis is a bug, please report it (after making sure it isn't already reported)\n", __FILE__, __LINE__);
      return 1;
    }

  theserver = remoteio_getserver(opts, servername);
  if(!theserver)
    {
      fprintf(stderr, "%s:%d: Could not find server named ``%s''\n", __FILE__, __LINE__, servername);
      return 1;
    }

  if(theserver->method >= REMOTEIO_METHOD_MAX
     || theserver->method < 0)
    {
      fprintf(stderr, "%s:%d: Unsupported remoteio method %d\n\tThis is a bug, probably indicating memory corruption. This is, of course, probably my fault (not your hardware's) ;-)\n", __FILE__, __LINE__, theserver->method);
      return 1;
    }

  if(remoteio_open_common(remoteio, theserver->method, opts, read_handler, read_handler_data, close_handler))
    return 1;
  rem = *remoteio;

  tmp = funcmap[theserver->method].open_func(rem, theserver);
  if(tmp)
    {
      fprintf(stderr, "Error using method %s for server ``%s''", funcmap[theserver->method].name, servername);
      free(rem->inbuf.data);
      q_free(rem->outmsgs, QUEUE_NODEALLOC);
      free(rem);
      *remoteio = NULL;
      return tmp;
    }

  /**
   * @todo make this code slightly more generic... able to handle
   * execio's multi-sockets by letting execio register itself with
   * multiio instead of us registering here perhaps
   */
  multiio_socket_add(opts->multiio, rem->sock, opts->socket_type, rem, POLLIN);
  
  return 0;
}

/**
 * Implementation of multiio_event_handler_func_t
 */
int _remoteio_handle_read(multiio_context_t multiio,
			  int fd,
			  short revent,
			  struct remoteio_opts *opts,
			  struct remoteio *rem)
{
  struct remoteio_packet packet;
  size_t readlen;
  char buf[8192];

  int tmp;

  packet.len = 0;
  packet.data = NULL;

  if(rem->sock != fd)
    fprintf(stderr, "%d != %d\n", rem->sock, fd);

  tmp = funcmap[rem->method].read_func(rem, buf, sizeof(buf), &readlen);
  if(tmp)
    {
      remoteio_close(rem);
      return 1;
    }

  /* expand the input buffer */
  packet.len = rem->inbuf.len + readlen;
  packet.data = malloc(rem->inbuf.len + readlen);
  if(!packet.data)
    {
      fprintf(stderr, "OOM!\n");

      return 1;
    }
  if(rem->inbuf.data)
    memcpy(packet.data, rem->inbuf.data, rem->inbuf.len);
  memcpy(packet.data + rem->inbuf.len, buf, readlen);
  free(rem->inbuf.data);
  memcpy(&rem->inbuf, &packet, sizeof(struct remoteio_packet));

  /*
   * readlen wil now keeps track of how many bytes the handler
   * function has read.
   *
   * Call the read_handler. Set careful_free, see remoteio_close(), so
   * that rem->read_handler() may call remoteio_close() without
   * segfaulting us ;-).
   */
  rem->careful_free = 1;
  readlen = (*rem->read_handler)(rem, rem->opts->generic_handler_data, rem->inbuf.data, rem->inbuf.len, rem->read_handler_data);
  if(rem->careful_free == 2)
    {
      rem->careful_free = 0;
      remoteio_close(rem);

      return 0;
    }
  rem->careful_free = 0;

  memmove(rem->inbuf.data, rem->inbuf.data + readlen, rem->inbuf.len - readlen);
  rem->inbuf.len -= readlen;

  return 0;
}


int remoteio_write(struct remoteio *rem, const void *buf, size_t len)
{
  struct remoteio_packet *packet;
  struct pollfd pollfd;
  size_t bytes_written;
  int err;

  /**
   * This is probably about the only optimization that exists in
   * distren.... :-D
   *
   * Write to the client immediately if there are no other messages
   * waiting and if the client will accept it.
   */
  if(q_empty(rem->outmsgs))
    {
      pollfd.fd = rem->sock;
      pollfd.revents = POLLOUT;
      pollfd.events = 0;
      poll(&pollfd, 1, 0);
      if(pollfd.events & POLLOUT)
	{
	  err = funcmap[rem->method].write_func(rem, buf, len, &bytes_written);
	  if(bytes_written > 0)
	    {
	      len -= bytes_written;
	      buf += bytes_written;
	    }
	}
    }

  /**
   * zero length is easy... and might be possible if the above
   * optimization works ;-)
   */
  if(!len)
    return 0;

  packet = malloc(sizeof(struct remoteio_packet));
  if(!packet)
    {
      fprintf(stderr, "OOM\n");
      return 1;
    }

  packet->len = len;
  packet->data = malloc(len);
  if(!packet->data)
    {
      free(packet);
      fprintf(stderr, "OOM\n");
      return 1;
    }

  memcpy(packet->data, buf, len);

  q_enqueue(rem->outmsgs, packet, 0);
  multiio_socket_event_enable(rem->opts->multiio, rem->sock, POLLOUT);

  return 0;
}

int _remoteio_handle_write(multiio_context_t multiio,
			   int fd,
			   short revent,
			   struct remoteio_opts *opts,
			   struct remoteio *rem)
{
  struct remoteio_packet *packet;
  size_t written_amount;

  int tmp;

  fprintf(stderr, "%s:%d: <del>My</del><ins>Someone else's</ins> traversal says that sock %d is available for writing\n",
	  __FILE__, __LINE__, fd);

  /*
   * check if we're out of stuff to write.
   */
  if(q_empty(rem->outmsgs))
    {
      multiio_socket_event_disable(multiio, fd, POLLOUT);
      return 0;
    }

  packet = q_front(rem->outmsgs);
  tmp = funcmap[rem->method].write_func(rem, packet->data, packet->len, &written_amount);

  /**
     Disconnect in case of write error.
  */
  if(tmp)
    {
      fprintf(stderr, __FILE__ ":%d: error handling for write() needs to be inserted into remoteio.... perhaps.. ;-)\n", __LINE__);
    }
  if(packet->len == written_amount)
    {
      q_dequeue(rem->outmsgs);
      remoteio_packet_free(packet);

      if(q_empty(rem->outmsgs))
	multiio_socket_event_disable(multiio, fd, POLLOUT);
    }
  else
    {
      /**
       * shifting seems the simplest solution.
       */
      packet->len -= written_amount;
      memmove(packet->data, packet->data + written_amount, packet->len);
    }

  return 0;
}


int remoteio_close(struct remoteio *rem)
{
  int rtn;
  
  /**
   * See careful_free's and _remoteio_handle_read()'s docs.  If
   * careful_free is nonzero, then we shouldn't free it here because
   * such a free would cause a segfault. However, whoever set
   * rem->careful_free to nonzero will handle resetting
   * rem->careful_free to zero and calling remoteio_close() if
   * necessary.
   */
  if(rem->careful_free)
    {
      rem->careful_free = 2;
      return 0;
    }

  /* call close handler */
  if(rem->close_handler)
    (*rem->close_handler)(rem->opts->generic_handler_data, rem->read_handler_data);

  /* cleanup multiiio stuff */
  multiio_socket_del(rem->opts->multiio, rem->sock);

  /* backend-specific cleanup */
  rtn = funcmap[rem->method].close_func(rem);

  /* this part is normal ;-) */
  free(rem->inbuf.data);
  q_free(rem->outmsgs, (list_dealloc_func_t)remoteio_packet_free);

  free(rem);

  return rtn;
}

/**
 * Frees an entire packet, including the passed pointer. If you just
 * want the contents of the packet free()ed, just do
 * free(packet.data);
 */
void remoteio_packet_free(struct remoteio_packet *packet)
{
  free(packet->data);
  free(packet); 
}


int _remoteio_getserver_traverse(char *servername, struct remoteio_server *aserver)
{
  if(!strcmp(aserver->name, servername))
    return FALSE; /* stop traversal */

  return TRUE;
}

struct remoteio_server *remoteio_getserver(const struct remoteio_opts *opts, const char *servername)
{
  int traversal_result;
  char *dispensible_servername;

  dispensible_servername = strdup(servername); /* for the sake of constness... */
  traversal_result = list_traverse(opts->servers, dispensible_servername, (list_traverse_func_t)&_remoteio_getserver_traverse, LIST_FRNT | LIST_ALTR);
  free(dispensible_servername);

  if(traversal_result == LIST_OK)
    return (struct remoteio_server *)list_curr(opts->servers);

  return (struct remoteio_server *)NULL;
}

size_t remoteio_sendq_len(const struct remoteio *rem)
{
  return (size_t)q_size(rem->outmsgs);
}

/**
   different remoteio methods' implementations:
 */

/*
  SSH, via execio
*/

int _remoteio_ssh_open(struct remoteio *rem, struct remoteio_server *server)
{
  char *userhost;
  char *sshargs[] = {rem->opts->ssh_command, NULL /* userhost */, "distrend", "-d", (char *)NULL};

  int rtn;

  if(server->username)
    _distren_asprintf(&userhost, "%s@%s", server->username, server->hostname);
  else
    userhost = strdup(server->hostname);
  sshargs[1] = userhost;

  rtn = execio_open( &rem->execio, "ssh", sshargs);
  if(rtn)
    {
      fprintf(stderr, "error opening remoteio channel to ssh userhost ``%s''\n" , userhost);
      free(userhost);
      return 1;
    }
  free(userhost);
  
  return 0;
}

int _remoteio_ssh_read(struct remoteio *rem, void *buf, size_t len, size_t *bytesread)
{
  return execio_read(rem->execio, buf, len, bytesread);
}

int _remoteio_ssh_write(struct remoteio *rem, const void *buf, size_t len, size_t *byteswritten)
{
  return execio_write(rem->execio, buf, len, byteswritten);
}

int _remoteio_ssh_close(struct remoteio *rem)
{
  int rtn;
  
  rtn = execio_close(rem->execio);
  if(rtn)
    fprintf(stderr, "%s:%d: error closing execio\n", __FILE__, __LINE__);
  
  return rtn;
}

#ifndef _WIN32
/*
  local sockets implementation (``named pipes''), unix-only
 */
int _remoteio_sock_open(struct remoteio *rem, struct remoteio_server *server)
{
  int sock;
  struct sockaddr_un sockaddr;

  /*
    The POSIX docs pretty much say that I can't depend on sockpath being able to be longer than 
    some proprietary length. So, if the compiler specifies a long path for RUNSTATEDIR, it could
    cause a buffer overflow.
   */
  char *sockpath = RUNSTATEDIR "/distrend.sock";
  unsigned int sockaddr_len;

  sock = socket(AF_UNIX, SOCK_STREAM, 0);
  if(sock == -1)
    {
      perror("socket");
      return 1;
    }

  sockaddr.sun_family = AF_UNIX;
  /*
    The terminating NULL should not be included in what's copied to sun_path,
    although it won't hurt as long as strlen(sockpath) < max socket length
   */
  for(sockaddr_len = 0; sockpath[sockaddr_len]; sockaddr_len ++)
    sockaddr.sun_path[sockaddr_len] = sockpath[sockaddr_len];

  if(connect(sock, (struct sockaddr *)&sockaddr, sockaddr_len) == -1)
    {
      perror("connect");
      close(sock);
      return 1;
    }

  rem->sock = sock;

  return 0;
}

#endif /* _WIN32 */

int _remoteio_sock_close(struct remoteio *rem)
{
  close(rem->sock);

  return 0;
}

int _remoteio_sock_read(struct remoteio *rem, void *buf, size_t len, size_t *bytesread)
{
  ssize_t readrtn;

  *bytesread = 0;
  readrtn = read(rem->sock, buf, len);
  /*
    The following is valid for blocking sockets:
   */
  if(readrtn == -1)
    {
      /*
	in this case, we may have been interrupted by a signal and errno == EINTR
	or the connection was reset and errno = ECONNRESET

	Some of these are not error conditions:
       */
      perror("read");

      if(errno != EINTR)
	{
	  fprintf(stderr, "error reading socket in remoteio_sock_read\n");
	  return 1;
	}

      return 0;
    }

  *bytesread = readrtn;
  if(!readrtn)
    {
      /*
	means EOF except when FD is in ``message-nondiscard'' or ``message-discard''
	modes.
       */
      return 1;
    }

  return 0;
}

int _remoteio_sock_write(struct remoteio *rem, const void *buf, size_t len, size_t *byteswritten)
{
  ssize_t writertn;

  writertn = write(rem->sock, buf, len);

  if(writertn == -1)
    {
      perror("write");
      if(errno != EINTR)
	{
	  fprintf(stderr, "error writing to socket in remoteio_sock_write()\n");
	  return 1;
	}
    }

  *byteswritten = writertn;
  if(!writertn)
    {
      /*
	should we consider this an error? I'm pretty
	sure we should :-/
       */
      fprintf(stderr, "write() returned 0 in remoteio_sock_write()\n");
      return 1;
    }

  return 0;
}

/**
   TCP, taking advantage of the two generic socks read/write functions:
 */
int _remoteio_tcp_open(struct remoteio *rem, struct remoteio_server *server)
{
  int tmp;
  int tmp2;

  int sock;

  char *hostname;
  char *port;

  struct addrinfo addrinfo_hints;
  struct addrinfo *addrinfo_res;

  static char *default_port = REMOTEIO_DEFAULT_PORT;

  /**
     only hostname should be free()-ed, not port,
     because both are from the same block of malloc()-ed
     memory
   */
  hostname = strdup(server->hostname);
  for(port = hostname;
      *port && *port != ':';
      port ++)
    ;
  if(*port)
    {
      *port = '\0';
      port ++;
    }
  else
    port = default_port;

  memset(&addrinfo_hints, '\0', sizeof(struct addrinfo));
  addrinfo_hints.ai_family = AF_UNSPEC;
#ifdef _WIN32
  /* windows lacks stuff documented in POSIX, I guess :-( */
  addrinfo_hints.ai_flags = 0;
#else
  addrinfo_hints.ai_flags = AI_V4MAPPED | AI_ADDRCONFIG;
#endif
  addrinfo_hints.ai_socktype = SOCK_STREAM;

  tmp = getaddrinfo(hostname, port, &addrinfo_hints, &addrinfo_res);
  if(tmp)
    {
      fprintf(stderr, "error resolving %s:%s: %s\n", server->hostname, port, gai_strerror(tmp));
      free(hostname);
      return 1;
    }
  fprintf(stderr, "connecting to %s[%s]:%s\n", server->hostname, addrinfo_res->ai_canonname, port);

  free(hostname);

  sock = socket(addrinfo_res->ai_family, SOCK_STREAM, addrinfo_res->ai_protocol);
  if(sock == -1)
    {
      perror("socket");
      freeaddrinfo(addrinfo_res);
  
      return 1;
    }

  tmp = connect(sock, addrinfo_res->ai_addr, addrinfo_res->ai_addrlen);
  tmp2 = errno;
  freeaddrinfo(addrinfo_res);
  errno = tmp2;
  if(tmp == -1)
    {
      perror("connect");
      close(sock);

      return 1;
    }

  rem->sock = sock;

  return 0;
}

int _remoteio_tcp_close(struct remoteio *rem)
{
  close(rem->sock);

  return 0;
}