#include "SendMail.h"
#include "Log.h"
#include <algorithm>
#include <map>
#include <arpa/inet.h>
#include <arpa/nameser.h>
#include <cstdlib>
#include <netdb.h>
#include <netinet/in.h>
#include <resolv.h>
#include <sstream>
#include <sys/socket.h>
#include <unistd.h>
#include <utility>

namespace mail {

   SendMail::SendMail(std::string hostName, std::string sendLogPath, std::string messageId,
                       coreutils::MString from, coreutils::MString recipient, coreutils::MString mailData,
                       std::map<std::string, DKIMSigner*> *dkimSigners)
     : hostName(hostName), sendLogPath(sendLogPath), messageId(messageId),
       from(from), recipient(recipient), mailData(mailData), dkimSigners(dkimSigners) {}

   std::string SendMail::domainOf(coreutils::MString &address) {
      std::string addr = address.str();
      size_t at = addr.find_last_of('@');
      return (at == std::string::npos) ? addr : addr.substr(at + 1);
   }

   std::vector<std::string> SendMail::lookupMX(const std::string &domain) {
      std::vector<std::pair<int, std::string>> records;
      unsigned char answer[4096];
      int len = res_query(domain.c_str(), C_IN, T_MX, answer, sizeof(answer));
      if (len > 0) {
         ns_msg handle;
         if (ns_initparse(answer, len, &handle) == 0) {
            int count = ns_msg_count(handle, ns_s_an);
            for (int i = 0; i < count; ++i) {
               ns_rr rr;
               if (ns_parserr(&handle, ns_s_an, i, &rr) == 0 && ns_rr_type(rr) == ns_t_mx) {
                  const unsigned char *rdata = ns_rr_rdata(rr);
                  int preference = (rdata[0] << 8) | rdata[1];
                  char host[NS_MAXDNAME];
                  if (dn_expand(ns_msg_base(handle), ns_msg_end(handle), rdata + 2, host, sizeof(host)) >= 0)
                     records.emplace_back(preference, std::string(host));
               }
            }
         }
      }
      std::sort(records.begin(), records.end(),
                [](const std::pair<int, std::string> &a, const std::pair<int, std::string> &b) {
                   return a.first < b.first;
                });
      std::vector<std::string> hosts;
      for (auto &r : records)
         hosts.push_back(r.second);
      // RFC 5321 5.1: a domain with no MX records is treated as a
      // single implicit MX naming that domain itself.
      if (hosts.empty())
         hosts.push_back(domain);
      return hosts;
   }

   std::string SendMail::readLine(int fd) {
      std::string line;
      char ch;
      while (::read(fd, &ch, 1) == 1) {
         if (ch == '\n')
            break;
         if (ch != '\r')
            line += ch;
      }
      return line;
   }

   void SendMail::sendLine(int fd, const std::string &line, coreutils::File &log) {
      std::string wire = line + "\r\n";
      ::write(fd, wire.c_str(), wire.length());
      log.write(">> " + line + "\n");
   }

   int SendMail::readResponse(int fd, coreutils::File &log, std::string &fullText) {
      // A multi-line SMTP reply repeats the same 3-digit code on every
      // line, with "-" (not a space) after the code on every line but
      // the last (RFC 5321 4.2.1) - keep reading until a line uses the
      // space form.
      int code = -1;
      fullText.clear();
      while (true) {
         std::string line = readLine(fd);
         log.write("<< " + line + "\n");
         if (line.length() < 4)
            return code;
         code = atoi(line.substr(0, 3).c_str());
         fullText += line + "\n";
         if (line[3] == ' ')
            break;
      }
      return code;
   }

   bool SendMail::deliverToHost(const std::string &host, coreutils::File &log) {

      struct addrinfo hints {};
      struct addrinfo *result = nullptr;
      // AF_UNSPEC let this pick a destination MX host's IPv6 address
      // whenever one existed, which - since the OS picks the LOCAL
      // source address to match the destination's address family - has
      // this server originate the connection from its own IPv6 address
      // instead of its IPv4 one. Confirmed live: every one of Gmail's 5
      // MX hosts rejected delivery with "550 5.7.25 ... does not have a
      // PTR record" for that IPv6 address, which really doesn't have
      // one (nor is it covered by this domain's own SPF record, which
      // only lists the IPv4 address - a second, independent reason
      // IPv6 delivery from here can never pass real-world receiver
      // policy). The IPv4 address DOES have a correct PTR (resolves
      // back to barant.com) and IS the address barant.com's SPF record
      // authorizes - AF_INET forces every outbound relay connection to
      // originate from it, matching what this domain's actual DNS
      // records promise. Revisit if IPv6 PTR + an ip6: SPF entry are
      // ever added for this host.
      hints.ai_family = AF_INET;
      hints.ai_socktype = SOCK_STREAM;
      if (getaddrinfo(host.c_str(), "25", &hints, &result) != 0 || !result) {
         log.write("connect: could not resolve " + host + "\n");
         return false;
      }

      int fd = -1;
      for (struct addrinfo *rp = result; rp != nullptr; rp = rp->ai_next) {
         fd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
         if (fd < 0)
            continue;
         if (::connect(fd, rp->ai_addr, rp->ai_addrlen) == 0)
            break;
         ::close(fd);
         fd = -1;
      }
      freeaddrinfo(result);

      if (fd < 0) {
         log.write("connect: unable to reach " + host + " on port 25\n");
         return false;
      }

      struct timeval timeout {60, 0};
      setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));
      setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout));

      std::string text;

      // Everything but the final QUIT lives in this lambda so a failure
      // partway through the conversation can bail out with a plain
      // `return false` while the caller still always sends QUIT and
      // closes the socket afterward, regardless of how this ended.
      auto converse = [&]() -> bool {

         if (readResponse(fd, log, text) / 100 != 2)
            return false;

         sendLine(fd, "EHLO " + hostName, log);
         if (readResponse(fd, log, text) / 100 != 2) {
            // Some older servers only speak HELO.
            sendLine(fd, "HELO " + hostName, log);
            if (readResponse(fd, log, text) / 100 != 2)
               return false;
         }

         sendLine(fd, "MAIL FROM:<" + from.str() + ">", log);
         if (readResponse(fd, log, text) / 100 != 2)
            return false;

         sendLine(fd, "RCPT TO:<" + recipient.str() + ">", log);
         if (readResponse(fd, log, text) / 100 != 2)
            return false;

         sendLine(fd, "DATA", log);
         if (readResponse(fd, log, text) != 354)
            return false;

         std::string body = mailData.str();
         std::istringstream stream(body);
         std::string bodyLine, wire;
         while (std::getline(stream, bodyLine)) {
            if (!bodyLine.empty() && bodyLine.back() == '\r')
               bodyLine.pop_back();
            // RFC 5321 4.5.2 transparency: a body line that legitimately
            // starts with '.' must be stuffed with an extra leading '.'
            // so the receiver never mistakes it for the terminating line.
            if (!bodyLine.empty() && bodyLine[0] == '.')
               wire += ".";
            wire += bodyLine;
            wire += "\r\n";
         }
         wire += ".\r\n";
         ::write(fd, wire.c_str(), wire.length());
         log.write(">> [message body, " + std::to_string(body.length()) + " bytes]\n");

         return readResponse(fd, log, text) / 100 == 2;
      };

      bool ok = converse();

      sendLine(fd, "QUIT", log);
      readResponse(fd, log, text);
      ::close(fd);

      return ok;
   }

   bool SendMail::send() {

      // ZString's std::string constructor takes a non-const reference,
      // which can't bind to a temporary - a named local is required
      // here rather than passing the concatenation expression directly.
      std::string logFileName = sendLogPath + "/" + messageId + ".log";
      coreutils::File log(logFileName, O_CREAT | O_WRONLY | O_APPEND, 0660);
      log.write("=== delivery attempt: message " + messageId + " -> " + recipient.str() + " ===\n");

      if (dkimSigners) {
         // Signed once here, not inside deliverToHost() - deliverToHost()
         // can run again against a second MX host if the first fails, and
         // the signature (over this exact message content) is the same
         // either way, so signing once upfront avoids redundant RSA ops
         // and keeps every delivery attempt for this message consistent.
         // Which signer (if any) is used depends on the sender's own
         // domain, not a single server-wide identity - a real multi-
         // domain deployment must never sign one domain's mail with a
         // different domain's key (DMARC alignment would fail it
         // anyway, and it would leak which private key production
         // uses for an unrelated domain).
         std::string fromDomain = domainOf(from);
         auto it = dkimSigners->find(fromDomain);
         if (it != dkimSigners->end()) {
            std::string signature = it->second->sign(mailData.str());
            coreutils::MString signedData;
            signedData << signature.c_str() << mailData;
            mailData = signedData;
            log.write("DKIM-signed (domain: " + fromDomain + ").\n");
         } else {
            log.write("No DKIM key configured for domain: " + fromDomain + " - sending unsigned.\n");
         }
      }

      std::string domain = domainOf(recipient);
      std::vector<std::string> hosts = lookupMX(domain);

      bool delivered = false;
      for (auto &host : hosts) {
         log.write("--- trying " + host + " ---\n");
         if (deliverToHost(host, log)) {
            log.write("delivered successfully via " + host + "\n");
            delivered = true;
            break;
         }
      }
      if (!delivered)
         log.write("all hosts for " + domain + " failed\n");

      return delivered;
   }

}