/**
 * ---------------------------------------------------------------------------
 * @notice This file was adapted from the Minnow networking framework,
 * originally developed for Stanford University's CS144 course.
 * It has been incorporated into the Ricochet project as foundational
 * system utility code.
 * ---------------------------------------------------------------------------
 */

#pragma once

#include "address.hh"
#include "file_descriptor.hh"

#include <functional>
#include <sys/socket.h>

//! \brief Base class for network sockets (TCP, UDP, etc.)
//! \details Socket is generally used via a subclass. See TCPSocket and UDPSocket for usage examples.
class Socket : public FileDescriptor
{
public:
  //! Bind a socket to a specified address with [bind(2)](\ref man2::bind), usually for listen/accept
  void bind( const Address& address );

  //! Bind a socket to a specified device
  void bind_to_device( std::string_view device_name );

  //! Connect a socket to a specified peer address with [connect(2)](\ref man2::connect)
  void connect( const Address& address );

  //! Shut down a socket via [shutdown(2)](\ref man2::shutdown)
  void shutdown( int how );

  //! Get local address of socket with [getsockname(2)](\ref man2::getsockname)
  Address local_address() const;
  //! Get peer address of socket with [getpeername(2)](\ref man2::getpeername)
  Address peer_address() const;

  //! Allow local address to be reused sooner via [SO_REUSEADDR](\ref man7::socket)
  void set_reuseaddr();

  //! Check for errors (will be seen on non-blocking sockets)
  void throw_if_error() const;

private:
  //! Get the local or peer address the socket is connected to
  Address get_address( const std::string& name_of_function,
                       const std::function<int( int, sockaddr*, socklen_t* )>& function ) const;

protected:
  //! Construct via [socket(2)](\ref man2::socket)
  Socket( int domain, int type, int protocol = 0 );

  //! Construct from a file descriptor.
  Socket( FileDescriptor&& fd, int domain, int type, int protocol = 0 );

  //! Wrapper around [getsockopt(2)](\ref man2::getsockopt)
  template<typename option_type>
  socklen_t getsockopt( int level, int option, option_type& option_value ) const;

  //! Wrappers around [setsockopt(2)](\ref man2::setsockopt)
  template<typename option_type>
  void setsockopt( int level, int option, const option_type& option_value );

  void setsockopt( int level, int option, std::string_view option_val );
};

class DatagramSocket : public Socket
{
public:
  // Receive a datagram and the Address of its sender, into a buffer or collection of buffers.
  void recv( Address& source_address, std::string& payload );
  void recv( Address& source_address, std::vector<std::string>& payloads );

  //! Send a datagram to specified Address (if destination is missing, sends to already-connected address)
  void send( std::string_view payload, const std::optional<Address>& destination = {} );
  void send( const StringViewRange auto&& payloads, const std::optional<Address>& destination = {} )
  {
    static thread_local std::vector<iovec> iovecs;
    const size_t total_size = to_iovecs( payloads, iovecs );
    send( iovecs, total_size, destination );
  }

private:
  void send( std::vector<iovec>& iovecs, size_t total_size, const std::optional<Address>& destination = {} );

protected:
  DatagramSocket( int domain, int type, int protocol = 0 ) : Socket( domain, type, protocol ) {}

  //! Construct from a file descriptor.
  DatagramSocket( FileDescriptor&& fd, int domain, int type, int protocol = 0 )
    : Socket( std::move( fd ), domain, type, protocol )
  {}
};

//! A wrapper around [UDP sockets](\ref man7::udp)
class UDPSocket : public DatagramSocket
{
  //! \param[in] fd is the FileDescriptor from which to construct
  explicit UDPSocket( FileDescriptor&& fd ) : DatagramSocket( std::move( fd ), AF_INET, SOCK_DGRAM ) {}

public:
  //! Default: construct an unbound, unconnected UDP socket
  UDPSocket() : DatagramSocket( AF_INET, SOCK_DGRAM ) {}
};

//! A wrapper around [TCP sockets](\ref man7::tcp)
class TCPSocket : public Socket
{
private:
  //! \brief Construct from FileDescriptor (used by accept())
  //! \param[in] fd is the FileDescriptor from which to construct
  explicit TCPSocket( FileDescriptor&& fd ) : Socket( std::move( fd ), AF_INET, SOCK_STREAM, IPPROTO_TCP ) {}

public:
  //! Default: construct an unbound, unconnected TCP socket
  TCPSocket() : Socket( AF_INET, SOCK_STREAM ) {}

  //! Mark a socket as listening for incoming connections
  void listen( int backlog = 16 );

  //! Accept a new incoming connection
  TCPSocket accept();
};

//! A wrapper around [packet sockets](\ref man7:packet)
class PacketSocket : public DatagramSocket
{
public:
  PacketSocket( const int type, const int protocol ) : DatagramSocket( AF_PACKET, type, protocol ) {}

  void set_promiscuous();
};

//! A wrapper around [raw sockets](\ref man7:raw)
class RawSocket : public DatagramSocket
{
public:
  RawSocket() : DatagramSocket( AF_INET, SOCK_RAW, IPPROTO_RAW ) {}
};

//! A wrapper around [Unix-domain stream sockets](\ref man7::unix)
class LocalStreamSocket : public Socket
{
public:
  //! Construct from a file descriptor
  explicit LocalStreamSocket( FileDescriptor&& fd ) : Socket( std::move( fd ), AF_UNIX, SOCK_STREAM ) {}
};

//! A wrapper around [Unix-domain datagram sockets](\ref man7::unix)
class LocalDatagramSocket : public DatagramSocket
{
  //! \param[in] fd is the FileDescriptor from which to construct
  explicit LocalDatagramSocket( FileDescriptor&& fd ) : DatagramSocket( std::move( fd ), AF_UNIX, SOCK_DGRAM ) {}

public:
  //! Default: construct an unbound, unconnected socket
  LocalDatagramSocket() : DatagramSocket( AF_UNIX, SOCK_DGRAM ) {}
};