#include "__IMAP_SEARCH.h"
#include "IMAPArgs.h"
#include "IMAPServer.h"
#include "MessageHeaders.h"
#include <algorithm>
#include <cctype>
#include <cstdlib>
#include <ctime>
#include <functional>
#include <memory>

namespace mail {

   namespace {

      using Predicate = std::function<bool(MailFolder &, const MailFolder::Message &)>;

      bool hasFlag(const MailFolder::Message &m, const char *flag) {
         return std::find(m.flags.begin(), m.flags.end(), flag) != m.flags.end();
      }

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

      bool containsCI(const std::string &haystack, const std::string &needle) {
         if (needle.empty())
            return true;
         return upperOf(haystack).find(upperOf(needle)) != std::string::npos;
      }

      time_t midnightUTC(time_t t) {
         struct tm tmValue;
         gmtime_r(&t, &tmValue);
         tmValue.tm_hour = 0;
         tmValue.tm_min = 0;
         tmValue.tm_sec = 0;
         return timegm(&tmValue);
      }

      Predicate headerContains(const std::string &headerName, const std::string &needle) {
         return [headerName, needle](MailFolder &folder, const MailFolder::Message &m) {
            std::string content = folder.readMessage(m.uid);
            std::string headerBlock, bodyBlock;
            splitMessage(content, headerBlock, bodyBlock);
            auto headers = parseHeaderLines(headerBlock);
            return containsCI(getHeader(headers, headerName), needle);
         };
      }

      Predicate parseOneKey(coreutils::ZString &request, MailFolder &folder) {
         coreutils::ZString keyArg = IMAPArgs::nextArgument(request);
         std::string keyText = keyArg.str();
         if (keyText.empty())
            return [](MailFolder &, const MailFolder::Message &) { return true; };

         if (std::isdigit((unsigned char)keyText[0]) || keyText[0] == '*') {
            coreutils::ZString setOnly(keyText);
            auto uids = std::make_shared<std::vector<uint32_t>>(IMAPArgs::parseSequenceSet(setOnly, folder, false));
            return [uids](MailFolder &, const MailFolder::Message &m) {
               return std::find(uids->begin(), uids->end(), m.uid) != uids->end();
            };
         }

         std::string key = upperOf(keyText);

         if (key == "ALL")
            return [](MailFolder &, const MailFolder::Message &) { return true; };
         if (key == "ANSWERED")
            return [](MailFolder &, const MailFolder::Message &m) { return hasFlag(m, "\\Answered"); };
         if (key == "UNANSWERED")
            return [](MailFolder &, const MailFolder::Message &m) { return !hasFlag(m, "\\Answered"); };
         if (key == "DELETED")
            return [](MailFolder &, const MailFolder::Message &m) { return hasFlag(m, "\\Deleted"); };
         if (key == "UNDELETED")
            return [](MailFolder &, const MailFolder::Message &m) { return !hasFlag(m, "\\Deleted"); };
         if (key == "SEEN")
            return [](MailFolder &, const MailFolder::Message &m) { return hasFlag(m, "\\Seen"); };
         if (key == "UNSEEN")
            return [](MailFolder &, const MailFolder::Message &m) { return !hasFlag(m, "\\Seen"); };
         if (key == "FLAGGED")
            return [](MailFolder &, const MailFolder::Message &m) { return hasFlag(m, "\\Flagged"); };
         if (key == "UNFLAGGED")
            return [](MailFolder &, const MailFolder::Message &m) { return !hasFlag(m, "\\Flagged"); };
         if (key == "DRAFT")
            return [](MailFolder &, const MailFolder::Message &m) { return hasFlag(m, "\\Draft"); };
         if (key == "UNDRAFT")
            return [](MailFolder &, const MailFolder::Message &m) { return !hasFlag(m, "\\Draft"); };
         if (key == "NEW")
            return [](MailFolder &, const MailFolder::Message &m) { return m.recent && !hasFlag(m, "\\Seen"); };
         if (key == "OLD")
            return [](MailFolder &, const MailFolder::Message &m) { return !m.recent; };
         if (key == "RECENT")
            return [](MailFolder &, const MailFolder::Message &m) { return m.recent; };

         if (key == "UID") {
            coreutils::ZString setArg = IMAPArgs::nextArgument(request);
            std::string setText = setArg.str();
            coreutils::ZString setOnly(setText);
            auto uids = std::make_shared<std::vector<uint32_t>>(IMAPArgs::parseSequenceSet(setOnly, folder, true));
            return [uids](MailFolder &, const MailFolder::Message &m) {
               return std::find(uids->begin(), uids->end(), m.uid) != uids->end();
            };
         }

         if (key == "FROM" || key == "TO" || key == "CC" || key == "BCC" || key == "SUBJECT") {
            coreutils::ZString valueArg = IMAPArgs::nextArgument(request);
            std::string needle = IMAPArgs::unwrap(valueArg);
            std::string headerName = key == "FROM" ? "From" : key == "TO" ? "To" :
                                      key == "CC" ? "Cc" : key == "BCC" ? "Bcc" : "Subject";
            return headerContains(headerName, needle);
         }

         if (key == "BODY") {
            coreutils::ZString valueArg = IMAPArgs::nextArgument(request);
            std::string needle = IMAPArgs::unwrap(valueArg);
            return [needle](MailFolder &folder, const MailFolder::Message &m) {
               std::string content = folder.readMessage(m.uid);
               std::string headerBlock, bodyBlock;
               splitMessage(content, headerBlock, bodyBlock);
               return containsCI(bodyBlock, needle);
            };
         }

         if (key == "TEXT") {
            coreutils::ZString valueArg = IMAPArgs::nextArgument(request);
            std::string needle = IMAPArgs::unwrap(valueArg);
            return [needle](MailFolder &folder, const MailFolder::Message &m) {
               return containsCI(folder.readMessage(m.uid), needle);
            };
         }

         if (key == "BEFORE" || key == "ON" || key == "SINCE") {
            coreutils::ZString dateArg = IMAPArgs::nextArgument(request);
            time_t target = IMAPArgs::parseDate(dateArg);
            if (target == (time_t)-1)
               return [](MailFolder &, const MailFolder::Message &) { return true; };
            target = midnightUTC(target);
            if (key == "BEFORE")
               return [target](MailFolder &, const MailFolder::Message &m) { return midnightUTC(m.internalDate) < target; };
            if (key == "SINCE")
               return [target](MailFolder &, const MailFolder::Message &m) { return midnightUTC(m.internalDate) >= target; };
            return [target](MailFolder &, const MailFolder::Message &m) { return midnightUTC(m.internalDate) == target; };
         }

         if (key == "LARGER" || key == "SMALLER") {
            coreutils::ZString numArg = IMAPArgs::nextArgument(request);
            long value = std::strtol(numArg.str().c_str(), nullptr, 10);
            if (key == "LARGER")
               return [value](MailFolder &, const MailFolder::Message &m) { return (long)m.size > value; };
            return [value](MailFolder &, const MailFolder::Message &m) { return (long)m.size < value; };
         }

         if (key == "NOT") {
            Predicate inner = parseOneKey(request, folder);
            return [inner](MailFolder &folder, const MailFolder::Message &m) { return !inner(folder, m); };
         }

         if (key == "OR") {
            Predicate left = parseOneKey(request, folder);
            Predicate right = parseOneKey(request, folder);
            return [left, right](MailFolder &folder, const MailFolder::Message &m) {
               return left(folder, m) || right(folder, m);
            };
         }

         // Unrecognized keyword (CHARSET, HEADER, KEYWORD, SENT*,
         // SEQUENCE numbers already handled above - see the plan's
         // deferred-items list) - treated as always-true rather than
         // aborting the whole SEARCH.
         return [](MailFolder &, const MailFolder::Message &) { return true; };
      }

   }

   int searchImpl(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<Predicate> predicates;
      while (true) {
         request.skipWhitespace();
         if (request.eod())
            break;
         predicates.push_back(parseOneKey(request, session.mailbox));
      }

      session.out << "* SEARCH";
      for (const MailFolder::Message &message : session.mailbox.messages()) {
         bool matched = true;
         for (const Predicate &predicate : predicates) {
            if (!predicate(session.mailbox, message)) {
               matched = false;
               break;
            }
         }
         if (matched)
            session.out << " " << (int)(uidMode ? message.uid : session.mailbox.seq(message.uid));
      }
      session.out << CRLF;
      session.out << session.tag << " OK " << (uidMode ? "UID SEARCH" : "SEARCH") << " completed" << CRLF;
      return 1;
   }

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

}