#include <stdio.h>
#include <string.h>
#include <locale.h>

#include "filter.h"
#include "filter-manager.h"
#include "filter-utils.h"
#include "filter-kvs.h"

#define USE_QDBM 1
#ifdef USE_QDBM
#  include "filter-kvs-qdbm.h"
#endif

#include "textcontent-filter.h"
#include "blacklist-filter.h"
#include "whitelist-filter.h"
#include "wordsep-filter.h"
#include "ngram-filter.h"
#include "bayes-filter.h"

enum {
	MODE_TEST_JUNK,
	MODE_LEARN_JUNK,
	MODE_LEARN_CLEAN,
	MODE_SHOW_STATUS
};


static int verbose = 0;

static int learn_filter(int mode, const char *file);
static int test_filter(int mode, const char *file);
static void print_message_data(XMessageData *msgdata);


int main(int argc, char *argv[])
{
	int retval = 0;
	int i;
	int mode = MODE_TEST_JUNK;

	setlocale(LC_ALL, "");

	xfilter_init();

#ifdef USE_QDBM
	xfilter_kvs_qdbm_set_engine();
#endif
	xfilter_bayes_db_init();

	for (i = 1; i < argc; i++) {
		if (!strcmp(argv[i], "-j"))
			mode = MODE_LEARN_JUNK;
		else if (!strcmp(argv[i], "-c"))
			mode = MODE_LEARN_CLEAN;
		else if (!strcmp(argv[i], "-t"))
			mode = MODE_TEST_JUNK;
		else if (!strcmp(argv[i], "-s"))
			mode = MODE_SHOW_STATUS;
		else if (!strcmp(argv[i], "-v"))
			verbose = 1;
		else if (!strcmp(argv[i], "-d"))
			xfilter_set_debug_mode(1);
	}

	if (mode == MODE_SHOW_STATUS) {
		retval = xfilter_bayes_db_show_contents();
	} else if (mode == MODE_LEARN_JUNK || mode == MODE_LEARN_CLEAN) {
		for (i = 1; i < argc; i++) {
			if (argv[i][0] != '-') {
				retval = learn_filter(mode, argv[i]);
				if (retval != 0)
					break;
			}
		}
	} else {
		for (i = 1; i < argc; i++) {
			if (argv[i][0] != '-') {
				retval = test_filter(mode, argv[i]);
				if (retval == 127)
					break;
			}
		}
	}

	xfilter_bayes_db_done();
	xfilter_done();

	return retval;
}

static int learn_filter(int mode, const char *file)
{
	XFilterManager *mgr;
	XMessageData *msgdata;
	XMessageData *resdata;
	XFilterResult *res;
	XFilterStatus status;
	int retval = 0;

	XFilterConstructorFunc learn_junk_ctors[] = {
		xfilter_textcontent_new,
		xfilter_wordsep_new,
		xfilter_ngram_new,
		xfilter_bayes_learn_junk_new,
		NULL
	};

	XFilterConstructorFunc learn_nojunk_ctors[] = {
		xfilter_textcontent_new,
		xfilter_wordsep_new,
		xfilter_ngram_new,
		xfilter_bayes_learn_nojunk_new,
		NULL
	};

	if (verbose)
		printf("learning message file: %s\n", file);

	mgr = xfilter_manager_new();

	if (mode == MODE_LEARN_JUNK)
		xfilter_manager_add_filters(mgr, learn_junk_ctors);
	else
		xfilter_manager_add_filters(mgr, learn_nojunk_ctors);

	msgdata = xfilter_message_data_read_file(file, "message/rfc822");

	res = xfilter_manager_run(mgr, msgdata);
	if (verbose)
		xfilter_result_print(res);
	status = xfilter_result_get_status(res);
	if (status == XF_UNSUPPORTED_TYPE || status == XF_ERROR) {
		printf("%s: Error on learning mail\n", file);
		retval = 127;
	}

	if (xfilter_get_debug_mode()) {
		resdata = xfilter_result_get_message_data(res);
		print_message_data(resdata);
	}

	xfilter_result_free(res);
	xfilter_message_data_free(msgdata);
	xfilter_manager_free(mgr);

	return retval;
}

static int test_filter(int mode, const char *file)
{
	XFilterManager *mgr;
	XMessageData *msgdata;
	XMessageData *resdata;
	XFilterResult *res;
	XFilterStatus status;
	int retval = 0;

	XFilterConstructorFunc ctors[] = {
		xfilter_textcontent_new,
		xfilter_blacklist_new,
		xfilter_whitelist_new,
		xfilter_wordsep_new,
		xfilter_ngram_new,
		xfilter_bayes_new,
		NULL
	};

	if (verbose)
		printf("testing message file: %s\n", file);

	mgr = xfilter_manager_new();
	xfilter_manager_add_filters(mgr, ctors);

	msgdata = xfilter_message_data_read_file(file, "message/rfc822");

	res = xfilter_manager_run(mgr, msgdata);
	if (verbose)
		xfilter_result_print(res);
	status = xfilter_result_get_status(res);
	if (status == XF_JUNK) {
		printf("%s: This is a junk mail (prob: %f)\n", file, xfilter_result_get_probability(res));
		retval = 0;
	} else if (status == XF_UNCERTAIN) {
		printf("%s: This mail could not be classified (prob: %f)\n", file, xfilter_result_get_probability(res));
		retval = 2;
	} else if (status == XF_UNSUPPORTED_TYPE || status == XF_ERROR) {
		printf("%s: Error on testing mail\n", file);
		retval = 127;
	} else {
		printf("%s: This is a clean mail (prob: %f)\n", file, xfilter_result_get_probability(res));
		retval = 1;
	}

	if (xfilter_get_debug_mode()) {
		resdata = xfilter_result_get_message_data(res);
		print_message_data(resdata);
	}

	xfilter_result_free(res);
	xfilter_message_data_free(msgdata);

	xfilter_manager_free(mgr);

	return retval;
}

static void print_message_data(XMessageData *msgdata)
{
	const char *content;

	if (!msgdata)
		return;

	printf("\n");

	content = xfilter_message_data_get_attribute(msgdata, XM_FROM);
	if (content)
		printf("from: %s\n", content);
	content = xfilter_message_data_get_attribute(msgdata, XM_TO);
	if (content)
		printf("to: %s\n", content);
	content = xfilter_message_data_get_attribute(msgdata, XM_CC);
	if (content)
		printf("cc: %s\n", content);
	content = xfilter_message_data_get_attribute(msgdata, XM_SUBJECT);
	if (content)
		printf("subject: %s\n", content);
	content = xfilter_message_data_get_content(msgdata);
	printf("content: %s\n", content);
}