#include "__IMAP_FETCH.h"
#include "IMAPArgs.h"
#include "IMAPServer.h"
#include "MessageHeaders.h"
#include <algorithm>
#include <cctype>
#include <cstdlib>
#include <ctime>
#include <sstream>

namespace mail {

   namespace {

      // Reads one FETCH data-item token, which - unlike an ordinary IMAP
      // atom - may carry a "[section]" and/or "<start.length>" suffix
      // with no space before them (BODY[HEADER.FIELDS (SUBJECT)]), and
      // may itself contain embedded spaces INSIDE that bracketed
      // section. IMAPArgs::nextArgument() can't handle this shape (it
      // only recognizes a container when the very first character is
      // one of "([{<" - "BODY[...]" starts with 'B').
      std::string nextFetchItem(coreutils::ZString &items) {
         items.skipWhitespace();
         if (items.eod())
            return "";
         std::string token;
         while (!items.eod() && items.getChar() != ' ' && items.getChar() != '[' && items.getChar() != '<')
            token += items.nextChar();
         if (!items.eod() && items.getChar() == '[') {
            coreutils::ZString section = items.getContainer();
            token += section.str();
         }
         if (!items.eod() && items.getChar() == '<') {
            token += items.nextChar();
            while (!items.eod() && items.getChar() != '>')
               token += items.nextChar();
            if (!items.eod())
               token += items.nextChar();
         }
         return token;
      }

      std::string upper(const std::string &s) {
         std::string result = s;
         for (char &c : result)
            c = std::toupper((unsigned char)c);
         return result;
      }

      std::string formatInternalDate(time_t when) {
         struct tm tmValue;
         gmtime_r(&when, &tmValue);
         char buf[64];
         strftime(buf, sizeof(buf), "%d-%b-%Y %H:%M:%S +0000", &tmValue);
         return buf;
      }

      std::string joinFlags(const std::vector<std::string> &flags) {
         std::string out;
         for (size_t i = 0; i < flags.size(); ++i) {
            if (i > 0)
               out += " ";
            out += flags[i];
         }
         return out;
      }

      // Formats a section's content as an IMAP literal ({n}\r\n<bytes>),
      // always - simpler and always correct, vs. quoted-string escaping
      // for content that can contain arbitrary bytes/CRLFs.
      std::string asLiteral(const std::string &content) {
         std::ostringstream out;
         out << "{" << content.size() << "}\r\n" << content;
         return out.str();
      }

      // Extracts the section name out of a token like
      // "BODY[HEADER.FIELDS (SUBJECT TO)]<0.100>" -> "HEADER.FIELDS (SUBJECT TO)",
      // and separately reports any <start.length>/<start> range.
      struct FetchItem {
         std::string name;     // e.g. "FLAGS", "BODY.PEEK", "RFC822.SIZE"
         std::string section;  // interior of [...], if any
         bool hasSection = false;
         bool hasRange = false;
         long rangeStart = 0;
         long rangeLength = -1; // -1 == to end
      };

      FetchItem parseFetchItem(const std::string &token) {
         FetchItem item;
         size_t bracket = token.find('[');
         size_t angle = token.find('<');
         size_t nameEnd = std::min(bracket, angle);
         item.name = (nameEnd == std::string::npos) ? token : token.substr(0, nameEnd);
         if (bracket != std::string::npos) {
            size_t closeBracket = token.find(']', bracket);
            if (closeBracket != std::string::npos) {
               item.section = token.substr(bracket + 1, closeBracket - bracket - 1);
               item.hasSection = true;
            }
         }
         if (angle != std::string::npos) {
            size_t closeAngle = token.find('>', angle);
            if (closeAngle != std::string::npos) {
               std::string rangeText = token.substr(angle + 1, closeAngle - angle - 1);
               size_t dot = rangeText.find('.');
               item.hasRange = true;
               item.rangeStart = std::strtol(rangeText.substr(0, dot).c_str(), nullptr, 10);
               if (dot != std::string::npos)
                  item.rangeLength = std::strtol(rangeText.substr(dot + 1).c_str(), nullptr, 10);
            }
         }
         return item;
      }

      std::string applyRange(const std::string &content, const FetchItem &item) {
         if (!item.hasRange)
            return content;
         if (item.rangeStart < 0 || (size_t)item.rangeStart >= content.size())
            return "";
         size_t length = item.rangeLength < 0 ? std::string::npos : (size_t)item.rangeLength;
         return content.substr(item.rangeStart, length);
      }

      std::string sectionContent(const std::string &fullMessage, const std::string &sectionSpec) {
         std::string headerBlock, bodyBlock;
         splitMessage(fullMessage, headerBlock, bodyBlock);
         std::string spec = upper(sectionSpec);
         // Strip trailing whitespace introduced by our own tokenizing.
         while (!spec.empty() && spec.back() == ' ')
            spec.pop_back();

         if (spec.empty())
            return fullMessage;
         if (spec == "TEXT")
            return bodyBlock;
         if (spec == "HEADER")
            return headerBlock + "\r\n\r\n";
         if (spec.rfind("HEADER.FIELDS", 0) == 0) {
            size_t paren = sectionSpec.find('(');
            size_t closeParen = sectionSpec.find(')', paren);
            std::string fieldsText = (paren != std::string::npos && closeParen != std::string::npos)
              ? sectionSpec.substr(paren + 1, closeParen - paren - 1) : "";
            std::istringstream fieldStream(fieldsText);
            std::string fieldName;
            auto headers = parseHeaderLines(headerBlock);
            std::string result;
            // Only fields actually present in the message are emitted -
            // matching common server behavior for a requested field the
            // message simply lacks.
            while (fieldStream >> fieldName) {
               for (auto &h : headers) {
                  if (h.first.size() != fieldName.size())
                     continue;
                  bool match = true;
                  for (size_t i = 0; i < fieldName.size(); ++i)
                     if (std::toupper((unsigned char)h.first[i]) != std::toupper((unsigned char)fieldName[i])) {
                        match = false;
                        break;
                     }
                  if (match)
                     result += h.first + ": " + h.second + "\r\n";
               }
            }
            result += "\r\n";
            return result;
         }
         // Unrecognized section shape (MIME part numbers, HEADER.FIELDS.NOT,
         // etc. - out of scope, see plan's deferred-items list) - fall
         // back to the whole message rather than fail outright.
         return fullMessage;
      }

   }

   int fetchImpl(coreutils::ZString &request, IMAPSession &session, IMAPServer &server, bool uidMode) {
      if (session.state != IMAPSession::SELECTED) {
         session.out << session.tag << " NO No mailbox selected" << CRLF;
         return 1;
      }

      std::vector<uint32_t> uids = IMAPArgs::parseSequenceSet(request, session.mailbox, uidMode);

      request.skipWhitespace();
      std::vector<std::string> itemTokens;
      if (!request.eod() && request.getChar() == '(') {
         coreutils::ZString listArg = request.getContainer();
         std::string interior = IMAPArgs::unwrap(listArg);
         coreutils::ZString interiorZ(interior);
         while (true) {
            std::string item = nextFetchItem(interiorZ);
            if (item.empty())
               break;
            itemTokens.push_back(item);
         }
      } else {
         std::string item = nextFetchItem(request);
         if (!item.empty())
            itemTokens.push_back(item);
      }

      if (itemTokens.size() == 1) {
         std::string macro = upper(itemTokens[0]);
         if (macro == "ALL")
            itemTokens = {"FLAGS", "INTERNALDATE", "RFC822.SIZE", "ENVELOPE"};
         else if (macro == "FAST")
            itemTokens = {"FLAGS", "INTERNALDATE", "RFC822.SIZE"};
         else if (macro == "FULL")
            itemTokens = {"FLAGS", "INTERNALDATE", "RFC822.SIZE", "ENVELOPE"};
      }

      for (uint32_t uid : uids) {
         MailFolder::Message *message = session.mailbox.find(uid);
         if (!message)
            continue;
         uint32_t seqNum = session.mailbox.seq(uid);

         bool needsContent = false;
         for (const std::string &token : itemTokens) {
            std::string name = upper(parseFetchItem(token).name);
            if (name == "ENVELOPE" || name == "RFC822" || name == "RFC822.HEADER" ||
                name == "RFC822.TEXT" || name.rfind("BODY", 0) == 0)
               needsContent = true;
         }
         std::string content = needsContent ? session.mailbox.readMessage(uid) : "";

         bool markSeen = false;
         std::ostringstream line;
         bool first = true;
         bool sawFlags = false;
         auto emit = [&](const std::string &text) {
            if (!first)
               line << " ";
            first = false;
            line << text;
         };

         for (const std::string &token : itemTokens) {
            FetchItem item = parseFetchItem(token);
            std::string name = upper(item.name);

            if (name == "FLAGS") {
               emit("FLAGS (" + joinFlags(message->flags) + ")");
               sawFlags = true;
            } else if (name == "UID") {
               emit("UID " + std::to_string(uid));
            } else if (name == "INTERNALDATE") {
               emit("INTERNALDATE " + imapQuoteOrNil(formatInternalDate(message->internalDate)));
            } else if (name == "RFC822.SIZE") {
               emit("RFC822.SIZE " + std::to_string(message->size));
            } else if (name == "ENVELOPE") {
               emit("ENVELOPE " + buildEnvelope(content));
            } else if (name == "RFC822") {
               emit("RFC822 " + asLiteral(applyRange(content, item)));
               if (!session.readOnly) markSeen = true;
            } else if (name == "RFC822.HEADER") {
               std::string headerBlock, bodyBlock;
               splitMessage(content, headerBlock, bodyBlock);
               emit("RFC822.HEADER " + asLiteral(applyRange(headerBlock + "\r\n\r\n", item)));
            } else if (name == "RFC822.TEXT") {
               std::string headerBlock, bodyBlock;
               splitMessage(content, headerBlock, bodyBlock);
               emit("RFC822.TEXT " + asLiteral(applyRange(bodyBlock, item)));
               if (!session.readOnly) markSeen = true;
            } else if (name == "BODY" || name == "BODY.PEEK") {
               std::string sectionText = item.hasSection ? sectionContent(content, item.section) : content;
               std::string responseKey = "BODY[" + item.section + "]";
               if (item.hasRange)
                  responseKey += "<" + std::to_string(item.rangeStart) + ">";
               emit(responseKey + " " + asLiteral(applyRange(sectionText, item)));
               bool isPeek = (name == "BODY.PEEK");
               std::string upperSection = upper(item.section);
               bool headerOnly = (upperSection == "HEADER" || upperSection.rfind("HEADER.FIELDS", 0) == 0);
               if (!isPeek && !headerOnly && !session.readOnly)
                  markSeen = true;
            } else {
               // Unrecognized item (BODYSTRUCTURE, bare BODY macro, etc. -
               // see the plan's explicitly-deferred list) - acknowledge
               // with NIL rather than silently drop the field.
               emit(upper(token) + " NIL");
            }
         }

         if (markSeen && std::find(message->flags.begin(), message->flags.end(), "\\Seen") == message->flags.end()) {
            message->flags.push_back("\\Seen");
            session.mailbox.setFlags(uid, message->flags);
            if (!sawFlags)
               emit("FLAGS (" + joinFlags(message->flags) + ")");
         }

         session.out << "* " << (int)seqNum << " FETCH (" << line.str() << ")" << CRLF;
      }

      session.out << session.tag << " OK " << (uidMode ? "UID FETCH" : "FETCH") << " completed" << CRLF;
      return 1;
   }

   int __IMAP_FETCH::processCommand(coreutils::ZString &request, IMAPSession &session, IMAPServer &server) {
      return fetchImpl(request, session, server, false);
   }

}