/*
 * Copyright ©2026 Soham Pardeshi. All rights reserved.
 * Permission is hereby granted to students registered for University of
 * Washington CSE 333 for use solely during Summer Quarter 2026 for
 * purposes of the course. No other use, copying, distribution, or
 * modification is permitted without prior written consent. Copyrights
 * for third-party components of this work must be honored. Instructors
 * interested in reusing these course materials should contact the author.
 */

#include <arpa/inet.h>
#include <errno.h>
#include <netdb.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#include <cassert>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>

// Once connected, a socket is just a file descriptor: read() and write()
// as usual (or send()/recv() for socket-specific flags).
//
//   ./sendreceive attu.cs.washington.edu 22
//
// The important part is the loops. On a socket:
//   * read() can return FEWER bytes than you asked for, and write() can
//     write fewer than you gave it -- always loop until you're done.
//   * read()/write() returning -1 with errno == EINTR is not an error;
//     a signal interrupted the call and you should retry.
//   * read() returning 0 means EOF: the peer closed the connection.
// Getting these three cases wrong is the most common networking bug.
//
// LookupName() is the same helper as in connect.cc, repeated so each
// program stands alone.
#define BUF 256

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

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

  struct sockaddr_storage addr;
  size_t addrlen;
  if (!LookupName(argv[1], port, &addr, &addrlen)) {
    Usage(argv[0]);
  }

  int socket_fd;
  if (!Connect(addr, addrlen, &socket_fd)) {
    Usage(argv[0]);
  }

  // Read from the socket, echo it back to the server, and print it,
  // until the peer closes the connection.
  char buf[BUF];
  while (true) {
    ssize_t rres = read(socket_fd, buf, BUF);
    if (rres == 0) {
      std::cerr << "peer closed the connection" << std::endl;
      break;
    }
    if (rres == -1) {
      if (errno == EINTR) continue;  // interrupted, not an error: retry
      std::cerr << "socket read failure: " << strerror(errno) << std::endl;
      close(socket_fd);
      return EXIT_FAILURE;
    }

    // Loop: a single write() is not guaranteed to consume the whole buffer.
    ssize_t written = 0;
    while (written < rres) {
      ssize_t wres = write(socket_fd, buf + written, rres - written);
      if (wres == -1) {
        if (errno == EINTR) continue;
        std::cerr << "socket write failure: " << strerror(errno) << std::endl;
        close(socket_fd);
        return EXIT_FAILURE;
      }
      written += wres;
    }

    std::cout.write(buf, rres);
    std::cout.flush();
  }

  close(socket_fd);
  return EXIT_SUCCESS;
}

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

bool LookupName(char* name, unsigned short port,
                struct sockaddr_storage* ret_addr, size_t* ret_addrlen) {
  struct addrinfo hints, *results;
  memset(&hints, 0, sizeof(hints));
  hints.ai_family = AF_UNSPEC;
  hints.ai_socktype = SOCK_STREAM;

  int retval = getaddrinfo(name, nullptr, &hints, &results);
  if (retval != 0) {
    std::cerr << "getaddrinfo failed: " << gai_strerror(retval) << std::endl;
    return false;
  }

  // getaddrinfo fills in the address but not the port.
  if (results->ai_family == AF_INET) {
    struct sockaddr_in* v4addr =
        reinterpret_cast<sockaddr_in*>(results->ai_addr);
    v4addr->sin_port = htons(port);
  } else if (results->ai_family == AF_INET6) {
    struct sockaddr_in6* v6addr =
        reinterpret_cast<sockaddr_in6*>(results->ai_addr);
    v6addr->sin6_port = htons(port);
  } else {
    std::cerr << "getaddrinfo returned neither IPv4 nor IPv6" << std::endl;
    freeaddrinfo(results);
    return false;
  }

  assert(results != nullptr);
  memcpy(ret_addr, results->ai_addr, results->ai_addrlen);
  *ret_addrlen = results->ai_addrlen;

  freeaddrinfo(results);
  return true;
}

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

  int res = connect(socket_fd,
                    reinterpret_cast<const struct sockaddr*>(&addr),
                    addrlen);
  if (res == -1) {
    std::cerr << "connect() failed: " << strerror(errno) << std::endl;
    close(socket_fd);
    return false;
  }

  *ret_fd = socket_fd;
  return true;
}