#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <limits.h>
#include <assert.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/time.h>

#include "stellar/packet.h"
#include "stellar/utils.h"
#include "stellar/mq.h"
#include "stellar/session.h"
#include "stellar/mail.h"

#include "mail_decoder_util.h"
#include "mail_decoder_smtp.h"
//#include "mail_decoder_pop3.h"
//#include "mail_decoder_imap.h"
#include "mail_decoder_module.h"

const char * g_mail_topic_name[MAIL_TOPIC_MAX] = {
    "MAIL_COMMAND",
    "MAIL_HEADER",
    "MAIL_BODY",
    "MAIL_ATTACHMENT",
    "MAIL_EML",
};

static unsigned short get_session_dest_port(struct session *sess)
{
    const struct layer *first_pkt_layer;
    const struct packet *first_pkt;
    int layer_count;

    first_pkt = session_get_current_packet(sess);
    layer_count = packet_get_layer_count(first_pkt);
    first_pkt_layer = packet_get_layer_by_idx(first_pkt, layer_count - 1);

    switch (first_pkt_layer->proto) {
        case LAYER_PROTO_TCP:
            return ntohs(first_pkt_layer->hdr.tcp->dest);
        case LAYER_PROTO_UDP:
            return ntohs(first_pkt_layer->hdr.udp->dest);;
        default:
            break;
    }
    return 0;
}

static int is_current_packet_c2s(struct session *sess)
{
    const struct layer *first_pkt_layer, *curr_pkt_layer;
    
    const struct packet *first_pkt, *curr_pkt;

    first_pkt = session_get_first_packet(sess, FLOW_TYPE_C2S);
    curr_pkt = session_get_current_packet(sess);

    int i = 0;
    int layer_count = packet_get_layer_count(first_pkt);

    for (i = 0; i < 2 /* l3/l4 */; i++) {
        first_pkt_layer = packet_get_layer_by_idx(first_pkt, layer_count - 1 -i);
        curr_pkt_layer = packet_get_layer_by_idx(curr_pkt, layer_count - 1 -i);

        if (first_pkt_layer->proto != curr_pkt_layer->proto) {
            return 0;
        }
        switch (first_pkt_layer->proto) {
            case LAYER_PROTO_TCP:
                if (first_pkt_layer->hdr.tcp->dest != curr_pkt_layer->hdr.tcp->dest) {
                    return 0;
                }
                break;
            case LAYER_PROTO_UDP:
                if (first_pkt_layer->hdr.udp->dest != curr_pkt_layer->hdr.udp->dest) {
                    return 0;
                }
                break;
            case LAYER_PROTO_IPV4:
                if (0 != memcmp(&first_pkt_layer->hdr.ip4->ip_dst, &curr_pkt_layer->hdr.ip4->ip_dst, sizeof(struct in_addr))) {
                    return 0;
                }
                break;
            case LAYER_PROTO_IPV6:
                if (0 != memcmp(&first_pkt_layer->hdr.ip6->ip6_dst, &curr_pkt_layer->hdr.ip6->ip6_dst, sizeof(struct in6_addr))) {
                    return 0;
                }
                break;
            default:
                return 0;
        }

    }

    return 1;
}

static int mail_subscribe_command(struct mq_schema *schema, mail_command_callback_func *cb, void *arg)
{
    int topic;

    topic = mq_schema_get_topic_id(schema, g_mail_topic_name[MAIL_TOPIC_COMMAND]);
    if (topic < 0) {
        return -1;
    }

    return mq_schema_subscribe(schema, topic, (on_msg_cb_func *)(void *)cb, arg);
}

static int mail_subscribe_header(struct mq_schema *schema, mail_header_callback_func *cb, void *arg)
{
    int topic;

    topic = mq_schema_get_topic_id(schema, g_mail_topic_name[MAIL_TOPIC_HEADER]);
    if (topic < 0) {
        return -1;
    }

    return mq_schema_subscribe(schema, topic, (on_msg_cb_func *)(void *)cb, arg);
}

static int mail_subscribe_body(struct mq_schema *schema, mail_body_callback_func *cb, void *arg)
{
    int topic;

    topic = mq_schema_get_topic_id(schema, g_mail_topic_name[MAIL_TOPIC_BODY]);
    if (topic < 0) {
        return -1;
    }

    return mq_schema_subscribe(schema, topic, (on_msg_cb_func *)(void *)cb, arg);
}

static int mail_subscribe_attachment(    struct mq_schema *schema, mail_attachment_callback_func *cb, void *arg)
{
    int topic;

    topic = mq_schema_get_topic_id(schema, g_mail_topic_name[MAIL_TOPIC_ATTACHMENT]);
    if (topic < 0) {
        return -1;
    }

    return mq_schema_subscribe(schema, topic, (on_msg_cb_func *)(void *)cb, arg);
}

static int mail_subscribe_eml(struct mq_schema *schema, mail_eml_callback_func *cb, void *arg)
{
    int topic;

    topic = mq_schema_get_topic_id(schema, g_mail_topic_name[MAIL_TOPIC_EML]);
    if (topic < 0) {
        return -1;
    }

    return mq_schema_subscribe(schema, topic, (on_msg_cb_func *)(void *)cb, arg);
}

static void mail_dispatch(int topic, void *msg, on_msg_cb_func *msg_cb, void *arg, void *dispatch_arg)
{
    struct mail_message *mail_msg;
    struct mail_decoder *decoder;
    struct mail_topics *topics;

    mail_msg = (struct mail_message *)msg;
    decoder = (struct mail_decoder *)dispatch_arg;
    topics = &decoder->topics;

    //char *msg_json = mail_message_to_json(mail_msg);
    //printf("%s\n", msg_json);
    //free(msg_json);

    if (topic == topics->topic_ids[MAIL_TOPIC_COMMAND]) {
        mail_command_callback_func *command_cb = (mail_command_callback_func *)(void *)msg_cb;
        command_cb(mail_msg->sess_ref, mail_msg->mail_protocol, mail_msg->mail_seq,
                mail_msg->command->cmd, (const char *)mail_msg->command->arg, mail_msg->command->arg_len,
                (const char *)mail_msg->command->cmd_line, mail_msg->command->cmd_line_len,
                arg);
        return;
    }

    if (topic == topics->topic_ids[MAIL_TOPIC_HEADER]) {
        mail_header_callback_func *header_cb = (mail_header_callback_func *)(void *)msg_cb;
        header_cb(mail_msg->sess_ref, mail_msg->mail_protocol, mail_msg->mail_seq, (const struct mail_header *)mail_msg->header, arg);
        return;
    }

    if (topic == topics->topic_ids[MAIL_TOPIC_BODY]) {
        mail_body_callback_func *body_cb = (mail_body_callback_func *)(void *)msg_cb;
        body_cb(mail_msg->sess_ref, mail_msg->mail_protocol, mail_msg->mail_seq,
                (const char *)mail_msg->body->body, mail_msg->body->body_len,
                mail_msg->body->body_offset, mail_msg->body->is_body_finished,
                arg);
        return;
    }

    if (topic == topics->topic_ids[MAIL_TOPIC_ATTACHMENT]) {
        mail_attachment_callback_func *attachment_cb = (mail_attachment_callback_func *)(void *)msg_cb;
        attachment_cb(mail_msg->sess_ref, mail_msg->mail_protocol, mail_msg->mail_seq,
                (const char *)mail_msg->attachment->attachment_name, mail_msg->attachment->attachment_name_len,
                (const char *)mail_msg->attachment->attachment, mail_msg->attachment->attachment_len,
                mail_msg->attachment->attachment_offset, mail_msg->attachment->is_attachment_finished,
                arg);
        return;
    }

    if (topic == topics->topic_ids[MAIL_TOPIC_EML]) {
        mail_eml_callback_func *eml_cb = (mail_eml_callback_func *)(void *)msg_cb;
        eml_cb(mail_msg->sess_ref, mail_msg->mail_protocol,
                (const char *)mail_msg->eml->eml, mail_msg->eml->eml_len,
                mail_msg->eml->eml_offset, mail_msg->eml->is_eml_finished,
                arg);
        return;
    }
}

int mail_subscribe(struct mail_decoder *decoder,
                        mail_command_callback_func command_cb,
                        mail_header_callback_func *header_cb,
                        mail_body_callback_func *body_cb,
                        mail_attachment_callback_func *attachment_cb,
                        mail_eml_callback_func *eml_cb,
                        void *arg)
{
    int ret;

    ret = mail_subscribe_command(decoder->mq_schema, command_cb, arg);
    if (ret < 0) {
        return ret;
    }

    ret = mail_subscribe_header(decoder->mq_schema, header_cb, arg);
    if (ret < 0) {
        return ret;
    }

    ret = mail_subscribe_body(decoder->mq_schema, body_cb, arg);
    if (ret < 0) {
        return ret;
    }

    ret = mail_subscribe_attachment(decoder->mq_schema, attachment_cb, arg);
    if (ret < 0) {
        return ret;
    }

    ret = mail_subscribe_eml(decoder->mq_schema, eml_cb, arg);
    if (ret < 0) {
        return ret;
    }

    return 0;
}

static void mail_decoder_session_ctx_free(struct mail_session_ctx *session_ctx)
{
    switch (session_ctx->protocol) {
        case MAIL_PROTOCOL_SMTP:
            smtp_parser_free(session_ctx->smtp_parser);
            break;
        case MAIL_PROTOCOL_POP3:
            //mail_decoder_pop3_exit(session_ctx->mail_pme);
            break;
        case MAIL_PROTOCOL_IMAP:
            //mail_decoder_imap_exit(session_ctx->mail_pme);
            break;
        default:
            break;
    }
    free(session_ctx);
}

static struct mail_session_ctx *mail_decoder_session_ctx_new(void)
{
    struct mail_session_ctx *session_ctx;
    session_ctx = (struct mail_session_ctx *)calloc(1, sizeof(struct mail_session_ctx));
    return session_ctx;
}

//static int mail_decoder_imap_process(struct mail_decoder *env, struct mail_session_ctx *session_ctx, const char *payload, size_t payload_len, int is_c2s, int tid)
//{
//    return 0;
//}
//
//static int mail_decoder_pop3_process(struct mail_decoder *env, struct mail_session_ctx *session_ctx, const char *payload, size_t payload_len, int is_c2s, int tid)
//{
//    return 0;
//}

static int mail_decoder_smtp_process(struct mail_decoder *decoder, struct mail_session_ctx *session_ctx, const char *payload, size_t payload_len, int is_c2s)
{
    char ret;
    struct mq_runtime *mq;

    mq = module_manager_get_mq_runtime(decoder->mod_mgr);

    if (session_ctx->smtp_parser == NULL && mq != NULL) {
        session_ctx->smtp_parser = smtp_parser_new(mq, &decoder->topics, is_c2s);
    }

    ret = smtp_parser_entry(session_ctx->smtp_parser, session_ctx->sess_ref, payload, payload_len, is_c2s);
    if (ret != 0) {
        return -1;
    }

    return 0;
}

static void mail_decoder_on_tcp_payload(struct session *sess, enum session_state state, const char *payload, uint32_t payload_len, void *arg)
{
    (void)state;
    int ret;
    int is_packet_c2s;
    struct mail_session_ctx *session_ctx;
    struct mail_decoder *decoder = (struct mail_decoder *)arg;

    if (payload == NULL || payload_len == 0) {
        return;
    }

    is_packet_c2s = is_current_packet_c2s(sess);

    session_ctx = (struct mail_session_ctx *)session_get_exdata(sess, decoder->exdata_id);
    if (session_ctx == NULL) {
        session_ctx = mail_decoder_session_ctx_new();
        session_ctx->sess_ref = sess;
        session_ctx->protocol = MAIL_PROTOCOL_MAX;
        session_ctx->is_droped = 0;
        session_set_exdata(sess, decoder->exdata_id, session_ctx);

        unsigned short session_dest_port = get_session_dest_port(sess);

        switch (session_dest_port) {
            case PORT_SMTP:
                session_ctx->protocol = MAIL_PROTOCOL_SMTP;
                break;
            case PORT_POP3:
                session_ctx->protocol = MAIL_PROTOCOL_POP3;
                break;
            case PORT_IMAP:
                session_ctx->protocol = MAIL_PROTOCOL_IMAP;
                break;
            default:
                if (smtp_identify(payload, payload_len, is_packet_c2s)) {
                    session_ctx->protocol = MAIL_PROTOCOL_SMTP;
                    break;
                }
                //if (pop3_identify((char *)payload, payload_len, is_packet_c2s)) {
                //    session_ctx->protocol = MAIL_PROTOCOL_POP3;
                //    break;
                //}
                //if (imap_identify((char *)payload, payload_len, is_packet_c2s)) {
                //    session_ctx->protocol = MAIL_PROTOCOL_IMAP;
                //    break;
                //}
                break;
        }
    }

    if (session_ctx->is_droped) {
        return;
    }

    switch (session_ctx->protocol) {
        case MAIL_PROTOCOL_SMTP:
            ret = mail_decoder_smtp_process(decoder, session_ctx, payload, payload_len, is_packet_c2s);
            break;
        case MAIL_PROTOCOL_POP3:
            //ret = mail_decoder_pop3_process(decoder, session_ctx, payload, payload_len , is_packet_c2s);
            break;
        case MAIL_PROTOCOL_IMAP:
            //ret = mail_decoder_imap_process(decoder, session_ctx, payload, payload_len , is_packet_c2s);
            break;
        default:
            break;
    }

    if (ret != 0) {
        session_ctx->is_droped = 1;
    }
}

void mail_decoder_on_exdata_free(int idx, void *ex_ptr, void *arg)
{
    (void)(idx);
    (void)(arg);
    struct mail_session_ctx *session_ctx = (struct mail_session_ctx *)ex_ptr;

    if (session_ctx) {
        mail_decoder_session_ctx_free(session_ctx);
    }
}

void mail_exit(struct module_manager *mod_mgr, struct module *mod)
{
    (void)(mod_mgr);
    struct mail_decoder *decoder;
    struct mq_schema *schema;

    if (mod) {
        decoder = (struct mail_decoder *)module_get_ctx(mod);
        if (decoder) {
            schema = module_manager_get_mq_schema(mod_mgr);
            int i;
            for (i = 0; i < MAIL_TOPIC_MAX; i++) {
                mq_schema_destroy_topic(schema, decoder->topics.topic_ids[i]);
            }
            free(decoder);
        }

        module_free(mod);
    }
}

struct module* mail_init(struct module_manager *mod_mgr)
{
    int ret;
    struct module *mod;
    struct session_manager *sess_mgr;
    struct mail_decoder *decoder;
    struct mail_topics *topics;

    decoder = (struct mail_decoder *)calloc(1, sizeof(struct mail_decoder));
    decoder->mod_mgr = mod_mgr;
    mod = module_new(MAIL_MODULE_NAME, decoder);

    sess_mgr = module_to_session_manager(module_manager_get_module(mod_mgr, SESSION_MANAGER_MODULE_NAME));
    if (sess_mgr == NULL) {
        goto exit;
    }

    ret = session_manager_subscribe_tcp_stream(sess_mgr, mail_decoder_on_tcp_payload, decoder);
    if (ret < 0) {
        goto exit;
    }

    decoder->exdata_id = session_manager_new_session_exdata_index(sess_mgr, MAIL_EXDATA_NAME, mail_decoder_on_exdata_free, NULL);	
    if (decoder->exdata_id < 0) {
        goto exit;
    }

    decoder->mq_schema = module_manager_get_mq_schema(mod_mgr);

    int i;
    topics = &decoder->topics;
    for (i = 0; i < MAIL_TOPIC_MAX; i++) {
    topics->topic_ids[i] = mq_schema_create_topic(decoder->mq_schema, g_mail_topic_name[i], mail_dispatch, decoder, mail_message_free, NULL);
        if (topics->topic_ids[i] < 0) {
            goto exit;
        }
        topics->topic_names[i] = g_mail_topic_name[i];
    }

    return mod;
exit:
    mail_exit(mod_mgr, mod);
    return NULL;
}

struct mail_decoder *module_to_mail_decoder(struct module *mod)
{
    assert(mod);
    assert(strcmp(module_get_name(mod), MAIL_MODULE_NAME) == 0);
    return module_get_ctx(mod);
}