/* Copyright (C) 2008 Timo Sirainen, LGPLv2.1 */

/*
   export DOVECOT=~/src/dovecot-1.1.0
   gcc -fPIC -g -shared -Wall -I$DOVECOT -I$DOVECOT/src/lib \
     -I$DOVECOT/src/lib-storage -I$DOVECOT/src/lib-mail \
     -I$DOVECOT/src/lib-imap -I$DOVECOT/src/lib-index -DHAVE_CONFIG_H \
     penalty.c -o penalty_plugin.so
*/

#include "lib.h"
#include "array.h"
#include "module-context.h"
#include "mail-search.h"
#include "mail-storage-private.h"

#include <time.h>

#define PENALTY_CONTEXT(obj) \
	MODULE_CONTEXT(obj, penalty_storage_module)

/* penalties in milliseconds */
#define MAIL_CACHE_LOOKUP_PENALTY 200
#define MAIL_BODY_LOOKUP_PENALTY 1000
/* Allow penalty to grow up to n milliseconds until we actually start
   delaying. */
#define PENALTY_MAX_MSECS (1000*10)

struct penalty {
	unsigned int penalty;
	time_t time;
};

const char *penalty_plugin_version = PACKAGE_VERSION;

static void (*penalty_next_hook_mail_storage_created)
	(struct mail_storage *storage);

static MODULE_CONTEXT_DEFINE_INIT(penalty_storage_module,
				  &mail_storage_module_register);

static struct penalty user_penalty;

static unsigned int
search_args_get_penalty(const struct mail_search_arg *args)
{
	unsigned int sub_penalty, penalty = 0;

	for (; args != NULL; args = args->next) {
		switch (args->type) {
		case SEARCH_OR:
		case SEARCH_SUB:
			sub_penalty =
				search_args_get_penalty(args->value.subargs);
			if (sub_penalty > penalty)
				penalty = sub_penalty;
			break;
		case SEARCH_ALL:
		case SEARCH_SEQSET:
		case SEARCH_FLAGS:
		case SEARCH_KEYWORDS:
			/* can be looked up from index - no penalty */
			break;

		case SEARCH_BEFORE:
		case SEARCH_ON:
		case SEARCH_SINCE:
		case SEARCH_SENTBEFORE:
		case SEARCH_SENTON:
		case SEARCH_SENTSINCE:
		case SEARCH_SMALLER:
		case SEARCH_LARGER:
		case SEARCH_HEADER:
		case SEARCH_HEADER_ADDRESS:
		case SEARCH_HEADER_COMPRESS_LWSP:
			/* these could be looked up from cache */
			penalty = I_MAX(penalty, MAIL_CACHE_LOOKUP_PENALTY);
			break;
		case SEARCH_BODY:
		case SEARCH_TEXT:
		case SEARCH_BODY_FAST:
		case SEARCH_TEXT_FAST:
			/* message body lookup */
			penalty = I_MAX(penalty, MAIL_BODY_LOOKUP_PENALTY);
			break;
		}
	}
	return penalty;
}

static void penalty_wait(struct penalty *penalty, unsigned int new_penalty)
{
	time_t now = time(NULL);
	unsigned int drop_penalty;

	if (penalty->penalty == 0)
		penalty->penalty += new_penalty;
	else {
		if (now > penalty->time) {
			/* reduce penalty since the last update */
			drop_penalty = (now - penalty->time) * 1000;
			if (penalty->penalty > drop_penalty)
				penalty->penalty -= drop_penalty;
			else
				penalty->penalty = 0;
		}
		penalty->penalty += new_penalty;
		if (penalty->penalty > PENALTY_MAX_MSECS) {
			/* too much penalty - wait */
			usleep((penalty->penalty - PENALTY_MAX_MSECS) * 1000);
			penalty->penalty = PENALTY_MAX_MSECS;
			now = time(NULL);
		}
	}
	penalty->time = now;
}

static struct mail_search_context *
penalty_search_init(struct mailbox_transaction_context *t,
		    const char *charset, struct mail_search_arg *args,
		    const enum mail_sort_type *sort_program)
{
	union mailbox_module_context *mbox = PENALTY_CONTEXT(t->box);
	struct mail_search_context *ctx;

	ctx = mbox->super.search_init(t, charset, args, sort_program);
	penalty_wait(&user_penalty, search_args_get_penalty(args));
	return ctx;
}

static struct mailbox *
penalty_mailbox_open(struct mail_storage *storage, const char *name,
			      struct istream *input,
			      enum mailbox_open_flags flags)
{
	union mail_storage_module_context *mstorage = PENALTY_CONTEXT(storage);
	struct mailbox *box;
	union mailbox_module_context *mbox;

	box = mstorage->super.mailbox_open(storage, name, input, flags);
	if (box == NULL)
		return NULL;

	mbox = p_new(box->pool, union mailbox_module_context, 1);
	mbox->super = box->v;

	box->v.search_init = penalty_search_init;
	MODULE_CONTEXT_SET_SELF(box, penalty_storage_module, mbox);
	return box;
}

static void penalty_mail_storage_created(struct mail_storage *storage)
{
	union mail_storage_module_context *mstorage;

	mstorage = p_new(storage->pool, union mail_storage_module_context, 1);
	mstorage->super = storage->v;
	storage->v.mailbox_open = penalty_mailbox_open;

	MODULE_CONTEXT_SET_SELF(storage, penalty_storage_module, mstorage);

	if (penalty_next_hook_mail_storage_created != NULL)
		penalty_next_hook_mail_storage_created(storage);
}

void penalty_plugin_init(void);
void penalty_plugin_deinit(void);

void penalty_plugin_init(void)
{
	penalty_next_hook_mail_storage_created = hook_mail_storage_created;
	hook_mail_storage_created = penalty_mail_storage_created;
}

void penalty_plugin_deinit(void)
{
	hook_mail_storage_created = penalty_next_hook_mail_storage_created;
}