/* Copyright (c) 2011-2018 Dovecot authors, see the included COPYING file */

#include "lib.h"
#include "str.h"
#include "imap-util.h"
#include "mail-namespace.h"
#include "mail-search.h"
#include "mailbox-list-iter.h"
#include "lang-tokenizer.h"
#include "lang-filter.h"
#include "lang-user.h"
#include "language.h"
#include "fts-storage.h"
#include "fts-search-args.h"
#include "fts-user.h"
#include "doveadm-print.h"
#include "doveadm-mail.h"
#include "doveadm-mailbox-list-iter.h"
#include "doveadm-fts.h"

const char *doveadm_fts_plugin_version = DOVECOT_ABI_VERSION;

struct fts_tokenize_cmd_context {
	struct doveadm_mail_cmd_context ctx;
	const char *language;
	const char *tokens;
};

struct fts_namespace_cmd_context {
	struct doveadm_mail_cmd_context ctx;
	const char *namespace;
};

static int
cmd_search_box(struct doveadm_mail_cmd_context *ctx,
	       const struct mailbox_info *info)
{
	struct event *event = mailbox_list_get_event(info->ns->list);
	struct mailbox *box;
	struct fts_backend *backend;
	struct fts_result result;
	int ret = 0;

	backend = fts_list_backend(info->ns->list);
	if (backend == NULL) {
		e_error(event, "fts not enabled for %s", info->vname);
		ctx->exit_code = EX_CONFIG;
		return -1;
	}

	i_zero(&result);
	result.pool = pool_alloconly_create("doveadm", 512);
	i_array_init(&result.definite_uids, 16);
	i_array_init(&result.maybe_uids, 16);
	i_array_init(&result.scores, 16);

	box = mailbox_alloc(info->ns->list, info->vname, 0);
	if (fts_backend_lookup(backend, box, ctx->search_args->args,
				      FTS_LOOKUP_FLAG_AND_ARGS, &result) < 0) {
		e_error(event, "fts lookup failed");
		doveadm_mail_failed_error(ctx, MAIL_ERROR_TEMP);
		ret = -1;
	} else {
		printf("%s: ", info->vname);
		if (array_count(&result.definite_uids) == 0)
			printf("no results\n");
		else T_BEGIN {
			string_t *str = t_str_new(128);
			imap_write_seq_range(str, &result.definite_uids);
			printf("%s\n", str_c(str));
		} T_END;
		if (array_count(&result.maybe_uids) > 0) T_BEGIN {
			string_t *str = t_str_new(128);
			imap_write_seq_range(str, &result.maybe_uids);
			printf(" - maybe: %s\n", str_c(str));
		} T_END;
		fts_backend_lookup_done(backend);
	}
	mailbox_free(&box);
	array_free(&result.definite_uids);
	array_free(&result.maybe_uids);
	array_free(&result.scores);
	pool_unref(&result.pool);
	return ret;
}

static int
cmd_fts_lookup_run(struct doveadm_mail_cmd_context *ctx,
		   struct mail_user *user)
{
	const enum mailbox_list_iter_flags iter_flags =
		MAILBOX_LIST_ITER_NO_AUTO_BOXES |
		MAILBOX_LIST_ITER_RETURN_NO_FLAGS;
	struct doveadm_mailbox_list_iter *iter;
	const struct mailbox_info *info;
	int ret = 0;

	iter = doveadm_mailbox_list_iter_init(ctx, user, ctx->search_args,
					      iter_flags);
	while ((info = doveadm_mailbox_list_iter_next(iter)) != NULL) T_BEGIN {
		if (cmd_search_box(ctx, info) < 0)
			ret = -1;
	} T_END;
	if (doveadm_mailbox_list_iter_deinit(&iter) < 0)
		ret = -1;
	return ret;
}

static void
cmd_fts_lookup_init(struct doveadm_mail_cmd_context *_ctx)
{
	struct doveadm_cmd_context *cctx = _ctx->cctx;

	const char *const *query;
	if (!doveadm_cmd_param_array(cctx, "query", &query))
		doveadm_mail_help_name("fts lookup");
	_ctx->search_args = doveadm_mail_build_search_args(query);
}

static struct doveadm_mail_cmd_context *
cmd_fts_lookup_alloc(void)
{
	struct doveadm_mail_cmd_context *ctx;

	ctx = doveadm_mail_cmd_alloc(struct doveadm_mail_cmd_context);
	ctx->v.run = cmd_fts_lookup_run;
	ctx->v.init = cmd_fts_lookup_init;
	return ctx;
}

static int
cmd_fts_expand_run(struct doveadm_mail_cmd_context *ctx,
		   struct mail_user *user)
{
	struct mail_namespace *ns = mail_namespace_find_inbox(user->namespaces);
	struct mailbox *box;
	struct fts_backend *backend;
	string_t *str = t_str_new(128);

	backend = fts_list_backend(ns->list);
	if (backend == NULL) {
		e_error(user->event, "fts not enabled for INBOX");
		ctx->exit_code = EX_CONFIG;
		return -1;
	}

	box = mailbox_alloc(ns->list, "INBOX", 0);
	mail_search_args_init(ctx->search_args, box, FALSE, NULL);

	if (fts_search_args_expand(backend, ctx->search_args) < 0)
		i_fatal("Couldn't expand search args");
	mail_search_args_to_cmdline(str, ctx->search_args->args);
	printf("%s\n", str_c(str));
	mail_search_args_deinit(ctx->search_args);
	mailbox_free(&box);
	return 0;
}

static void
cmd_fts_expand_init(struct doveadm_mail_cmd_context *_ctx)
{
	struct doveadm_cmd_context *cctx = _ctx->cctx;

	const char *const *query;
	if (!doveadm_cmd_param_array(cctx, "query", &query))
		doveadm_mail_help_name("fts expand");
	_ctx->search_args = doveadm_mail_build_search_args(query);
}

static struct doveadm_mail_cmd_context *
cmd_fts_expand_alloc(void)
{
	struct doveadm_mail_cmd_context *ctx;

	ctx = doveadm_mail_cmd_alloc(struct doveadm_mail_cmd_context);
	ctx->v.run = cmd_fts_expand_run;
	ctx->v.init = cmd_fts_expand_init;
	return ctx;
}

static int
cmd_fts_tokenize_run(struct doveadm_mail_cmd_context *_ctx,
		     struct mail_user *user)
{
	struct fts_tokenize_cmd_context *ctx =
		container_of(_ctx, struct fts_tokenize_cmd_context, ctx);

	struct mail_namespace *ns = mail_namespace_find_inbox(user->namespaces);
	struct fts_backend *backend;
	struct language_user *user_lang;
	const struct language *lang = NULL;
	int ret, ret2;
	bool final = FALSE;

	backend = fts_list_backend(ns->list);
	if (backend == NULL) {
		e_error(user->event, "fts not enabled for INBOX");
		_ctx->exit_code = EX_CONFIG;
		return -1;
	}

	if (ctx->language == NULL) {
		struct language_list *lang_list =
			lang_user_get_language_list(user);
		enum language_detect_result result;
		const char *error;

		result = language_detect(lang_list,
		    (const unsigned char *)ctx->tokens, strlen(ctx->tokens),
                    &lang, &error);
		if (lang == NULL)
			lang = language_list_get_first(lang_list);
		switch (result) {
		case LANGUAGE_DETECT_RESULT_SHORT:
			e_warning(user->event,
				  "Text too short, can't detect its language - assuming %s",
				  lang->name);
			break;
		case LANGUAGE_DETECT_RESULT_UNKNOWN:
			e_warning(user->event,
				  "Can't detect its language - assuming %s",
				  lang->name);
			break;
		case LANGUAGE_DETECT_RESULT_OK:
			break;
		case LANGUAGE_DETECT_RESULT_ERROR:
			e_error(user->event,
				"Language detection library initialization failed: %s",
				error);
			_ctx->exit_code = EX_CONFIG;
			return -1;
		default:
			i_unreached();
		}
	} else {
		lang = language_find(ctx->language);
		if (lang == NULL) {
			e_error(user->event,
				"Unknown language: %s", ctx->language);
			_ctx->exit_code = EX_USAGE;
			return -1;
		}
	}
	user_lang = lang_user_language_find(user, lang);
	if (user_lang == NULL) {
		e_error(user->event,
			"Language not enabled for user: %s", ctx->language);
		_ctx->exit_code = EX_USAGE;
		return -1;
	}

	lang_tokenizer_reset(user_lang->index_tokenizer);
	for (;;) {
		const char *token, *error;

		if (!final) {
			ret = lang_tokenizer_next(user_lang->index_tokenizer,
				(const unsigned char *)ctx->tokens, strlen(ctx->tokens),
				&token, &error);
		} else {
			ret = lang_tokenizer_final(user_lang->index_tokenizer,
						  &token, &error);
		}
		if (ret < 0)
			break;
		if (ret > 0 && user_lang->filter != NULL) {
			ret2 = lang_filter(user_lang->filter, &token, &error);
			if (ret2 > 0)
				doveadm_print(token);
			else if (ret2 < 0)
				e_error(user->event,
					"Couldn't create indexable tokens: %s",
					error);
		}
		if (ret == 0) {
			if (final)
				break;
			final = TRUE;
		}
	}
	return 0;
}

static void
cmd_fts_tokenize_init(struct doveadm_mail_cmd_context *_ctx)
{
	struct doveadm_cmd_context *cctx = _ctx->cctx;
	struct fts_tokenize_cmd_context *ctx =
		container_of(_ctx, struct fts_tokenize_cmd_context, ctx);

	(void)doveadm_cmd_param_str(cctx, "language", &ctx->language);

	const char *const *args;
	if (!doveadm_cmd_param_array(cctx, "text", &args))
		doveadm_mail_help_name("fts tokenize");

	ctx->tokens = p_strdup(_ctx->pool, t_strarray_join(args, " "));

	doveadm_print_header("token", "token", DOVEADM_PRINT_HEADER_FLAG_HIDE_TITLE);
}

static struct doveadm_mail_cmd_context *
cmd_fts_tokenize_alloc(void)
{
	struct fts_tokenize_cmd_context *ctx;

	ctx = doveadm_mail_cmd_alloc(struct fts_tokenize_cmd_context);
	ctx->ctx.v.run = cmd_fts_tokenize_run;
	ctx->ctx.v.init = cmd_fts_tokenize_init;
	doveadm_print_init(DOVEADM_PRINT_TYPE_FLOW);
	return &ctx->ctx;
}

static int
fts_namespace_find(struct mail_user *user, const char *ns_prefix,
		   struct mail_namespace **ns_r)
{
	struct mail_namespace *ns;

	if (ns_prefix == NULL)
		ns = mail_namespace_find_inbox(user->namespaces);
	else {
		ns = mail_namespace_find_prefix(user->namespaces, ns_prefix);
		if (ns == NULL) {
			e_error(user->event,
				"Namespace prefix not found: %s", ns_prefix);
			return -1;
		}
	}

	if (fts_list_backend(ns->list) == NULL) {
		e_error(user->event, "fts not enabled for user's namespace %s",
			ns_prefix != NULL ? ns_prefix : "INBOX");
		return -1;
	}
	*ns_r = ns;
	return 0;
}

static int
cmd_fts_optimize_run(struct doveadm_mail_cmd_context *_ctx,
		     struct mail_user *user)
{
	struct fts_namespace_cmd_context *ctx =
		container_of(_ctx, struct fts_namespace_cmd_context, ctx);

	const char *ns_prefix = ctx->namespace;
	struct mail_namespace *ns;
	struct fts_backend *backend;

	if (fts_namespace_find(user, ns_prefix, &ns) < 0) {
		doveadm_mail_failed_error(_ctx, MAIL_ERROR_NOTFOUND);
		return -1;
	}
	backend = fts_list_backend(ns->list);
	if (fts_backend_optimize(backend) < 0) {
		e_error(user->event, "fts optimize failed");
		doveadm_mail_failed_error(_ctx, MAIL_ERROR_TEMP);
		return -1;
	}
	return 0;
}

static void
cmd_fts_optimize_init(struct doveadm_mail_cmd_context *_ctx)
{
	struct doveadm_cmd_context *cctx = _ctx->cctx;
	struct fts_namespace_cmd_context *ctx =
		container_of(_ctx, struct fts_namespace_cmd_context, ctx);

	(void)doveadm_cmd_param_str(cctx, "namespace", &ctx->namespace);
}

static struct doveadm_mail_cmd_context *
cmd_fts_optimize_alloc(void)
{
	struct fts_namespace_cmd_context *ctx =
		doveadm_mail_cmd_alloc(struct fts_namespace_cmd_context);
	ctx->ctx.v.run = cmd_fts_optimize_run;
	ctx->ctx.v.init = cmd_fts_optimize_init;
	return &ctx->ctx;
}

static int
cmd_fts_rescan_run(struct doveadm_mail_cmd_context *_ctx, struct mail_user *user)
{
	struct fts_namespace_cmd_context *ctx =
		container_of(_ctx, struct fts_namespace_cmd_context, ctx);


	const char *ns_prefix = ctx->namespace;
	struct mail_namespace *ns;
	struct fts_backend *backend;

	if (fts_namespace_find(user, ns_prefix, &ns) < 0) {
		doveadm_mail_failed_error(_ctx, MAIL_ERROR_NOTFOUND);
		return -1;
	}
	backend = fts_list_backend(ns->list);
	if (fts_backend_rescan(backend) < 0) {
		e_error(user->event, "fts rescan failed");
		doveadm_mail_failed_error(_ctx, MAIL_ERROR_TEMP);
		return -1;
	}
	return 0;
}

static void
cmd_fts_rescan_init(struct doveadm_mail_cmd_context *_ctx)
{
	struct doveadm_cmd_context *cctx = _ctx->cctx;
	struct fts_namespace_cmd_context *ctx =
		container_of(_ctx, struct fts_namespace_cmd_context, ctx);

	(void)doveadm_cmd_param_str(cctx, "namespace", &ctx->namespace);
}

static struct doveadm_mail_cmd_context *
cmd_fts_rescan_alloc(void)
{
	struct fts_namespace_cmd_context *ctx =
		doveadm_mail_cmd_alloc(struct fts_namespace_cmd_context);
	ctx->ctx.v.run = cmd_fts_rescan_run;
	ctx->ctx.v.init = cmd_fts_rescan_init;
	return &ctx->ctx;
}

static struct doveadm_cmd_ver2 fts_commands[] = {
{
	.name = "fts lookup",
	.mail_cmd = cmd_fts_lookup_alloc,
	.usage = DOVEADM_CMD_MAIL_USAGE_PREFIX "<search query>",
DOVEADM_CMD_PARAMS_START
DOVEADM_CMD_MAIL_COMMON
DOVEADM_CMD_PARAM('\0', "query", CMD_PARAM_ARRAY, CMD_PARAM_FLAG_POSITIONAL)
DOVEADM_CMD_PARAMS_END
},
{
	.name = "fts expand",
	.mail_cmd = cmd_fts_expand_alloc,
	.usage = DOVEADM_CMD_MAIL_USAGE_PREFIX "<search query>",
DOVEADM_CMD_PARAMS_START
DOVEADM_CMD_MAIL_COMMON
DOVEADM_CMD_PARAM('\0', "query", CMD_PARAM_ARRAY, CMD_PARAM_FLAG_POSITIONAL)
DOVEADM_CMD_PARAMS_END
},
{
	.name = "fts tokenize",
	.mail_cmd = cmd_fts_tokenize_alloc,
	.usage = DOVEADM_CMD_MAIL_USAGE_PREFIX "<text>",
DOVEADM_CMD_PARAMS_START
DOVEADM_CMD_MAIL_COMMON
DOVEADM_CMD_PARAM('l', "language", CMD_PARAM_STR, 0)
DOVEADM_CMD_PARAM('\0', "text", CMD_PARAM_ARRAY, CMD_PARAM_FLAG_POSITIONAL)
DOVEADM_CMD_PARAMS_END
},
{
	.name = "fts optimize",
	.mail_cmd = cmd_fts_optimize_alloc,
	.usage = DOVEADM_CMD_MAIL_USAGE_PREFIX "[<namespace>]",
DOVEADM_CMD_PARAMS_START
DOVEADM_CMD_MAIL_COMMON
DOVEADM_CMD_PARAM('\0', "namespace", CMD_PARAM_STR, CMD_PARAM_FLAG_POSITIONAL)
DOVEADM_CMD_PARAMS_END
},
{
	.name = "fts rescan",
	.mail_cmd = cmd_fts_rescan_alloc,
	.usage = DOVEADM_CMD_MAIL_USAGE_PREFIX "[<namespace>]",
DOVEADM_CMD_PARAMS_START
DOVEADM_CMD_MAIL_COMMON
DOVEADM_CMD_PARAM('\0', "namespace", CMD_PARAM_STR, CMD_PARAM_FLAG_POSITIONAL)
DOVEADM_CMD_PARAMS_END
},
};

void doveadm_fts_plugin_init(struct module *module ATTR_UNUSED)
{
	unsigned int i;

	for (i = 0; i < N_ELEMENTS(fts_commands); i++)
		doveadm_cmd_register_ver2(&fts_commands[i]);
}

void doveadm_fts_plugin_deinit(void)
{
}
