#include "__SMTP_DATA.h"
#include "SMTPServer.h"
#include "DKIMVerifier.h"
#include "DMARCChecker.h"
#include "Log.h"
#include <cctype>
#include <chrono>
#include <cstdlib>
#include <ctime>
#include <iomanip>
#include <vector>

namespace mail {

   namespace {

      // Pulls the domain out of the message's own From: header (the
      // identity DMARC alignment actually judges) - handles both
      // "Name <user@domain>" and a bare "user@domain" value, and
      // RFC 5322 3.2.2 folding (a continuation line is joined in before
      // looking for the address).
      std::string extractFromDomain(const std::string &message) {
         size_t headerEnd = message.find("\r\n\r\n");
         std::string headerBlock = (headerEnd == std::string::npos) ? message : message.substr(0, headerEnd);

         std::vector<std::string> lines;
         size_t pos = 0;
         while (pos < headerBlock.size()) {
            size_t next = headerBlock.find("\r\n", pos);
            lines.push_back(next == std::string::npos ? headerBlock.substr(pos) : headerBlock.substr(pos, next - pos));
            if (next == std::string::npos)
               break;
            pos = next + 2;
         }

         std::string fromValue;
         for (size_t i = 0; i < lines.size(); ++i) {
            if (lines[i].size() < 5)
               continue;
            std::string prefix = lines[i].substr(0, 5);
            for (char &c : prefix)
               c = std::tolower((unsigned char)c);
            if (prefix != "from:")
               continue;
            fromValue = lines[i].substr(5);
            for (size_t j = i + 1; j < lines.size() && !lines[j].empty() && (lines[j][0] == ' ' || lines[j][0] == '\t'); ++j)
               fromValue += " " + lines[j];
            break;
         }
         if (fromValue.empty())
            return "";

         size_t lt = fromValue.find('<');
         size_t gt = fromValue.find('>');
         std::string address = (lt != std::string::npos && gt != std::string::npos && gt > lt)
           ? fromValue.substr(lt + 1, gt - lt - 1)
           : fromValue;

         size_t at = address.find_last_of('@');
         if (at == std::string::npos)
            return "";
         std::string domain = address.substr(at + 1);
         size_t start = domain.find_first_not_of(" \t");
         size_t end = domain.find_last_not_of(" \t");
         if (start == std::string::npos)
            return "";
         return domain.substr(start, end - start + 1);
      }

   }

   // RFC 5321 4.4: every SMTP server MUST prepend a Received trace
   // header to a message it processes - this codebase never wrote one
   // at all before. Deliberately no "for <recipient>" clause: one
   // already-written queue file gets hardlinked to every recipient (see
   // queueMail() below), and RFC 5321 4.4 itself says that clause
   // SHOULD be omitted once there's more than one anyway.
   coreutils::MString __SMTP_DATA::receivedHeader(SMTPSession &session, SMTPServer &server) {
      time_t now = time(nullptr);
      struct tm tmResult;
      localtime_r(&now, &tmResult);
      char dateBuf[64];
      strftime(dateBuf, sizeof(dateBuf), "%a, %d %b %Y %H:%M:%S %z", &tmResult);

      coreutils::MString header;
      header << "Received: from " << session.clientDomainName << CRLF;
      header << "\tby " << server.hostName.c_str() << " with SMTP;" << CRLF;
      header << "\t" << dateBuf << CRLF;
      return header;
   }

   int __SMTP_DATA::processCommand(coreutils::ZString &request, SMTPSession &session, SMTPServer &server) {

       switch(session.mode) {

        case WAIT_FOR_DATA:

	  switch(session.state) {

	   case RCPT:
	     session.out << "354 Enter the mail message terminated by <CRLF>.<CRLF>" << CRLF;
	     // RFC 5321 4.4 requires a Received trace header on every message
	     // this server processes - previously absent entirely. Written
	     // once here as the seed of mailData so the body lines that
	     // follow just accumulate after it; omits a "for <recipient>"
	     // clause since one already-written file gets hardlinked to every
	     // recipient (see queueMail() below) and RFC 5321 4.4 itself says
	     // that clause SHOULD be left out once there's more than one.
	     // session.spfHeader was computed once already, back at MAIL
	     // FROM time (see __SMTP_MAIL.cpp) - reused here rather than
	     // re-running the SPF DNS lookups a second time for the same
	     // transaction. Placed directly after Received:, matching
	     // where real-world MTAs conventionally put it.
	     session.mailData = receivedHeader(session, server);
	     session.mailData << session.spfHeader.c_str();
	     session.mode = RECEIVE_DATA;
	     grabInput(session);
//	     session.setTimer(120.0f);
	     break;

	   case MAIL:
	     session.out << "503 Please use RCPT first" << CRLF;
	     break;

	   default:
	     session.out << "503 Please use MAIL first" << CRLF;
	     break;

	  }
	  break;

		case RECEIVE_DATA:
		  if(request != ".") {
		     // RFC 5321 4.5.2 transparency: a line starting with '.' that
		     // isn't the bare terminator has been stuffed with an extra
		     // leading '.' by a compliant sender - strip exactly one before
		     // storing, or every such line comes out with a spurious
		     // leading dot baked permanently into the stored message body.
		     coreutils::ZString line = request.startsWith(".")
		       ? coreutils::ZString(request.getData() + 1, request.getLength() - 1)
		       : request;
		     // Cap accumulation at the advertised SIZE limit rather than
		     // growing mailData without bound - still consuming every line
		     // the client sends either way (RFC 5321 doesn't support a
		     // synchronous mid-DATA reply; the client isn't reading
		     // anything back until its own terminating "." goes out), so
		     // this only bounds our own memory use. The client actually
		     // finds out via the 552 reply below, once that terminator
		     // arrives.
		     if (session.mailData.getLength() < SMTPServer::maxMessageSize)
		       session.mailData << line << CRLF;
		  }
		  else {
		     session.mode = WAIT_FOR_DATA;
		     session.state = READY;
		     clearGrab(session);
		     if (session.mailData.getLength() >= SMTPServer::maxMessageSize) {
			session.out << "552 5.3.4 Message size exceeds fixed maximum message size" << CRLF;
		     } else if(session.recipientList.size() > 0) {
				if(session.mailData.getLength() > 0) {
				   std::string message = session.mailData.str();

				   // DKIM verification + DMARC alignment, now that the full
				   // message (including whatever From:/DKIM-Signature headers
				   // the client sent) is actually available - neither can run
				   // any earlier than this.
				   std::vector<DKIMVerifier::SignatureResult> dkimResults = DKIMVerifier::verify(message);
				   std::string fromDomain = extractFromDomain(message);
				   DMARCChecker::Policy dmarcPolicy;
				   DMARCChecker::Result dmarcResult = DMARCChecker::evaluate(
				      fromDomain, session.spfDomain, session.spfResult, dkimResults, dmarcPolicy);
				   std::string authResults = DMARCChecker::authenticationResultsHeader(
				      session.spfResult, session.spfDomain, dkimResults, dmarcResult, fromDomain, server.hostName);

				   // RFC 7489 6.3: pct= applies the published disposition to only
				   // that percentage of otherwise-failing messages - the
				   // remainder still gets the Authentication-Results header
				   // recording the real (failing) result, just isn't rejected
				   // for it. p=none/quarantine never reject outright either way
				   // (quarantine has no separate mailbox to route into here, so
				   // it degrades to "accept, but the header says fail" - still a
				   // meaningful signal for whatever reads the header downstream).
				   bool sampled = (dmarcPolicy.percent >= 100) || ((rand() % 100) < dmarcPolicy.percent);
				   bool reject = (dmarcResult == DMARCChecker::Result::Fail) && dmarcPolicy.policy == "reject" && sampled;

				   if (reject) {
				      coreutils::Log(coreutils::LOG_INFO) << "Rejected message from " << session.from
				         << " - DMARC fail, p=reject for " << fromDomain.c_str() << ".";
				      session.out << "550 5.7.1 Message rejected: failed DMARC policy for " << fromDomain.c_str() << CRLF;
				   } else {
				      coreutils::MString withAuthResults(authResults);
				      withAuthResults << session.mailData;
				      session.mailData = withAuthResults;

				      coreutils::MString ID = queueMail(server, session.from, session.recipientList, session.mailData);
				      if(ID != "") {
				         coreutils::Log(coreutils::LOG_INFO) << "Queued message " << ID << ".";
				         session.out << "250 OK Queued message " << ID << CRLF;
				      }
				      else
				        session.out << "550 Mail message too big" << CRLF;
				   }
				} else
				  session.out << "550 Mail message was empty" << CRLF;
		     } else
		       session.out << "250 OK Queued message " << CRLF;
		  }
		      break;
       }
      return 1;
   }

   coreutils::MString __SMTP_DATA::generateMailFileName() {
      struct timespec timex;
      clock_gettime(CLOCK_REALTIME, &timex);
      std::stringstream temp;
		temp << timex.tv_sec << "." << std::setfill('0') << std::setw(9) << timex.tv_nsec;
      return coreutils::MString(temp.str());
   }

   coreutils::MString __SMTP_DATA::queueMail(SMTPServer &server,
                                             coreutils::MString &sender,
                                             std::vector<coreutils::MString> &recipientList,
                                             coreutils::MString &mailData) {
      coreutils::MString fileName;
      fileName << server.mailFileSystem.getMailPath() << "/.queue/" << generateMailFileName();
      coreutils::File mailFile(fileName, O_CREAT | O_WRONLY, 0660);
      mailFile.write(mailData);

      for(coreutils::MString recipient: recipientList) {
	 coreutils::MString newName;
         // Also encode the envelope sender, not just the recipient -
         // SMTPServer::inCreate() needs the real MAIL FROM address to
         // hand outbound (non-local) mail to SendMail, and the filename
         // is the only thing that survives from this queuing step to
         // that point.
         newName << fileName << "|" << recipient << "|" << sender;
	 link(fileName.c_str(), newName.c_str());
      }

      unlink(fileName.c_str());

      return fileName;
   }

}