// TODO FOR TOMORROW:
// - FIGURE OUT WHEN WE'RE REGISTERED (NICK and USER from user are replied to with commands 001, 002, 003, and 004 from the ircd when registration successful)
// - int registered = 0; already created to track when this has happened (assuming still want to do it this way)
// 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

#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 SOURCE_SERVER 0
#define SOURCE_CLIENT 1
#define EXCEPT_NONE 0

#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 (CRLF or space) separated tokens per server response we expect (TODO - check this is reasonable) (CRLF and spaces just grouped here due to laziness)
#define MAXPONGSIZE 32 // let's assume PING/PONG responses can't be larger than this (TODO - check this [it's totally made up]!)

#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 (EXCEPT_NONE) 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) {
  appendcrlf(str); // Do this just before sending so callers don't need to worry about it
  str_len = strlen(str); // Recalculate str_len in case it changed (TODO: so do we even need to pass it to this function?)
  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 IRC message (if anything)
// by splitting out the different components by space character (ASCII 0x20).
// Get serversockfd and clientsockfd in case we need to send a response.
//
// str is the raw string
//
// source is either 0 (SOURCE_SERVER) (server - from the real IRC server) or 1 (SOURCE_CLIENT) (client - from a real IRC client)
//
// 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 processircmessage(int *serversockfd, int *clientsockfd, char *str, int source) {
  // 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("  >> processircmessage(): Processing source %d message \"%s\"...\n", source, str);

  char *token;

  while ((token = strsep(&str, " ")) != NULL) {
    if (*token  == '\0') continue; // Skip consecutive matches
    printf("  >> Message 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++;
  }

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

  switch(source) {
    case SOURCE_SERVER: // If message(s) were from the real IRC 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);
        }
        sendtoserver(serversockfd, outgoingmsg, strlen(outgoingmsg));

        // We processed something so return true
        return 1;
      }
      break;
    case SOURCE_CLIENT: // If message(s) were from a real IRC client
      break;
    default:
      fprintf(stderr, "Unexpected raw IRC string source!\n");
      exit(1);
  }
  // =============================================>


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

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

// Figure out what the client or server sent and decide what to do with it (if anything)
// First by splitting out the different messages by CRLF, process if we know how to,
// or send on to all other clients (if from the server) or to the server (if from a client)
// if we don't know how to.
//
// serversockfd, clientsockfd, fdmax, arr_clients, except all passed to here so we can
// send/relay a response to server or (other) client(s) for things we don't know how to process here,
// or for any other reasons we need to send stuff for
//
// str is the raw string
//
// source is either 0 (SOURCE_SERVER) (server - from the real IRC server) or 1 (SOURCE_CLIENT) (client - from a real IRC client)
//
// Return 0 if something went wrong
// Return 1 if everything OK
int processrawstring(int *serversockfd, int *clientsockfd, char *str, int source, int fdmax, int arr_clients[], int except) {
  // Copy to a temporary string so we still have the original in case it's not processed
  char *strcopy = strdup(str);

  printf("processrawstring(): Source type %d sent: \"%s\".  Processing it...\n", source, strcopy);

  // Track which CLRF-separated message within this string we're on
  int messagecount = 0;

  // Build array of each space-separated token
  char messages[MAXTOKENS][MAXDATASIZE];

  // Split the string by CRLF and add each CRLF-separated IRC message to an array
  char *token;

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

  free(strcopy);

  // Go through each message, figure out what it is and if we're doing anything with it
  for (int i = 0; i < messagecount; i++) {
    // Copy to a temporary string so we still have the original in case it's not processed
    char *messagecopy = strdup(messages[i]);
    if (processircmessage(serversockfd, clientsockfd, messagecopy, source)) {
      printf("Message processed: \"%s\", NULLing...\n", messages[i]);
      messages[i][0] = '\0';
    }
    free(messagecopy);
  }

  // Deal with any unprocessed messages (send on to server/client(s) since we haven't dealt with them)
  for (int i = 0; i < messagecount; i++) {
    if (messages[i][0] != '\0') {
      printf("Message %d (\"%s\") remains unprocessed...\n", i, messages[i]);
      switch(source) {
        case SOURCE_SERVER: // If message(s) were from the real IRC server
          // Relay/send to all clients ("except" = 0 because this should send to all clients)
          // TODO - Is this really going to send the original string if we have messed it with it in processrawstring() and friends!?
          printf("bouncer-server: sending \"%s\" to all clients, length %zd.\n", messages[i], strlen(messages[i]));
          sendtoallclients(clientsockfd, fdmax, arr_clients, messages[i], strlen(messages[i]), EXCEPT_NONE);
          break;
        case SOURCE_CLIENT: // If message(s) were from a real IRC client
          // Send to server
          printf("bouncer-client: sending \"%s\" to the server, length %zd.\n", messages[i], strlen(messages[i]));
          sendtoserver(serversockfd, messages[i], strlen(messages[i]));

          // send the same thing to all *other* clients (all except for fd "i")
          sendtoallclients(clientsockfd, fdmax, arr_clients, messages[i], strlen(messages[i]), except);
          break;
        default:
          fprintf(stderr, "Unexpected raw IRC string source!\n");
          exit(1);
      }
    }
  }

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

  return 1;
}

// 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)
  sendtoserver(serversockfd, outgoingmsg, strlen(outgoingmsg));

  // Send our USER

  snprintf(outgoingmsg, MAXDATASIZE, "USER %s 8 * : %s", IRCUSER, IRCREALNAME); // TODO - Check for success (with return code)
                                                                                // TODO - Send a more intelligent/correct USER string
  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);

      // Try to process received string (which should contain one or more server responses/commands)
      // TODO - What if there were two server respones/commands and only one didn't need relaying?
      if (!processrawstring(serversockfd, clientsockfd, serverbuf, SOURCE_SERVER, fdmax, arr_clients, EXCEPT_NONE)) {
        fprintf(stderr, "Error: bouncer-server failed to process raw string.\n");
        exit(1);
      }
    }

    // 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);
      }

      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 is the current iteration of this loop, since 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);

            // Try to process received string (which should contain one or more client responses/commands)
            // TODO - What if there were two server respones/commands and only one didn't need relaying?
            if (!processrawstring(serversockfd, clientsockfd, clientbuf, SOURCE_CLIENT, fdmax, arr_clients, i)) {
              fprintf(stderr, "Error: bouncer-client failed to process raw string.\n");
              exit(1);
            }
          }
        }
      }
    }
  }
}

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;
}