class ImapBuffer {
    U8* rx;
    U8* tx;
};

class ImapClient {
    CTask* mem_task;
    TlsSocket* s;
    JsonObject* o;
    ImapBuffer buf;
    U8 response_text[1024];
    U8 selected[256];
    I64 state;
    I64 tag;
    U0 (*close)();
    U0 (*connect)(U8* host, I64 port);
    U0 (*fetch)(I64 uid);
    JsonArray* (*fetch_array)(U8* sequence_set, U8* data_items);
    U0 (*select)(U8* mailbox);
    U0 (*login)(U8* userid, U8* password);
    U0 (*logout)();
};

#define IMAP_BUFFER_SIZE 65536

#define IMAP_LINE_EQ 1
#define IMAP_LINE_IQ 2

#define IMAP_STATE_NOT_CONNECTED 0
#define IMAP_STATE_NOT_AUTHENTICATED 1
#define IMAP_STATE_AUTHENTICATED 2
#define IMAP_STATE_SELECTED 3
#define IMAP_STATE_LOGOUT 128

#define IMAP_STATE_ERROR 255

#define IMAP_WRAPPER_MAGIC_NUMBER 0xCAFEFACADEDEFACE

JsonArray* @imap_json_array_from_param_list(ImapClient* c, U8* line)
{
    I64 i = 0;
    I64 state = 0;
    JsonArray* a = NULL;
    U8* out = CAlloc(StrLen(line) * 2, c->mem_task);
    while (line[i]) {
        switch (state) {
        case 0:
            if (line[i] < ' ') {
                ++i;
                goto @imap_convert_process_next_ch;
            }
            if (line[i] == ' ') {
                String.Append(out, ",");
                ++i;
                goto @imap_convert_process_next_ch;
            }
            if (line[i] == '(') {
                String.Append(out, "[");
                ++i;
                goto @imap_convert_process_next_ch;
            }
            if (line[i] == ')') {
                String.Append(out, "]");
                ++i;
                goto @imap_convert_process_next_ch;
            }
            if (line[i] == '"') {
                state = IMAP_LINE_EQ;
                String.Append(out, "%c", line[i]);
                ++i;
                goto @imap_convert_process_next_ch;
            }
            state = IMAP_LINE_IQ;
            String.Append(out, "\"");
            String.Append(out, "%c", line[i]);
            ++i;
            break;
        case IMAP_LINE_EQ:
            if (line[i] == '"') {
                state = 0;
            }
            String.Append(out, "%c", line[i]);
            ++i;
            break;
        case IMAP_LINE_IQ:
            if (line[i] == ')' || line[i] == ' ') {
                state = 0;
                String.Append(out, "\"");
                goto @imap_convert_process_next_ch;
            }
            String.Append(out, "%c", line[i]);
            ++i;
            break;
        default:
            // We shouldn't be here! Return empty Array
            StrCpy(out, "[]");
            i = StrLen(line);
            break;
        }
        @imap_convert_process_next_ch:
    }
    a = Json.Parse(out, c->mem_task);
    Free(out);
    return a;
}

U8** @imap_lineate_buffer(ImapClient* c, U8* line_cnt_out)
{
    if (!c || !line_cnt_out)
        return NULL;
    I64 line_cnt = StrOcc(c->buf.rx, '\n');
    *line_cnt_out = line_cnt;
    if (!line_cnt)
        return NULL;
    I64 i = 0;
    U8* p = c->buf.rx;
    U8** lines = CAlloc(sizeof(U8*) * line_cnt, c->mem_task);
    lines[i++] = p++;
    while (*p) {
        switch (*p) {
        case '\r':
            *p = NULL;
            p += 2;
            lines[i++] = p++;
            break;
        default:
            ++p;
            break;
        }
    }
    return lines;
}

U0 @imap_init_buffers(ImapClient* c)
{
    c->buf.rx = CAlloc(IMAP_BUFFER_SIZE, c->mem_task);
    c->buf.tx = CAlloc(IMAP_BUFFER_SIZE, c->mem_task);
}

U0 @imap_free_buffers(ImapClient* c)
{
    if (c->buf.rx)
        Free(c->buf.rx);
    if (c->buf.tx)
        Free(c->buf.tx);
}

U0 @imap_send(ImapClient* c)
{
    c->s->send(c->buf.tx, StrLen(c->buf.tx));
}

U0 @imap_set_response_text(ImapClient* c, U8* s)
{
    // Trim whitespace at both ends
    Bool trim = TRUE;
    while (trim) {
        switch (*s) {
        case ' ':
        case '\t':
        case '\r':
        case '\n':
            ++s;
            break;
        default:
            trim = FALSE;
            break;
        }
    }
    StrCpy(c->response_text, s);
    s = c->response_text + StrLen(c->response_text) - 1;
    trim = TRUE;
    while (trim) {
        switch (*s) {
        case ' ':
        case '\t':
        case '\r':
        case '\n':
            *s = NULL;
            --s;
            break;
        default:
            trim = FALSE;
            break;
        }
    }
}

U0 @imap_error(ImapClient* c, U8* s)
{
    c->state = IMAP_STATE_ERROR;
    @imap_set_response_text(c, s);
}

U0 @imap_receive(ImapClient* c)
{
    U8 buf[16];
    I64 cnt = 0;

    U8* msg;

    while (c->s->state != TCP_SOCKET_STATE_CLOSED) {
        cnt += c->s->receive(c->buf.rx + cnt, IMAP_BUFFER_SIZE);

        if (StrFind("* OK", c->buf.rx)) {
            @imap_set_response_text(c, StrFind("* OK", c->buf.rx) + 4);
            return;
        }
        StrPrint(buf, "A%03d OK", c->tag);
        if (StrFind(buf, c->buf.rx)) {
            @imap_set_response_text(c, StrFind(buf, c->buf.rx) + 7);
            return;
        }

        if (StrFind("* NO", c->buf.rx)) {
            @imap_error(c, StrFind("* NO", c->buf.rx) + 4);
            return;
        }
        StrPrint(buf, "A%03d NO", c->tag);
        if (StrFind(buf, c->buf.rx)) {
            @imap_error(c, StrFind(buf, c->buf.rx) + 7);
            return;
        }

        if (StrFind("* BAD", c->buf.rx)) {
            @imap_error(c, StrFind("* BAD", c->buf.rx) + 5);
            return;
        }
        StrPrint(buf, "A%03d BAD", c->tag);
        if (StrFind(buf, c->buf.rx)) {
            @imap_error(c, StrFind(buf, c->buf.rx) + 8);
            return;
        }
    }
}

U0 @imap_logout(ImapClient* c)
{
    StrPrint(c->buf.tx, "A%03d LOGOUT\r\n", ++c->tag);
    @imap_send(c);

    // Receive data
    @imap_receive(c);

    c->state = IMAP_STATE_LOGOUT;
}

U0 @imap_close(ImapClient* c)
{
    c->s->close();
    @imap_free_buffers(c);
    Free(c->s);

    Free(c->close);
    Free(c->connect);
    Free(c->fetch);
    Free(c->fetch_array);
    Free(c->select);
    Free(c->login);
    Free(c->logout);
    Free(c);
}

U0 @imap_login(ImapClient* c, U8* userid, U8* password)
{
    StrPrint(c->buf.tx, "A%03d LOGIN %s %s\r\n", ++c->tag, userid, password);
    @imap_send(c);

    // Receive data
    @imap_receive(c);
    if (c->state == IMAP_STATE_ERROR)
        return;

    c->state = IMAP_STATE_AUTHENTICATED;
}

JsonArray* @imap_fetch_array(ImapClient* c, U8* sequence_set, U8* data_items)
{
    StrPrint(c->buf.tx, "A%03d FETCH %s %s\r\n", ++c->tag, sequence_set, data_items);
    @imap_send(c);

    // Receive data
    @imap_receive(c);
    if (c->state == IMAP_STATE_ERROR)
        return Json.Parse("[]", c->mem_task);

    // Split buffer into lines
    I64 i = 0;
    I64 lines_cnt = 0;
    U8** lines = @imap_lineate_buffer(c, &lines_cnt);
    U8* param_list = NULL;

    if (!lines_cnt) {
        Free(lines);
        return Json.Parse("[]", c->mem_task);
    }

    // Convert each param_list to a JSON array
    JsonArray* a = Json.CreateArray(c->mem_task);
    for (i = 0; i < lines_cnt; i++) {
        param_list = StrFind(" FETCH (", lines[i]);
        if (param_list) {
            a->append(@imap_json_array_from_param_list(c, param_list + 7), JSON_ARRAY);
        }
    }

    // Return array of arrays
    return a;
}

U0 @imap_fetch(ImapClient* c, I64 uid)
{
    StrPrint(c->buf.tx, "A%03d UID FETCH %d BODY.PEEK[]\r\n", ++c->tag, uid);
    @imap_send(c);

    @imap_receive(c);
}

U0 @imap_select(ImapClient* c, U8* mailbox)
{
    StrPrint(c->buf.tx, "A%03d SELECT %s\r\n", ++c->tag, mailbox);
    @imap_send(c);

    @imap_receive(c);
    if (c->state == IMAP_STATE_ERROR)
        return;

    StrCpy(&c->selected, mailbox);
    c->state = IMAP_STATE_SELECTED;
}

U0 @imap_connect(ImapClient* c, U8* host, I64 port)
{
    if (!host || !port || !StrLen(host))
        return;

    U32 addr = @dns_query(host);
    if (addr == U32_MAX)
        return;

    TlsSocket* s = NULL;

    switch (port) {
    case 143:
        s = @tcp_socket_create(host, port);
        break;
    case 993:
        s = @tls_socket_create(host, port);
        break;
    default:
        return;
    }

    if (!s)
        return;

    c->s = s;

    // Receive data
    @imap_receive(c);
    if (c->state == IMAP_STATE_ERROR)
        return;

    c->state = IMAP_STATE_NOT_AUTHENTICATED;
}

U0 @imap_close_wrapper_function()
{
    @imap_close(IMAP_WRAPPER_MAGIC_NUMBER);
}

U0 @imap_connect_wrapper_function(U8* host, I64 port)
{
    @imap_connect(IMAP_WRAPPER_MAGIC_NUMBER, host, port);
}

U0 @imap_select_wrapper_function(U8* mailbox)
{
    @imap_select(IMAP_WRAPPER_MAGIC_NUMBER, mailbox);
}

U0 @imap_login_wrapper_function(U8* userid, U8* password)
{
    @imap_login(IMAP_WRAPPER_MAGIC_NUMBER, userid, password);
}

U0 @imap_logout_wrapper_function()
{
    @imap_logout(IMAP_WRAPPER_MAGIC_NUMBER);
}

U0 @imap_fetch_wrapper_function(I64 uid)
{
    @imap_fetch(IMAP_WRAPPER_MAGIC_NUMBER, uid);
}

U0 @imap_fetch_array_wrapper_function(U8* sequence_set, U8* data_items)
{
    @imap_fetch_array(IMAP_WRAPPER_MAGIC_NUMBER, sequence_set, data_items);
}

ImapClient* @imap_new(CTask* mem_task)
{
    if (!mem_task) {
        return NULL;
    }

    ImapClient* c = CAlloc(sizeof(ImapClient), mem_task);
    c->mem_task = mem_task;
    c->o = Json.CreateObject(mem_task);
    c->state = 0;
    StrCpy(&c->selected, "");
    @imap_init_buffers(c);
    c->tag = 1;

    U64 a;
    I64 buffer_size = (MSize(&@imap_close_wrapper_function) + MSize(&@imap_connect_wrapper_function) + MSize(&@imap_fetch_wrapper_function) + MSize(&@imap_fetch_array_wrapper_function) + MSize(&@imap_select_wrapper_function) + MSize(&@imap_login_wrapper_function) + MSize(&@imap_logout_wrapper_function));
    buffer_size += buffer_size % 16;

    U64 code_ptr = CAlloc(buffer_size, c->mem_task->code_heap);
    I64 code_size = 0;

    // Create a copy of function and patch close
    code_size = MSize(&@imap_close_wrapper_function);
    c->close = code_ptr;
    MemCpy(c->close, &@imap_close_wrapper_function, code_size);
    code_ptr += code_size;

    a = c->close;
    while (a(U64*)[0] != IMAP_WRAPPER_MAGIC_NUMBER)
        ++a;
    MemSetI64(a, c, 1);
    a += 9;
    @patch_call_rel32(a, &@imap_close);

    // Create a copy of function and patch connect
    code_size = MSize(&@imap_connect_wrapper_function);
    c->connect = code_ptr;
    MemCpy(c->connect, &@imap_connect_wrapper_function, code_size);
    code_ptr += code_size;

    a = c->connect;
    while (a(U64*)[0] != IMAP_WRAPPER_MAGIC_NUMBER)
        ++a;
    MemSetI64(a, c, 1);
    a += 9;
    @patch_call_rel32(a, &@imap_connect);

    // Create a copy of function and patch select
    code_size = MSize(&@imap_select_wrapper_function);
    c->select = code_ptr;
    MemCpy(c->select, &@imap_select_wrapper_function, code_size);
    code_ptr += code_size;

    a = c->select;
    while (a(U64*)[0] != IMAP_WRAPPER_MAGIC_NUMBER)
        ++a;
    MemSetI64(a, c, 1);
    a += 9;
    @patch_call_rel32(a, &@imap_select);

    // Create a copy of function and patch login
    code_size = MSize(&@imap_login_wrapper_function);
    c->login = code_ptr;
    MemCpy(c->login, &@imap_login_wrapper_function, code_size);
    code_ptr += code_size;

    a = c->login;
    while (a(U64*)[0] != IMAP_WRAPPER_MAGIC_NUMBER)
        ++a;
    MemSetI64(a, c, 1);
    a += 9;
    @patch_call_rel32(a, &@imap_login);

    // Create a copy of function and patch logout
    code_size = MSize(&@imap_logout_wrapper_function);
    c->logout = code_ptr;
    MemCpy(c->logout, &@imap_logout_wrapper_function, code_size);
    code_ptr += code_size;

    a = c->logout;
    while (a(U64*)[0] != IMAP_WRAPPER_MAGIC_NUMBER)
        ++a;
    MemSetI64(a, c, 1);
    a += 9;
    @patch_call_rel32(a, &@imap_logout);

    // Create a copy of function and patch fetch
    code_size = MSize(&@imap_fetch_wrapper_function);
    c->fetch = code_ptr;
    MemCpy(c->fetch, &@imap_fetch_wrapper_function, code_size);
    code_ptr += code_size;

    a = c->fetch;
    while (a(U64*)[0] != IMAP_WRAPPER_MAGIC_NUMBER)
        ++a;
    MemSetI64(a, c, 1);
    a += 9;
    @patch_call_rel32(a, &@imap_fetch);

    // Create a copy of function and patch fetch_array
    code_size = MSize(&@imap_fetch_array_wrapper_function);
    c->fetch_array = code_ptr;
    MemCpy(c->fetch_array, &@imap_fetch_array_wrapper_function, code_size);
    code_ptr += code_size;

    a = c->fetch_array;
    while (a(U64*)[0] != IMAP_WRAPPER_MAGIC_NUMBER)
        ++a;
    MemSetI64(a, c, 1);
    a += 9;
    @patch_call_rel32(a, &@imap_fetch_array);

    return c;
}

"imap ";