/* Copyright (c) 2012 Dovecot authors, see the included COPYING file */

#include "lib.h"
#include "strfuncs.h"
#include "str.h"
#include "array.h"
#include "net.h"
#include "istream.h"
#include "ostream.h"
#include "imap-url.h"
#include "imap-common.h"
#include "imap-commands.h"
#include "imap-urlauth.h"
#include "imap-urlauth-fetch.h"

struct cmd_urlfetch_context {
	struct imap_urlauth_fetch *ufetch;
	struct istream *input;

	unsigned int failed:1;
	unsigned int finished:1;
	unsigned int extended:1;
};

struct cmd_urlfetch_url {
	const char *url;
	
	enum imap_urlauth_fetch_flags flags;
};

static void cmd_urlfetch_finish
(struct client_command_context *cmd)
{
	struct cmd_urlfetch_context *ctx =
		(struct cmd_urlfetch_context *)cmd->context;

	if (ctx->finished)
		return;
	ctx->finished = TRUE;

	if (ctx->ufetch != NULL)
		imap_urlauth_fetch_deinit(&ctx->ufetch);

	if (ctx->failed) {
		client_send_internal_error(cmd);
		return;
	}

	client_send_tagline(cmd, "OK URLFETCH completed.");
}

static bool cmd_urlfetch_cancel(struct client_command_context *cmd)
{
	struct cmd_urlfetch_context *ctx =
		(struct cmd_urlfetch_context *)cmd->context;

	if (!cmd->cancel)
		return FALSE;

	if (ctx->ufetch != NULL) {
		if (cmd->client->user->mail_debug) {
			i_debug("URLFETCH: Canceling command; "
				"aborting URLAUTH fetch requests prematurely");
		}
		imap_urlauth_fetch_deinit(&ctx->ufetch);
	}
	return TRUE;
}

static int cmd_urlfetch_transfer_literal(struct client_command_context *cmd)
{
	struct client *client = cmd->client;
	struct cmd_urlfetch_context *ctx =
		(struct cmd_urlfetch_context *)cmd->context;
	const unsigned char *data;
	size_t size;
	int ret = 1;

	/* are we in the middle of an urlfetch literal? */
	if (ctx->input == NULL)
		return 1;

	/* flush output to client if buffer is filled above optimum */
	if (o_stream_get_buffer_used_size(client->output) >=
	    CLIENT_OUTPUT_OPTIMAL_SIZE) {
		if ((ret = o_stream_flush(client->output)) <= 0)
			return ret;
	}

	/* transfer literal to client */
	o_stream_set_max_buffer_size(client->output, 4096);
	while (i_stream_read_data(ctx->input, &data, &size, 0) > 0) {
		if ((ret = o_stream_send(client->output, data, size)) < 0)
			break;
		i_stream_skip(ctx->input, ret);
	}
	o_stream_set_max_buffer_size(client->output, (size_t)-1);

	if (ret < 0 || ctx->input->stream_errno != 0) {
		/* fatal error */
		return -1;
	}

	if (!ctx->input->eof)
		return 0;

	/* finished */
	i_stream_unref(&ctx->input);
	return 1;
}

static bool cmd_urlfetch_continue(struct client_command_context *cmd)
{
	struct client *client = cmd->client;
	struct cmd_urlfetch_context *ctx =
		(struct cmd_urlfetch_context *)cmd->context;
	int ret = 1;

	if (cmd->cancel)
		return cmd_urlfetch_cancel(cmd);

	i_assert(client->output_cmd_lock == NULL ||
		 client->output_cmd_lock == cmd);

	/* finish a pending literal transfer */
	ret = cmd_urlfetch_transfer_literal(cmd);
	if (ret < 0) {
		ctx->failed = TRUE;
		cmd_urlfetch_finish(cmd);
		return TRUE;
	}
	if (ret == 0) {
		/* not finished; apparently output blocked again */
		return FALSE;
	}

	if (ctx->extended)
		client_send_line(client, ")");
	else
		client_send_line(client, "");

	if (imap_urlauth_fetch_continue(ctx->ufetch)) {
		/* waiting for imap urlauth service */
		cmd->state = CLIENT_COMMAND_STATE_WAIT_EXTERNAL;
		cmd->func = cmd_urlfetch_cancel;
		client->output_cmd_lock = NULL;

		/* retrieve next url */
		return FALSE;
	}

	/* finished */
	cmd_urlfetch_finish(cmd);
	return TRUE;
}

static int cmd_urlfetch_url_sucess(struct client_command_context *cmd,
				   struct imap_urlauth_fetch_reply *reply)
{
	struct cmd_urlfetch_context *ctx = cmd->context;
	int ret;

	if ((reply->flags & IMAP_URLAUTH_FETCH_FLAG_EXTENDED) == 0) {
		/* simple */
		ctx->extended = FALSE;

		client_send_line(cmd->client, t_strdup_printf(
			"* URLFETCH %s {%"PRIuUOFF_T"}",
			reply->url, reply->size));
		i_assert(reply->size == 0 || reply->input != NULL);
	} else {
		string_t *response = t_str_new(256);
		bool metadata = FALSE;

		/* extended */
		ctx->extended = TRUE;

		str_append(response, "* URLFETCH ");
		str_append(response, reply->url);
		str_append(response, " (");
		if ((reply->flags & IMAP_URLAUTH_FETCH_FLAG_BODYPARTSTRUCTURE) != 0 &&
		    reply->bodypartstruct != NULL) {
			str_append(response, "BODYPARTSTRUCTURE (");
			str_append(response, reply->bodypartstruct);
			str_append_c(response, ')');
			metadata = TRUE;
		}
		if ((reply->flags & IMAP_URLAUTH_FETCH_FLAG_BODY) != 0 ||
		    (reply->flags & IMAP_URLAUTH_FETCH_FLAG_BINARY) != 0) {
			if (metadata)
				str_append_c(response, ' ');
			if ((reply->flags & IMAP_URLAUTH_FETCH_FLAG_BODY) != 0) {
				str_append(response, "BODY ");
			} else {
				str_append(response, "BINARY ");
				if (reply->binary_has_nuls)
					str_append_c(response, '~');
			}
			str_printfa(response, "{%"PRIuUOFF_T"}", reply->size);
			i_assert(reply->size == 0 || reply->input != NULL);
		} else {
			str_append_c(response, ')');
			ctx->extended = FALSE;
		}

		client_send_line(cmd->client, str_c(response));
	}

	if (reply->input != NULL) {
		ctx->input = reply->input;
		i_stream_ref(reply->input);

		ret = cmd_urlfetch_transfer_literal(cmd);
		if (ret < 0) {
			ctx->failed = TRUE;
			cmd_urlfetch_finish(cmd);
			return -1;
		}
		if (ret == 0) {
			/* not finished; apparently output blocked */
			cmd->state = CLIENT_COMMAND_STATE_WAIT_OUTPUT;
			cmd->func = cmd_urlfetch_continue;
			cmd->client->output_cmd_lock = cmd;
			return 0;
		}

		if (ctx->extended)
			client_send_line(cmd->client, ")");
		else
			client_send_line(cmd->client, "");
	}
	return 1;
}

static int
cmd_urlfetch_url_callback(struct imap_urlauth_fetch_reply *reply,
			  bool last, void *context)
{
	struct client_command_context *cmd = context;
	struct cmd_urlfetch_context *ctx = cmd->context;
	int ret;

	if (cmd->cancel)
		return 1;

	if (reply == NULL) {
		/* fatal failure */
		last = TRUE;
		ctx->failed = TRUE;
	} else if (reply->succeeded) {
		/* URL fetch succeeded */
		if ((ret = cmd_urlfetch_url_sucess(cmd, reply)) <= 0)
			return ret;
	} else {
		/* URL fetch failed */
		client_send_line(cmd->client,
			t_strdup_printf("* URLFETCH %s NIL", reply->url));
		if (reply->error != NULL) {
			client_send_line(cmd->client, t_strdup_printf(
				"* NO %s.", reply->error));
		}
	}

	if (last && cmd->state == CLIENT_COMMAND_STATE_WAIT_EXTERNAL) {
		cmd_urlfetch_finish(cmd);
		client_command_free(&cmd);
	}
	return 1;
}

static int
cmd_urlfetch_parse_arg(struct client_command_context *cmd,
		       const struct imap_arg *arg,
		       struct cmd_urlfetch_url *ufurl_r)
{
	enum imap_urlauth_fetch_flags url_flags = 0;
	const struct imap_arg *params;
	const char *url_text;

	if (imap_arg_get_list(arg, &params))
		url_flags |= IMAP_URLAUTH_FETCH_FLAG_EXTENDED;
	else
		params = arg;

	/* read url */
	if (!imap_arg_get_astring(params++, &url_text)) {
		client_send_command_error(cmd, "Invalid arguments.");
		return -1;
	}
	ufurl_r->url = t_strdup(url_text);
	if (url_flags == 0)
		return 0;

	while (!IMAP_ARG_IS_EOL(params)) {
		const char *fetch_param;

		if (!imap_arg_get_atom(params++, &fetch_param)) {
			client_send_command_error(cmd,
				"Invalid URL fetch parameter.");
			return -1;
		}

		if (strcasecmp(fetch_param, "BODY") == 0)
			url_flags |= IMAP_URLAUTH_FETCH_FLAG_BODY;
		else if (strcasecmp(fetch_param, "BINARY") == 0)
			url_flags |= IMAP_URLAUTH_FETCH_FLAG_BINARY;
		else if (strcasecmp(fetch_param, "BODYPARTSTRUCTURE") == 0)
			url_flags |= IMAP_URLAUTH_FETCH_FLAG_BODYPARTSTRUCTURE;
		else {
			client_send_command_error(cmd,
				t_strdup_printf("Unknown URL fetch parameter: %s",
						fetch_param));
			return -1;
		}
	}

	if ((url_flags & IMAP_URLAUTH_FETCH_FLAG_BODY) != 0 &&
	    (url_flags & IMAP_URLAUTH_FETCH_FLAG_BINARY) != 0) {
		client_send_command_error(cmd,
			"URL cannot have both BODY and BINARY fetch parameters.");
		return -1;
	}

	if (url_flags == IMAP_URLAUTH_FETCH_FLAG_EXTENDED)
		url_flags |= IMAP_URLAUTH_FETCH_FLAG_BODY;
	ufurl_r->flags = url_flags;
	return 0;
}

bool cmd_urlfetch(struct client_command_context *cmd)
{
	struct client *client = cmd->client;
	struct cmd_urlfetch_context *ctx;
	ARRAY(struct cmd_urlfetch_url) urls;
	const struct cmd_urlfetch_url *url;
	const struct imap_arg *args;
	struct cmd_urlfetch_url *ufurl;
	int ret;

	if (client->urlauth_ctx == NULL) {
		client_send_command_error(cmd, "URLAUTH disabled.");
		return TRUE;
	}

	if (!client_read_args(cmd, 0, 0, &args))
		return FALSE;

	if (IMAP_ARG_IS_EOL(args)) {
		client_send_command_error(cmd, "Invalid arguments.");
		return TRUE;
	}
	
	t_array_init(&urls, 32);

	/* parse url arguments and group them per userid */
	for (; !IMAP_ARG_IS_EOL(args); args++) {
		ufurl = array_append_space(&urls);
		if (cmd_urlfetch_parse_arg(cmd, args, ufurl) < 0)
			return TRUE;
	}
	cmd->context = ctx = p_new(cmd->pool, struct cmd_urlfetch_context, 1);
	cmd->func = cmd_urlfetch_cancel;
	cmd->state = CLIENT_COMMAND_STATE_WAIT_INPUT;

	ctx->ufetch = imap_urlauth_fetch_init(client->urlauth_ctx,
					      cmd_urlfetch_url_callback, cmd);

	ret = 1;
	array_foreach(&urls, url) {
		ret = imap_urlauth_fetch_url(ctx->ufetch, url->url, url->flags);
		if (ret < 0) {
			/* fatal error */
			ctx->failed = TRUE;
			break;
		}
	}

	if (ret != 0 && cmd->client->output_cmd_lock != cmd) {
		/* finished */
		cmd_urlfetch_finish(cmd);
		return TRUE;
	}

	if (cmd->client->output_cmd_lock != cmd)
		cmd->state = CLIENT_COMMAND_STATE_WAIT_EXTERNAL;
	return FALSE;	
}