#include "imap.hpp"
#include "main.hpp"
#include <stdarg.h>


#define SEQ "XX"
#define SEQ_LEN (sizeof(SEQ)-1)


class FreeBuf {
    public:
        char* buf;
        FreeBuf(): buf(NULL) {}
        FreeBuf(char* b): buf(b) {}
        ~FreeBuf() { free(buf); }
};


// https://tools.ietf.org/html/rfc3501#section-4.3
#define NOT_QUOTABLE(c) (NEVER_QUOTABLE(c) || (c) == '"' || (c) == '\\')
#ifndef NO_8BIT_LITERALS
#define NEVER_QUOTABLE(c) (!(c) || (c) == '\r' || (c) == '\n')
#else
#define NEVER_QUOTABLE(c) (!(c) || ((c) & 0x80u) || (c) == '\r' || (c) == '\n')
#endif

bool IMAP::is_quotable(const char* s) {
    const unsigned char* p = (const unsigned char*)s;
    while (*p) {
        if (NOT_QUOTABLE(*p)) return false;
        ++p;
    }
    return true;
}

bool IMAP::is_quotable(const char* s, size_t l) {
    const unsigned char* p = (const unsigned char*)s;
    while (l--) {
        if (!*p) return false;
        if (NOT_QUOTABLE(*p)) return false;
        ++p;
    }
    return true;
}

char* IMAP::do_quote(const char* s, size_t l) {
    size_t num = 0;
    for (size_t i=0; i<l; ++i) {
        if (!NOT_QUOTABLE(s[i])) continue;
        if (NEVER_QUOTABLE(s[i])) return NULL;
        num++;
    }
    if (!num) return strndup(s, l);
    char* rv = (char*)malloc(l+num+1);
    char* r = rv;
    while (l--) {
        if (NOT_QUOTABLE(*s)) {
            *(r++) = '\\';
        }
        *(r++) = *(s++);
    }
    *r = '\0';
    return rv;
}


/*
 * https://tools.ietf.org/html/rfc3501#section-6.4.4
 * CC <string>
 * FROM <string>
 * HEADER <field-name> <string> // e.g. Reply-To
 * NOT <search-key>
 * SEEN
 * SUBJECT <string>
 * TO <string>
// OR <search-key1> <search-key2>
 */
bool IMAP::check_match(const char*& m, unsigned& num, int expect) { // XXX: recursion
    enum { CMD=0, STR, FIELD };
    bool cmd = false;
    if (!strncmp(m, "CC ", 3) || !strncmp(m, "FROM ", 5) || !strncmp(m, "TO ", 3) || !strncmp(m, "SUBJECT ", 8)) {
        if (expect == CMD) {
            m = strchr(m, ' ')+1;
            if (!check_match(m, num, STR)) return false;
            cmd = true;
        } else if (expect == FIELD) {
            m = strchr(m, ' ');
            return true;
        } else {
            return false;
        }
    } else if (!strncmp(m, "NOT ", 4)) {
        if (expect != CMD) return false;
        m += 4;
        if (!check_match(m, num, CMD)) return false;
        cmd = true;
    } else if (!strncmp(m, "SEEN", 4)) {
        if (expect != CMD) return false;
        m += 4;
        cmd = true;
    } else if (!strncmp(m, "HEADER ", 7)) {
        if (expect != CMD) return false;
        m += 7;
        if (!check_match(m, num, FIELD)) return false;
        if (*(m++) != ' ') return false;
        if (!check_match(m, num, STR)) return false;
        cmd = true;
    }
    if (cmd) {
        num++;
        if (!*m) return true;
        if (*(m++) != ' ') return false;
        return check_match(m, num, CMD);
    } else if (expect == CMD) {
        return false;
    } else if (expect == FIELD) {
        const char* s = m;
        while ((*m >= 'a' && *m <= 'z') || (*m >= 'A' && *m <= 'Z') || (*m >= '0' && *m <= '9') || *m == '-' || *m == '_') ++m;
        if (s == m) return false; // empty
        return true; // m now hopefully at x20
    } else {
        assert(expect == STR);
        const char* s;
        if (*m == '"') {
            s = ++m;
            while (*m && *m != '"') ++m;
            if (*m != '"') return false;
            if (!is_quotable(s, m-s)) return false;
            m++;
        } else {
            s = m;
            while (*m && *m != ' ') ++m;
            if (!is_quotable(s, m-s)) return false;
        }
        return true;
    }
}


ssize_t IMAP::pop_line() {
    if (bufoff) {
        if (buflen) memmove(buf, buf+bufoff, buflen);
        bufoff = 0;
    }

    while (true) {
        if (buflen) {
            char* nl = (char*)memchr(buf, '\n', buflen);
            if (nl) {
                *nl = '\0';
#ifdef DEBUG
                LOG("< %s", buf);
#endif
                bufoff = nl-buf+1;
                buflen -= bufoff;
                return bufoff-1;
            }
        }

        if (buflen == sizeof(buf)) {
            LOG("cannot read response line, limit is %luKB", sizeof(buf)>>10);
            return -1;
        }
        ssize_t len = ssl->ssl_read(buf + buflen, sizeof(buf) - buflen);
        if (len <= 0) return -1;
        buflen += (size_t)len;
    }
}


bool IMAP::result_callback(const char* b, size_t l, void* ctx) {
    if (l < SEQ_LEN+3 || strncmp(b, SEQ " OK", SEQ_LEN+3) != 0) return false;
    return ctx? (strstr(b, (const char*)ctx) != NULL): true;
}


bool IMAP::response_callback(const char* b, size_t l, void* ctx) {
    if (l < 2 || strncmp(b, "* ", 2) != 0) return false;
    b+=2; l-=2;
    while (*b == ' ') {
        ++b; --l;
    }

    char** p = (char**)ctx;
    if (*p) free(*p);
    *p = strdup(b);

    while ((*p)[l-1] == '\n' || (*p)[l-1] == '\r' || (*p)[l-1] == ' ') {
        (*p)[(l--)-1] = '\0';
    }
    return true;
}


bool IMAP::response_print_callback(const char* b, size_t l, void*) {
    char* r = NULL;
    if (!response_callback(b, l, &r)) {
        return false;
    }
    puts(r);
    free(r);
    return true;
}


bool IMAP::send_cmd(line_cb_t response_cb, void* response_cb_ctx, line_cb_t result_cb, void* result_cb_ctx, const char* fmt, ...) {
    static char cmd[4096];

    va_list args;
    va_start(args, fmt);
    int l = vsnprintf(cmd+3, sizeof(cmd)-3-2, fmt, args);
    va_end (args);
    if (l < 0 || l >= (int)sizeof(cmd)-3-2) {
        LOG("cannot send command, limit is %luKB", sizeof(cmd)>>10);
        return false;
    }
    memcpy(cmd, SEQ " ", SEQ_LEN+1);
    l += SEQ_LEN+1;

    if (strchr(cmd, '\n') != NULL || strchr(cmd, '\r') != NULL) return false;
    cmd[l++] = '\r';
    cmd[l++] = '\n';

#ifdef DEBUG
    LOG("> %.*s", l-2, cmd);
#endif
    if (!ssl->ssl_write(cmd, l)) {
        return false;
    }

    do {
        ssize_t rv = pop_line();
        if (rv <= 0) return false;
        if (rv >= 2 && strncmp(buf, "* ", 2) == 0) {
            if (response_cb) {
                if (!response_cb(buf, (size_t)rv, response_cb_ctx)) return false;
            }
            continue;
        }
        if ((unsigned)rv >= SEQ_LEN+3 && strncmp(buf, SEQ " OK", SEQ_LEN+3) == 0) {
            if (result_cb) {
                return result_cb(buf, (size_t)rv, result_cb_ctx);
            }
            return true;
        }
        return false;
    } while (true);
}


IMAP::IMAP(SSLConn* s): ssl(s), bufoff(0), buflen(0) {
    assert(ssl);
}


IMAP::~IMAP() {
    if (send_cmd(NULL, NULL, NULL, NULL, "LOGOUT")) {
        LOG("logged out.");
    }
    delete ssl;
}


bool IMAP::init(bool interactive, const char* user, const char* pass) {
    ssize_t rv = pop_line();
    if (rv <= 0) return false;
    if (rv < 5 || strncmp(buf, "* OK ", 5) != 0) {
        LOG("cannot read IMAP OK");
        return false;
    }

    FreeBuf rbuf;
    if (!send_cmd(&response_callback, &rbuf.buf, NULL, NULL, "CAPABILITY")) {
        LOG("cannot fetch server capabilities");
        return false;
    } else if (!rbuf.buf || !strstr(rbuf.buf, " MOVE") || !strstr(rbuf.buf, " AUTH=PLAIN")) {
        LOG("plain auth and move capabilities not in: %s", rbuf.buf?:"-");
        return false;
    }

    if (*pass) {
        pass = do_quote(pass, strlen(pass));
    } else if (!interactive) {
        LOG("no password given for user %s", user);
        return false;
    } else {
        fprintf(stdout, "password for user %s? ", user);
        static char pwbuf[64];
        if (fgets(pwbuf, sizeof(pwbuf), stdin)) {
            *strchrnul(pwbuf, '\n') = '\0';
            pass = do_quote(pwbuf, strlen(pwbuf));
            memset(pwbuf, '\0', sizeof(pwbuf));
        } else {
            LOG("cannot read password");
            return false;
        }
    }

    if (!pass || !*pass) {
        LOG("invalid password");
        if (pass) {
            memset((void*)pass, '\0', strlen(pass));
            free((void*)pass);
        }
        return false;
    }

    if (!send_cmd(NULL, NULL, NULL, NULL, "LOGIN \"%s\" \"%s\"", user, pass)) {
        memset((void*)pass, '\0', strlen(pass));
        free((void*)pass);
        LOG("authentication failure");
        return false;
    }
    memset((void*)pass, '\0', strlen(pass));
    free((void*)pass);

    if (interactive) {
        (void)send_cmd(&response_print_callback, NULL, NULL, NULL, "LIST \"\" \"%\"");
    }

    return true;
}


bool IMAP::handle(bool interactive, const char* from, const char* search, const char* to) {
    assert(is_quotable(from));
    const char* rw = "[READ-WRITE]";
    if (!send_cmd(NULL, NULL, &result_callback, (void*)rw, "SELECT \"%s\"", from)) {
        LOG("cannot select mailbox %s", from);
        return false;
    }

    FreeBuf sbuf;
    if (!send_cmd(response_callback, &sbuf.buf, NULL, NULL, "UID SEARCH %s", search)) {
        LOG("cannot search in %s for %s", from, search);
        return false;
    }
    if (!sbuf.buf || strncmp(sbuf.buf, "SEARCH", 6) != 0) {
        LOG("no search result in %s for %s", from, search);
        return false;
    }
    if (!sbuf.buf[6]) {  // empty search result
        LOG("no search results in %s for %s", from, to);
        return true;
    }
    if (sbuf.buf[6] != ' ') {
        LOG("invalid search result in %s for %s", from, search);
        return false;
    }

    // replace spaces with colons for MOVE/sequence-set
    // https://tools.ietf.org/html/rfc6851#section-3 https://tools.ietf.org/html/rfc3501#section-9
    unsigned num = 0;
    char* p = sbuf.buf+6;
    do {
        ++num;
        *p = ',';
        p = strchr(p+1, ' ');
    } while (p);

    if (interactive) {
        fprintf(stdout, "move %u matching mails from %s to %s? ", num, from, to);
        static char a[3];
        if (!fgets(a, sizeof(a), stdin) || *a != 'y') {
            LOG("aborted due to user request");
            return true;
        }
    }

    assert(is_quotable(to));
    if (!send_cmd(NULL, NULL, NULL, NULL, "UID MOVE %s \"%s\"", sbuf.buf+7, to)) {
        LOG("cannot move %u mails from %s to %s", num, from, to);
        return false;
    }
    LOG("moved %u mails from %s to %s", num, from, to);
    return true;
}