/* This file is part of GNU Pies.
   Copyright (C) 2019 Sergey Poznyakoff

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

   GNU Pies 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 General Public License for more details.

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

#include <config.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <poll.h>
#include <errno.h>
#include <libpies.h>
#include <iobuf.h>

typedef struct netcat_server netcat_server_t;
typedef struct netcat_stream netcat_stream_t;
typedef ssize_t (*netcat_stream_io_t) (netcat_server_t *srv);
typedef int (*netcat_stream_disconnect_t) (netcat_server_t *srv);

enum { IN, OUT };

struct netcat_server
{
  char const *id;
  netcat_server_t *next, *prev;
  struct pollfd *pollfd;
  int state;
  netcat_stream_disconnect_t disconnect;
  netcat_server_t *peer;
  struct iobuf buf[2];  
};

static netcat_server_t *server_head, *server_tail;

netcat_server_t *
netcat_server_create (char const *id, struct pollfd *pfd, int state,
		      netcat_stream_disconnect_t dis, netcat_server_t *peer)
{
  netcat_server_t *srv = grecs_zalloc (sizeof (*srv));
  srv->id = id;
  srv->pollfd = pfd;
  srv->state = state;
  srv->pollfd->events |= srv->state;
  srv->disconnect = dis;
  srv->peer = peer;
  if (peer)
    srv->peer->peer = srv;
  iobuf_reset (&srv->buf[IN]);
  iobuf_reset (&srv->buf[OUT]);

  srv->prev = server_tail;
  srv->next = NULL;
  if (server_tail)
    server_tail->next = srv;
  else
    server_head = srv;
  server_tail = srv;

  return srv;
}

void
netcat_server_remove (netcat_server_t *srv)
{
  netcat_server_t *p;
  if ((p = srv->next) != NULL)
    p->prev = srv->prev;
  else
    server_tail = srv->prev;
  if ((p = srv->prev) != NULL)
    p->next = srv->next;
  else
    server_head = srv->next;
  if (srv->peer)
    srv->peer->peer = NULL;
  free (srv);
}

/*
net_reader
1. If there is free space in the IN buffer, fill it with the data read
from the fd and return
2. Otherwise, clear the POLLIN bit.

stdout_writer (peer -> net_reader)
1. If OUT buffer is empty
  1.1 If the peer's IN buffer is not empty, transfer data from it
  1.2 Else raise the peers POLLIN bit and return.
2. Write as much as possible from OUT to fd

stdin_reader (same as net_reader)

net_writer (peer -> stdin_reader)  same as stdout_writer
*/

ssize_t
netcat_stream_read (netcat_server_t *srv)
{
  ssize_t n;
  
  if (iobuf_avail_size (&srv->buf[IN]))
    {
      n = iobuf_fill (&srv->buf[IN], srv->pollfd->fd);
      if (n == -1)
	{
	  grecs_error (NULL, errno, "%s: read", srv->id);
	  srv->pollfd->events &= ~POLLIN;
	  return -1;
	}
      if (n == 0)
	{
	  /* No more input is expected */
	  srv->disconnect (srv);
	  srv->state &= ~POLLIN;
	  srv->pollfd->events &= ~POLLIN;
	  if (srv->pollfd->events == 0)
	    srv->pollfd->fd = -1;
	}
    }
  else
    {
      srv->pollfd->events &= ~POLLIN;
      n = 0; /* catch timeout? */
    }
  return n;
}

static inline int
peer_is_state (netcat_server_t *srv, int state)
{
  return srv->peer && (srv->peer->state & state);
}

ssize_t
netcat_stream_write (netcat_server_t *srv)
{
  ssize_t n;
  
  if (iobuf_data_size (&srv->buf[OUT]) == 0)
    {
      if (!peer_is_state (srv, POLLIN))
	{
	  // shutdown write end
	  srv->disconnect (srv);
	  srv->state &= ~POLLOUT;
	  srv->pollfd->events &= ~POLLOUT;
	  if (srv->pollfd->events == 0)
	    srv->pollfd->fd = -1;
	  return -1;
	}
      if (iobuf_copy (&srv->buf[OUT], &srv->peer->buf[IN]) == 0)
	{
	  srv->peer->pollfd->events |= POLLIN;
	  return 0;
	}
    }
  n = iobuf_flush (&srv->buf[OUT], srv->pollfd->fd);
  if (n == -1)
    {
      grecs_error (NULL, errno, "%s: write", srv->id);
      return -1;
    }
  if (n == 0)
    {
      // FIXME: eof
      return -1;
    }
  return 0;
}

int
disconnect_in (netcat_server_t *srv)
{
  return shutdown (srv->pollfd->fd, SHUT_RD);
}

int
disconnect_out (netcat_server_t *srv)
{
  return shutdown (srv->pollfd->fd, SHUT_WR);
}

int
disconnect_stdin (netcat_server_t *srv)
{
  return close (srv->pollfd->fd);
}

int
disconnect_stdout (netcat_server_t *srv)
{
  close (srv->pollfd->fd);
  exit (0);
}

static int
netcat (char const *urlstr)
{
  int fd;
  struct pies_url *url;
  struct pollfd pfd[3];
  int nfd = sizeof (pfd) / sizeof (pfd[0]);
  netcat_server_t *srvin, *srvout, *srv;
  
  if (pies_url_create (&url, urlstr))
    {
      perror (urlstr);
      return 64;
    }

  fd = url_connect (url, NULL);
  if (fd == -1)
    return 1;

  pfd[0].fd = 0;
  pfd[0].events = 0;
  srvin = netcat_server_create ("stdin", &pfd[0], POLLIN, disconnect_stdin, NULL);

  pfd[1].fd = 1;
  pfd[1].events = 0;
  srvout = netcat_server_create ("stdout", &pfd[1], POLLOUT, disconnect_stdout, NULL);

  pfd[2].fd = fd;
  pfd[2].events = 0;

  netcat_server_create ("netread", &pfd[2], POLLIN, disconnect_in, srvout);
  netcat_server_create ("netwrite", &pfd[2], POLLOUT, disconnect_out, srvin);

  while (server_head)
    {
      ssize_t n = poll (pfd, nfd, -1);
      if (n == -1)
	{
	  if (errno != EINTR)
	    grecs_error (NULL, errno, "poll");
	  continue;
	}
      for (srv = server_head; srv; )
	{
	  netcat_server_t *next = srv->next;
	  int events;
	  
	  events = (srv->pollfd->events|POLLHUP)
	           & (srv->pollfd->revents & (srv->state|POLLHUP));
	  if (events)
	    {
	      if (events & POLLHUP)
		{
		  //grecs_error (NULL, 0, "HUP on %s", srv->id);
		  srv->disconnect (srv);
		  srv->pollfd->events &= ~srv->state;
		  if (srv->pollfd->events == 0)
		    srv->pollfd->fd = -1;
		  srv->state = 0;
		}
	      else if (events & POLLIN)
		netcat_stream_read (srv);
	      else if (events & POLLOUT)
		netcat_stream_write (srv);
	    }
	  if (srv->state == 0 || srv->pollfd->fd == -1)
	    netcat_server_remove (srv);
	  srv = next;
	}
    }
  return 0;
}

static void
redirect (int sfd, char const *name)
{
  int fd;

  fd = open (name, sfd ? (O_CREAT | O_TRUNC | O_WRONLY) : O_RDONLY, 0644);
  if (fd == -1)
    {
      perror (name);
      exit (1);
    }
  if (dup2 (fd, sfd) == -1)
    {
      perror ("dup2");
      exit (1);
    }
}

static void
usage (FILE *fp)
{
  fprintf (fp, "usage: nt [-i FILE] [-o FILE] URL\n");
}

int
main (int argc, char **argv)
{
  int c;
  
  while ((c = getopt (argc, argv, "i:o:")) != EOF)
    {
      switch (c)
	{
	case 'i':
	  redirect (0, optarg);
	  break;

	case 'o':
	  redirect (1, optarg);
	  break;

	default:
	  exit (64);
	}
    }

  argc -= optind;
  argv += optind;

  if (argc != 1)
    {
      usage (stderr);
      exit (64);
    }

  return netcat (argv[0]);
}