// TODO FOR TOMORROW:
// - SPLIT THE RECEIVED MESSAGES INTO TOKENS AND FIGURE OUT WHEN WE'RE REGISTERED (NICK and USER from user are replied to with commands 001, 002, 003, and 004 from ircd when registration successful
// - perhaps split with strtok()
// - int registered = 0; already created to track when this has happened (assuming still want to do it this way)
// - Actually maybe let's just try to register with the server ourselves without a client (telnet might be more useful than irssi for debugging)
// - implement automatic PING/PONG responses
// Example WHOIS reply:
// BOUNCER-SERVER RECEIVED: :irc.tghost.co.uk 307 blabounce l_bratch :is identified for this nick

// "server" means the real IRC server
// "client" means bouncer clients

// uses socket demo code from https://beej.us/guide/bgnet/html/single/bgnet.html
// and getstdin() uses getLine() from https://stackoverflow.com/questions/4023895/

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <netdb.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <sys/select.h>

#include "functions.h"
#include "sockets.h"

#define MAXDATASIZE 513 // max number of bytes we can get at once (RFC2812 says 512, plus one for null terminator)
#define STDIN 0 // stdin is fd 0
#define MAXCLIENTS 32 // maximum number of clients that can connect to the bouncer at a time
#define MAXTOKENS 100 // maximum number of space-separated tokens per server response
#define MAXPONGSIZE 32 // let's assume PING/PONG responses can't be larger than this (TODO - check this!)

#define IRCNICK "blabounce" // TODO: change this to a config option!
#define IRCUSER "blabounce" // TODO: change this to a config option!
#define IRCREALNAME "Mr Bla Bouncer" // TODO: change this to a config option!

int debugmode = 0;

// Relay/send message to all clients (optionally except one)
// "except" is used to send to all clients _except_ the fd provided (except = 0 avoids this, i.e. sends to all)
int sendtoallclients(int *clientsockfd, int fdmax, int arr_clients[], char *str, int str_len, int except) {

  char *sendertype;

  // Decide what sort of text to prefix the debug output with
  // At the moment if non-zero "except" is specified then it must be a message from a bouncer client
  // and if "except" is zero then it must be a message from the real IRC server
  if (except) {
    sendertype = "bouncer-client";
  } else {
    sendertype = "bouncer-server";
  }

  // relay/send to all clients...
  for (int i = *clientsockfd + 1; i <= fdmax; i++) {
    // Skip the current client if "except" non-zero (no need to send back to itself)
    if (i == except) {
      continue;
    }
    // TODO maybe see if things are in rfds (not sure what this means any more - perhaps it was to do with only sending to connected things which is now solved)
    // ...but only if they are connected
    for (int j = 0; j < MAXCLIENTS; j++) {
      if (arr_clients[j] == i) {
        printf("%s: sending %s to client with fd %d.\n", sendertype, str, i);
        if (send(i, str, str_len, 0) == -1) {
          perror("send");
        }
      }
    }
  }

  return 0;
}

// Send whatever string to the real IRC server
int sendtoserver(int *serversockfd, char *str, int str_len) {
  printf("sendtoserver(): sending %s to IRC server (length %d).\n", str, str_len);
  if (send(*serversockfd, str, str_len, 0) == -1) { // 0 is bitwise OR for no send() flags
    printf("send error, exiting...\n");
    perror("send");
  }

  return 0;
}

// Figure out what to do with each CRLF-split server response (if anything)
// Now by splitting out the different components by space (ASCII 0x20)
// Get serversockfd in case we need to send a response
//
// Return 1 if we processed something and expect the caller to not need to do anything more
// Return 0 if we didn't process it and the caller might want to do something
int processserverresponse(int *serversockfd, char *str) {
  // Track which space-separated token within this response we're on
  int counter = 0;

  // Build array of each space-separated token (TODO - Use counter to stop splitting once we reach some reasonable value - i.e. once we're definitely past commands and into just free text)
  char tokens[MAXTOKENS][MAXDATASIZE];

  printf("  >> processserverresponse(): The server response is: \"%s\".  Processing it...\n", str);

  char *token;

  while ((token = strsep(&str, " ")) != NULL) {
    if (*token  == '\0') continue; // Skip consecutive matches
    printf("Response Token: \"%s\", length %zd.\n", token, strlen(token));
    // Copy into the token array (strlen + 1 to get the NULL terminator)
    strncpy(tokens[counter], token, strlen(token) + 1);
    counter++;
  }

  printf("We found these tokens:\n");
  for (int i = 0; i < counter; i++) {
    printf("Token %d: \"%s\"\n", i, tokens[i]);
  }

  // <=============================================
  // IRC command processing (commands from server)

  // PING received?  If so, send a PONG back with the next element as the argument.
  if (strncmp(tokens[0], "PING", strlen(tokens[0])) == 0) {
    printf("PING found and it is: %s with length %zd!  Sending response...\n", tokens[0], strlen(tokens[0]));

    char outgoingmsg[MAXDATASIZE]; // String to send to server
    if (!snprintf(outgoingmsg, MAXDATASIZE, "PONG %s", tokens[1])) { // TODO - Make sure tokens[1] actually has a token
      fprintf(stderr, "Error while preparing PONG response!\n");
      exit(1);
    }
    appendcrlf(outgoingmsg);
    sendtoserver(serversockfd, outgoingmsg, strlen(outgoingmsg));

    // We processed something so return true
    return 1;
  }
  // =============================================>

  printf("Done with processserverresponse()\n");

  // If we got here then we didn't process anything
  return 0;
}

// Figure out what the server sent and decide what to do with it (if anything)
// First by splitting out the different responses by CRLF
// Get serversockfd in case we need to send a response
//
// Return 1 if we processed something and expect the caller to not need to do anything more
// Return 0 if we didn't process it and the caller might want to do something
int processserverstring(int *serversockfd, char *str) {
  // Track if something in this string got processed
  int processed = 0;

  printf("processserverstring(): The server sent: \"%s\".  Processing it...\n", str);

  // Split the string by CRLF and send each CRLF-separated string (each "response") on to processserverresponse(s)
  char *token;

  while ((token = strsep(&str, "\r\n")) != NULL) {
    if (*token  == '\0') continue; // Skip consecutive matches
    printf("String Token: \"%s\", length %zd.\n", token, strlen(token));
    // Pass on for the individual response to be processed
    if (processserverresponse(serversockfd, token)) {
      processed = 1;
    }
  }

  printf("Done with processserverstring()\n");

  return processed;
}

// Where the big bouncing loop is
void dochat(int *serversockfd, int *clientsockfd) {
  char serverbuf[MAXDATASIZE]; // buffer for receiving data on server socket
  char clientbuf[MAXDATASIZE]; // buffer for receiving data on client socket(s)
  int servernumbytes; // Number of bytes received from remote server
  char outgoingmsg[MAXDATASIZE]; // String to send to server
  int outgoingmsgrc; // Return code from getstdin() for outgoing message
  int arr_clients[MAXCLIENTS]; // Array of all clients - 0 means not connected, greater than 0 means connected and the value is the fd number (so we know which ones to try to read)
  int num_clients = 0; // Current number of clients

  int fdmax; // highest numbered socket fd

  socklen_t addrlen; // client remote address size
  char remoteIP[INET6_ADDRSTRLEN]; // remote IP (assume up to IPv6 size)
  int newfd;  // newly accept()ed socket descriptor
  struct sockaddr_storage remoteaddr; // client address
  int clientnumbytes;

  fdmax = *clientsockfd; // keep track of highest fd number, currently client socket as created last (after server socket)

  fd_set rfds; // set of read fds to monitor with select() - 0: stdin, 1: stdout, 2: stderr, 3 and higher: sockets (at time of writing, 3: real IRC server, 4: client listener, 5 and higher: clients)

  // set all the clients to be "not connected"
  for (int i = 0; i < MAXCLIENTS; i++) {
    arr_clients[i] = 0;
  }

  // <=============================================
  // Initialise IRC connecting/registration state

  // Registered with the IRCd yet?
//////////////  int registered = 0;

  // Send our NICK
  snprintf(outgoingmsg, MAXDATASIZE, "NICK %s", IRCNICK); // TODO - Check for success (with return code)
  appendcrlf(outgoingmsg);
  sendtoserver(serversockfd, outgoingmsg, strlen(outgoingmsg));

  // Send our USER

  snprintf(outgoingmsg, MAXDATASIZE, "USER %s 8 * : %s", IRCUSER, IRCREALNAME); // TODO - Check for success (with return code)
  appendcrlf(outgoingmsg);
  sendtoserver(serversockfd, outgoingmsg, strlen(outgoingmsg));
  // =============================================>

  while (1) {
    printf("top of loop, fdmax %d.\n", fdmax);
    FD_ZERO(&rfds); // clear entries from fd set

    FD_SET(STDIN, &rfds); // add STDIN (fd 0) to read fds to monitor
    FD_SET(*serversockfd, &rfds); // add our server network socket to monitor
    FD_SET(*clientsockfd, &rfds); // add our client network socket to monitor

    // Add all connected clients to monitor (only add ones that are connected (clients[i] > 0))
    // TODO - make sure *serversockfd stays at the same value (probably 3?) in all cases - what if the server disconnects/reconnects/etc.
    // TODO - now that only connected clients are monitored, perhaps tracking using both fdmax and num_client loops is unnecessary?
    for (int i = 0; i < MAXCLIENTS; i++) {
      if (arr_clients[i] > 0) {
        printf("monitoring fd %d.\n", arr_clients[i]);
        FD_SET(arr_clients[i], &rfds);
      }
    }

    printf("select()ing...\n");
    // check to see if anything in the fd_set is waiting - waits here until one of the fds in the set does something
    if (select(fdmax + 1, &rfds, NULL, NULL, NULL) < 0) { // network socket + 1, rfds, no writes, no exceptions/errors, no timeout
      printf("receive error, exiting!?\n");
      perror("select");
    }

    // TODO - switch around the serversockfd and STDIN FD_ISSET if-statements?  They feel the wrong way round.  Are they like this on purpose?  I can't remember.
    // (although actually stdin may go once its not wanted for possible direct interaction for debugging)

    // See if there's anything to read from the server side (the real IRCd)
    if (FD_ISSET(*serversockfd, &rfds)) {
      printf("reading server socket!\n");

      if ((servernumbytes = recv(*serversockfd, serverbuf, MAXDATASIZE - 1, 0)) == -1) {
        printf("receive error (-1), exiting...\n");
        perror("recv");
        exit(1);
      } else if (servernumbytes == 0) {
        printf("socket closed (or no data received) (0), exiting...\n");
        perror("recv");
        exit(1);
      }
      serverbuf[servernumbytes] = '\0';

      printf("BOUNCER-SERVER RECEIVED: %s\n", serverbuf);
      printf("bouncer-server: sending it to all clients...\n");

      // Try to process received string (which should contain one or more server responses/commands)
      // If 0 returned, do nothing because we processed it elsewhere
      // If non-0 returned, send to all clients to see if they want to handle it
      // TODO - What if there were two server respones/commands and only one didn't need relaying?
      if (!processserverstring(serversockfd, serverbuf)) {
        // Relay/send to all clients ("except" = 0 because this should send to all clients)
        sendtoallclients(clientsockfd, fdmax, arr_clients, serverbuf, servernumbytes, 0);
      }
    }

    // see if there's anything from stdin
    if (FD_ISSET(STDIN, &rfds)) {
      printf("reading stdin!\n");

      outgoingmsgrc = getstdin(NULL, outgoingmsg, sizeof(outgoingmsg));

      if (outgoingmsgrc == NO_INPUT) {
        printf("\nError!  No input.\n");
      } else if (outgoingmsgrc == TOO_LONG) {
        printf ("Error!  Too long.  Would have allowed up to: [%s]\n", outgoingmsg);
      }

      appendcrlf(outgoingmsg);

      sendtoserver(serversockfd, outgoingmsg, strlen(outgoingmsg));
    }

    // go through all the remaining sockets to see if there's anything from the client sockets (either new connections or existing clients sending messages)
    // (clear newfd before doing this so we can tell if we're querying a new client or not)
    newfd = 0;
    for (int i = *clientsockfd; i <= fdmax; i++) {
      // skip if newfd as we know we have just accept()ed it
      if (i == newfd) {
        continue;
      }
      printf("checking client socket %d out of %d.\n", i, fdmax);
      if (FD_ISSET(i, &rfds)) {
        printf("fd %d is FD_ISSET and it is a...\n", i);
        // if value of clientsockfd then must be a new connection, if greater must be an existing connection
        if (i == *clientsockfd) {
          printf("...new connection!\n");
          // handle new connections
          if (num_clients >= MAXCLIENTS) {
            fprintf(stderr, "too many clients!\n");
            exit(1); // TODO - handle cleanly instead of exiting!
          }
          addrlen = sizeof remoteaddr;
          newfd = accept(*clientsockfd, (struct sockaddr *)&remoteaddr, &addrlen);
          if (newfd == -1) {
            // something went wrong when accept()ing
            perror("accept");
          } else {
            FD_SET(newfd, &rfds); // add to master set // TODO - needed?  going to be re-done at the top anyway...
            if (newfd > fdmax) {    // keep track of the max
              fdmax = newfd;
            }
            // Find a free element in the clients array and set to new fd value
            for (int j = 0; j < MAXCLIENTS; j++) {
                if (arr_clients[j] == 0) {
                  arr_clients[j] = newfd;
                  break;
                }
            }
            // TODO - Handle the "find a free element" loop not finding a free element
            num_clients++; // Track total number of clients
            printf("bouncer-client: new connection from %s on socket %d\n", inet_ntop(remoteaddr.ss_family, get_in_addr((struct sockaddr*)&remoteaddr), remoteIP, INET6_ADDRSTRLEN), newfd);
            printf("bouncer-client: total client connections: %d\n", num_clients);
          }
        } else {
          printf("...previous connection!\n");
          // handle data from a client
          if ((clientnumbytes = recv(i, clientbuf, sizeof clientbuf, 0)) <= 0) {
            // got error or connection closed by client
            if (clientnumbytes == 0) {
              // connection closed
              printf("bouncer-client: socket %d hung up\n", i);
            } else {
              perror("recv");
            }
            close(i); // bye!
            FD_CLR(i, &rfds); // remove from master set - TODO is this needed at the moment since we just add everything from *clientsockfd to fdmax to rfds
            // Remove the client from the clients array
            for (int j = 0; j < MAXCLIENTS; j++) {
                if (arr_clients[j] == i) {
                  printf("found and clearing fd %d from arr_clients[%d]\n", i, j);
                  arr_clients[j] = 0;
                  break;
                }
            }
            // TODO - Handle the "remove the client" loop not finding the old fd
            num_clients--; // Track total number of clients
            printf("bouncer-client: total client connections: %d\n", num_clients);
          } else {
            // we got some data from a client
            // null terminate that baby
            clientbuf[clientnumbytes] = '\0'; // TODO make sure this can't overrun if some super long line (max bytes?) was received
            // clear up any newlines
            while (clientbuf[strlen(clientbuf) - 1] == '\n' || clientbuf[strlen(clientbuf) - 1] == '\r') {
              clientbuf[strlen(clientbuf) - 1] = '\0';
            }
            printf("BOUNCER-CLIENT RECEIVED: '%s'\n", clientbuf);
            printf("bouncer-client: sending it to the server...\n");
            appendcrlf(clientbuf);
            sendtoserver(serversockfd, clientbuf, strlen(clientbuf));

            // send the same thing to all *other* clients (all except for fd "i")
            sendtoallclients(clientsockfd, fdmax, arr_clients, clientbuf, strlen(clientbuf), i);
          }
        }
      }
    }
  }
}

int main(int argc, char *argv[]) {
  if (argc < 3) {
    fprintf(stderr,"usage: %s hostname port [-d]\n", argv[0]);
    exit(1);
  }

  if (argc == 4) {
    if (!strcmp(argv[3], "-d")) {
      debugmode = 1;
      debug("debug mode enabled\n");
    }
  }

  // TODO: see if any of this can be shared (i.e. 1. avoid code duplication, and 2. see if variables can be shared between client/server sockets)

  // TODO: track fdmax - kind of doing this now with arr_clients and num_clients but might be pointlessly tracking both in some places (?)

  // I will try to keep to the notation of "server" meaning the real IRCd, "bouncer" meaning the bouncer, and "client" meaning the real IRC client

  // BOUNCER-TO-SERVER socket things

  // Create server socket
  int serversockfd = createserversocket(argv[1], argv[2]);

  // Create client socket (after server so we can use its fd number later as fdmax)
  int clientsockfd = createclientsocket(BOUNCERLISTENPORT);

  dochat(&serversockfd, &clientsockfd);

  printf("dochat() complete, closing socket...\n");

  close(serversockfd);

  return 0;
}