#include "shared.h"

void signal_handler(int sig); //handle signals for exit
void serverSYN(PacketManager* manager, Packet packet, int socketFd, struct sockaddr_in clientAddr);
void serverACK(std::unordered_map<int, PacketManager*>& connections, Packet packet, int socketFd, struct sockaddr_in clientAddr);
time_t serverFIN(std::unordered_map<int, PacketManager*>& connections, Packet packet, int socketFd, struct sockaddr_in clientAddr);


void serverResponse(Header head, int fd, sockaddr_in clientAddr);
void printServerOut(Packet packet, std::string type, bool dup=false);

int main(int argc, char *argv[])
{
  //Error and Signal checkin
  signal(SIGINT, signal_handler);
  signal(SIGQUIT, signal_handler);
  signal(SIGTERM, signal_handler);
  if(argc != 3) {
    exit(EXIT_FAILURE);
  }
  std::string portString = argv[1];
  for(char c : portString) {
    if(!isdigit(c)) {
      std::cerr << "ERROR: port has non digit" << std::endl;
      exit(EXIT_FAILURE);
    }
  }
  int portNum = std::stoi(argv[1]);
  if(portNum < PORT_MIN || portNum > PORT_MAX) {
    std::cerr << "ERROR: port number outside range" << std::endl;
    exit(EXIT_FAILURE);
  }
  std::string fileDir = argv[2];

  /*
  This part sets up the sockets needed to run the server and connect to the client
  */
  //Init socket
  int socketFd;
  struct sockaddr_in serverAddr;
  if ((socketFd = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP)) == -1) {
    exit(EXIT_FAILURE);
  }
  //Debugging purposes
  int optval = 1;
  setsockopt(socketFd, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval));
  //----------------
  //Bind socket call

  serverAddr.sin_family = AF_INET;
  serverAddr.sin_port = htons(portNum);
  int s_addrlen = sizeof(serverAddr);
  //serverAddr.sin_addr.s_addr = INADDR_ANY;
  serverAddr.sin_addr.s_addr = inet_addr("127.0.0.1"); //(using arpa/inet.h)
  memset(serverAddr.sin_zero, '\0', sizeof(serverAddr.sin_zero));
  if (bind(socketFd, (struct sockaddr *)&serverAddr, s_addrlen) == -1) {
    exit(EXIT_FAILURE);
  }
  int maxfdp1, nready;
  fd_set rset;
  FD_ZERO(&rset);
  maxfdp1 = socketFd + 1;
  /*
  Actual UDP portion
  */
  struct timeval timeout;
  timeout.tv_sec = 10;
  char msg[MAXBUFFER];
  socklen_t sin_size = sizeof(struct sockaddr_storage);
  int connectionCount = 1;
  //Loop through
  std::unordered_map<int, PacketManager*> connections;
  std::vector<FinCheck*> checker;
  while(true) {
    bzero(msg, sizeof(msg));
    struct sockaddr_in clientAddr;
    //Set up Header and Packet for easier analysis
    struct Header info = {0};
    ssize_t bytesRead;
    FD_SET(socketFd, &rset);
    nready = select(maxfdp1, &rset, NULL, NULL, &timeout);
    if(FD_ISSET(socketFd, &rset)) {
      if((bytesRead = recvfrom(socketFd, msg, sizeof(msg), 0, (struct sockaddr *)&clientAddr, &sin_size)) == -1) {
        std::cerr << "ERROR: first recv bad" << std::endl;
        exit(EXIT_FAILURE);
      }
      //Copy and convert to Host Byte Order, then stores in class for use
      memcpy(&info, msg, sizeof(Header));
      info = convertHeaderHost(info);
      std::vector<char> data(&msg[12], &msg[(int)bytesRead]); //Hopefully doesn't run into issue with \0
      Packet packet(info, data);
      //MAYBE ADD LINES HERE FOR WRONG FLAGS (DROP)?
      //SYN received
      if(packet.isSYN() && packet.getID() == 0) {
        //TODO maybe check if data is empty, otherwise DROP
        std::string filePath = fileDir + std::to_string(connectionCount) + ".file";
        PacketManager* current = new PacketManager(packet, filePath, connectionCount);
        connections.insert(std::pair<int, PacketManager*>(connectionCount, current));
        serverSYN(current, packet, socketFd, clientAddr);
        connectionCount++;
      }
      //ACK received
      if(packet.isACK() || packet.getHeader().mFlags == 0) {
        serverACK(connections, packet, socketFd, clientAddr);
      }
      //FIN received
      if(packet.isFIN()) {
        time_t res = serverFIN(connections, packet, socketFd, clientAddr);
        FinCheck* curr = new FinCheck;
        curr->mTimer = res;
        PacketManager* packetm = connections[packet.getID()];
        curr->mManager = packetm;
        Header connect = makeHeader(packetm->getACK(), packetm->getSeq(), packet.getID(), FIN_ACK);
        curr->mHeader = packet.getHeader();
        curr->mClientAddr = clientAddr;
        checker.push_back(curr);
      }
    }
    //Delete anything timeout 10s and ERROR in file
    for(auto it = connections.begin(); it != connections.end();) {
      PacketManager* temp = it->second;
      time_t curr = time(0);
      if(difftime(curr, temp->getTime()) >= 10.0) {
        temp->writeError();
        delete temp;
        temp=nullptr;
        it = connections.erase(it);
      } else {
        it++;
      }
    }
    //FIN re-transmit
    for(auto it = checker.begin(); it != checker.end();) {
      FinCheck* curr = *it;
      //delete from vector if already deleted, and delete if FINACKED from client
      if(curr->mManager == nullptr) {
        delete curr;
        it = checker.erase(it);
        continue;
      } else {
        if(curr->mManager->getFINACKED()) {
          int ID = curr->mManager->getID();
          if(curr->mManager != nullptr) {
            delete curr->mManager;
            curr->mManager = nullptr;
            connections.erase(ID);
          }
        }
        //Re-transmit every .5 seconds if no ACK, until 10s (timeout)
        if(difftime(time(0), curr->mTimer) >= 0.5) {
          curr->mTimer = time(0);
          Packet temp = Packet(curr->mHeader);
          serverResponse(curr->mHeader, socketFd, curr->mClientAddr);
          printServerOut(temp, SEND, true);
        }
        it++;
      }
    }
  }
  //End of while loop, clean up code.
  for(auto it = connections.begin(); it != connections.end();) { //PacketManager clean
    delete it->second;
    it = connections.erase(it);
  }
  for(auto it = checker.begin(); it != checker.end(); it++) { //FinCheck clean
    delete *it;
  }
  close(socketFd);
  return 0;
}

//Function to interrupt from loop and close program
void signal_handler(int sig) {
  exit(EXIT_SUCCESS);
}


//Function to create an empty file corresponding to the connection count and send a response
void serverSYN(PacketManager* manager, Packet packet, int socketFd, struct sockaddr_in clientAddr) {
  printServerOut(packet, RECV);
  //Create empty file
  std::ofstream fs(manager->getPath());
  if(!fs.is_open()) {
    std::cerr << "ERROR: FILE NOT OPEN" << std::endl;
    exit(EXIT_FAILURE);
  }
  fs.close();
  //Response
  Packet response = Packet(makeHeader(4321, packet.getSeq()+1, manager->getID(), SYN_ACK));
  serverResponse(response.getHeader(), socketFd, clientAddr);
  printServerOut(response, SEND);
  manager->updateACK(manager->getACK()+1);
  return;
}

//Function to ACK correct files and drop others
void serverACK(std::unordered_map<int, PacketManager*>& connections, Packet packet, int socketFd, struct sockaddr_in clientAddr) {
  printServerOut(packet, RECV);
  int ID = packet.getID();
  PacketManager* manager = connections[ID];
  if(manager->getFIN()) { //if FIN-ACKED, doesn't need to send anything and will be cleaned up in main loop
    manager->updateTime();
    manager->recvFINACKED();
    return;
  }
  if(packet.getData().size() == 0) {
    //manager->updateSEQ(packet.getAck());
    Packet temp = Packet(makeHeader(manager->getSeq(), manager->getACK(), manager->getID(), ACK));
    serverResponse(temp.getHeader(), socketFd, clientAddr);
    printServerOut(temp, SEND);
    return; //No data to be processed means we don't need to ACK it
  }
  if(connections.find(packet.getID()) == connections.end()) { //Connection not found, drop
    printServerOut(packet, DROP);
    return;
  } else { //Found connection
    manager->updateTime();
    //Check packet's order
    unsigned int curr = manager->getACK();
    bool dup = false;
    if(curr == packet.getSeq()) { //in-order ACK, write now
      manager->updateACK(packet.getSeq()+packet.getData().size());
      if(packet.getAck() != 0){
        manager->updateSEQ(packet.getAck());
      }
      manager->writeFileCurr(packet.getData());
    } else {
      if(packet.getSeq() > curr + RWND) { //Out of RWND, drop
        printServerOut(packet, DROP);
        return;
      }
      //Store into buffer if within RWND
      manager->addBuffer(packet.getSeq() % (MAX_NUM + 1), packet);
      dup = true;
    }
    //Response
    //std::cout << "here" << std::endl;
    Packet temp = Packet(makeHeader(manager->getSeq(), manager->getACK(), manager->getID(), ACK));
    serverResponse(temp.getHeader(), socketFd, clientAddr);
    printServerOut(temp, SEND, dup);
    return;
  }
}

//Function to mark a PacketManager as FIN and respond 
time_t serverFIN(std::unordered_map<int, PacketManager*>& connections, Packet packet, int socketFd, struct sockaddr_in clientAddr) {
  printServerOut(packet, RECV);
  int ID = packet.getID();
  Header response = makeHeader(connections[ID]->getSeq(), 0, ID, FIN);
  Packet out = Packet(response);
  //Mark as FIN received, respond then let main loop handle re-transmit and timeout
  time_t start = time(0);
  connections[ID]->updateTime();
  if(!connections[ID]->getFIN()) connections[ID]->updateACK(connections[ID]->getACK()+1);
  connections[ID]->recvFIN();
  Header temp = makeHeader(connections[ID]->getSeq(), connections[ID]->getACK(), ID, FIN_ACK);
  serverResponse(temp, socketFd, clientAddr);
  //serverResponse(response, socketFd, clientAddr);
  printServerOut(Packet(temp), SEND);
  //printServerOut(out, SEND);
  return start;
}

//Function to send a response to client (header only, no data attached)
void serverResponse(Header head, int fd, sockaddr_in clientAddr) {
  ssize_t bytesRead;
  socklen_t sin_size = sizeof(struct sockaddr_storage);
  char msg[12];
  head = convertHeaderNetwork(head);
  memcpy(&msg, &head, 12);
  if((bytesRead = sendto(fd, msg, sizeof(msg), 0, (struct sockaddr*) &clientAddr, sin_size)) == -1) {
  }
  return;
}

//Function to print cout statement for server (overload for simplicity)
void printServerOut(Packet packet, std::string type, bool dup) {
  /*
  Server: "RECV" <Sequence Number> <Acknowledgement Number>
  <Connection ID> ["ACK"] ["SYN"] ["FIN"]
  Server: "RECV" <Sequence Number> <Acknowledgement Number>
  <Connection ID> ["ACK"] ["SYN"] ["FIN"]
  Server: "SEND" <Sequence Number> <Acknowledgement Number>
  <Connection ID> ["ACK"] ["SYN"] ["FIN"] ["DUP"]
  */

  std::string result = type + " " + std::to_string(packet.getSeq()) + " " + std::to_string(packet.getAck())
  + " " + std::to_string(packet.getID());
  if(packet.isACK()) result += " ACK";
  if(packet.isSYN()) result += " SYN";
  if(packet.isFIN()) result += " FIN";
  if(dup) result += " DUP";
  std::cout << result << std::endl;
  return;
}