#include <arpa/inet.h>
#include <cassert>
#include <cerrno>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <netdb.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

using namespace std;

void Usage(char *progname);

bool LookupName(char *name,
                unsigned short port,
                struct sockaddr_storage *ret_addr,
                size_t *ret_addrlen);

bool Connect(const struct sockaddr_storage &addr,
             const size_t &addrlen,
             int *ret_fd);

// XXXXXXXXXXXXXXXXX YOUR NEW FUNCTION DECLARATIONS HERE

// Writes the given string to the given file descriptor.  Return true
// iff fully successful.  On an error, an unknown part of the string
// may have been already written.
bool WriteString(int fd, const string str);

// Reads bytes from the given file descriptor until a '\n' line break
// is found.  Returns the found line through the output parameter.
// The newline is part of the returned line.  If successful on finding
// a '\n', true will be returned.  False indicates failure and a
// possibly partial string has been read into the output parameter.
bool ReadLine(int fd, string *line);

int main(int argc, char *argv[]) {
  if (argc != 3) {
    Usage(argv[0]);
  }

  unsigned short port = 0;
  if (sscanf(argv[2], "%hu", &port) != 1) {
    Usage(argv[0]);
  }

  // XXXXXXXXXXXXXXXXX YOUR IMPLEMENTATION GOES HERE

  // Get an appropriate sockaddr structure.
  struct sockaddr_storage addr;
  size_t addrlen;
  if (!LookupName(argv[1], port, &addr, &addrlen)) {
    Usage(argv[0]);
  }

  // Connect to the remote host.
  int socket_fd;
  if (!Connect(addr, addrlen, &socket_fd)) {
    Usage(argv[0]);
  }

  // For each input line, write it to the socket.
  // Read in and print a returned line.
  while (!cin.eof()) {
    string line;
    getline(cin, line);

    if (line == "" && cin.eof())
      break;
    
    if (!WriteString(socket_fd, line+'\n'))
      Usage(argv[0]);

    string echoString;
    if (!ReadLine(socket_fd, &echoString))
      Usage(argv[0]);

    cout << echoString;
  }

  assert(close(socket_fd) == 0);

  return EXIT_SUCCESS;
}

// XXXXXXXXXXXXXXXXX YOUR NEW FUNCTION DEFINITIONS HERE

bool ReadLine(int fd, string *line) {
  line->clear();

  while (line->empty() || line->at(line->size()-1) != '\n') {
    char buf;
    ssize_t bytes_read = read(fd, &buf, 1);
    if (bytes_read == -1 && errno != EINTR && errno != EAGAIN) {
      perror(NULL);
      return false;
    }
    
    *line += buf;
  }

  return true;
}

bool WriteString(int fd, string str) {
  size_t written = 0;
  while (written < str.size()) {
    ssize_t res = write(fd, str.data()+written, str.size()-written);
    if (res == -1 && errno != EAGAIN && errno != EINTR) {
      perror(NULL);
      return false;
    }
    
    written += res;
  }

  return true;
}

void Usage(char *progname) {
  cerr << "usage: " << progname << " hostname port" << endl;
  exit(EXIT_FAILURE);
}

bool LookupName(char *name,
                unsigned short port,
                struct sockaddr_storage *ret_addr,
                size_t *ret_addrlen) {
  struct addrinfo hints, *results;
  int retval;

  memset(&hints, 0, sizeof(hints));
  hints.ai_family = AF_UNSPEC;
  hints.ai_socktype = SOCK_STREAM;

  // Do the lookup by invoking getaddrinfo().
  if ((retval = getaddrinfo(name, NULL, &hints, &results)) != 0) {
    cerr << "getaddrinfo failed: ";
    cerr << gai_strerror(retval) << endl;
    return false;
  }

  assert(results != NULL);

  // Set the port in the first result.
  if (results->ai_family == AF_INET) {
    struct sockaddr_in *v4addr = (struct sockaddr_in *) results->ai_addr;
    v4addr->sin_port = htons(port);
  } else if (results->ai_family == AF_INET6) {
    struct sockaddr_in6 *v6addr = (struct sockaddr_in6 *) results->ai_addr;
    v6addr->sin6_port = htons(port);
  } else {
    cerr << "getaddrinfo failed to provide an IPv4 or IPv6 address";
    cerr << endl;
    return false;
  }

  // Return the first result.
  memcpy(ret_addr, results->ai_addr, results->ai_addrlen);
  *ret_addrlen = results->ai_addrlen;

  // Clean up.
  freeaddrinfo(results);
  return true;
}

bool Connect(const struct sockaddr_storage &addr,
             const size_t &addrlen,
             int *ret_fd) {
  // Create the socket.
  int socket_fd = socket(addr.ss_family, SOCK_STREAM, 0);
  if (socket_fd == -1) {
    cerr << "socket() failed: " << strerror(errno) << endl;
    return false;
  }

  // Connect the socket to the remote host.
  int res = connect(socket_fd,
                    reinterpret_cast<const sockaddr *>(&addr),
                    addrlen);
  if (res == -1) {
    cerr << "connect() failed: " << strerror(errno) << endl;
    return false;
  }

  *ret_fd = socket_fd;
  return true;
}