/*
 * Copyright (c) 2009-2017 Satlab Aps <satlab@satlab.com>
 * All rights reserved. Do not distribute without permission.
 */

/* Get the GNU version of basename(3) */
#define _GNU_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#include <csp/csp.h>
#include <csp/arch/csp_time.h>

#include <btp/client.h>
#include <btp/error.h>
#include <btp/types.h>

#include <slash/slash.h>

#define DEFAULT_NODE		3
#define DEFAULT_PORT		13
#define DEFAULT_BACKEND		"file"
#define DEFAULT_TIMEOUT		30000
#define DEFAULT_BSIZE		60
#define DEFAULT_SIZE		-1
#define DEFAULT_OFFSET		0
#define DEFAULT_COUNT		25
#define DEFAULT_ATTEMPTS	10
#define DEFAULT_LIMIT		0

#define PROGRESS_WIDTH		40U

enum cmd_type {
	CMD_NONE,
	CMD_PUSH,
	CMD_PULL,
	CMD_LIST,
	CMD_RM,
	CMD_MV,
	CMD_CP,
	CMD_SHELL,
};

struct progress_state {
	int sizechars;
	uint32_t last;
	uint32_t lt;
	uint32_t bps;
	uint32_t last_print;
	uint32_t last_print_speed;
	uint32_t remain;
	bool quiet;
	struct slash *slash;
	bool abort;
};

#define for_each_attempt(_t, _a)	for (_t = 0; _t++ < _a;)
#define is_error(_r, _t, _a)	((_r != BTP_EOK && _r != BTP_ETIMEDOUT) || (_r == BTP_ETIMEDOUT && _t >= _a))

static inline int min(int a, int b)
{
	return a < b ? a : b;
}

static char *pretty_bytes(uint32_t bps)
{
	static char buf[25];

	if (bps >= 1048576)
		snprintf(buf, sizeof(buf), "%.1f Mbit", bps/1048576.0);
	else if (bps >= 1024)
		snprintf(buf, sizeof(buf), "%.1f kbit", bps/1024.0);
	else
		snprintf(buf, sizeof(buf), "%"PRIu32" bit", bps);

	return buf;
}

static char *pretty_seconds(uint32_t seconds)
{
	static char buf[25];
	uint32_t hrs = 0, min = 0, sec = 0;

	hrs = seconds / 3600;
	min = (seconds - hrs*3600) / 60;
	sec = seconds - hrs*3600 - min*60;

	if (seconds >= 3600)
		snprintf(buf, sizeof(buf), "%uh %um %us", hrs, min, sec);
	else if (seconds >= 60)
		snprintf(buf, sizeof(buf), "%um %us", min, sec);
	else
		snprintf(buf, sizeof(buf), "%us", sec);

	return buf;
}

static void pretty_progress(struct btp_context *ctx, void *arg)
{
	int i, prog, bps;
	struct progress_state *ps = arg;
	uint32_t now;

	/* Check for ^C */
	if (slash_wait_interruptible(ps->slash, 0) == '\x03') {
		printf("\nInterrupted by user\n");
		ps->abort = true;
	}

	if (ps->quiet || ps->abort)
		return;

	now = csp_get_ms();

	if ((now - ps->last_print) < 250 && ctx->next + 1 < ctx->blocks)
		return;

	ps->last_print = now;

	printf("\r\033[K;");

	if (!ps->sizechars)
		printf("%"PRIu32" %n\r", ctx->size, &ps->sizechars);

	printf("\r%3u%% [", ctx->next * 100 / ctx->blocks);
	prog = ctx->next * PROGRESS_WIDTH / ctx->blocks;
	for (i = 0; i < prog; i++)
		printf("#");
	for (i = 0; i < PROGRESS_WIDTH - prog; i++)
		printf(" ");
	printf("]");
	printf("%*"PRIu32" ", ps->sizechars,
			 min(ctx->next * ctx->block_size, ctx->size));

	if ((now - ps->last_print_speed) > 1000) {
		bps = ((ctx->progress - ps->last) * 1000 * 8) / (now - ps->lt);
		ps->bps = ps->bps ? 0.7 * ps->bps + 0.3 * bps : bps;

		ps->remain = (ctx->blocks - ctx->next) * ctx->block_size;
		ps->lt = now;
		ps->last_print_speed = ps->lt;
		ps->last = ctx->progress;
	}

	if (!ctx->complete && ps->bps) {
		printf("%s/s", pretty_bytes(ps->bps));
		if (ps->bps)
			printf(" eta %s ", pretty_seconds((ps->remain * 8) / ps->bps));
	}

	fflush(stdout);
}

static bool limit_reached(struct btp_context *ctx, uint32_t limit)
{
	if (!limit)
		return false;

	if (ctx->progress >= limit)
		return true;

	return false;
}

static int cmd_btp(struct slash *slash)
{
	int ret, remain, index, c, i;
	int cmd = CMD_NONE;
	uint32_t start, stop, tries;

	bool force = false, paranoid = false, singlechar = false;
	uint8_t node = DEFAULT_NODE;
	uint8_t port = DEFAULT_PORT;
	char *backend = DEFAULT_BACKEND;
	char *remotepath = NULL;
	char *localpath = NULL;
	char pathbuf[BLOB_NAME_SIZE];
	char shell[SHELL_INP_SIZE] = "\0";
	int shellp = 0;
	uint32_t timeout = DEFAULT_TIMEOUT;
	uint32_t options = CSP_O_NONE;
	uint32_t bsize = DEFAULT_BSIZE;
	uint32_t count = DEFAULT_COUNT;
	uint32_t offset = DEFAULT_OFFSET;
	uint32_t size = DEFAULT_SIZE;
	uint32_t attempts = DEFAULT_ATTEMPTS;
	uint32_t limit = DEFAULT_LIMIT;

	struct progress_state pgstate = {};
	pgstate.slash = slash;

	struct btp_context *ctx;
				
	while ((c = slash_getopt(slash, "n:p:b:s:c:o:z:t:a:l:qrfdj")) != -1) {
		switch (c) {
		case 'n':
			node = atoi(slash->optarg);
			break;
		case 'p':
			port = atoi(slash->optarg);
			break;
		case 'b':
			backend = slash->optarg;
			break;
		case 's':
			bsize = atoi(slash->optarg);
			break;
		case 'c':
			count = atoi(slash->optarg);
			break;
		case 'o':
			offset = atoi(slash->optarg);
			break;
		case 'z':
			size = atoi(slash->optarg);
			break;
		case 't':
			timeout = atoi(slash->optarg);
			break;
		case 'a':
			attempts = atoi(slash->optarg);
			break;
		case 'l':
			limit = atoi(slash->optarg);
			break;
		case 'q':
			pgstate.quiet = true;
			break;
		case 'r':
			options = CSP_O_RDP;
			break;
		case 'f':
			force = true;
			break;
		case 'd':
			paranoid = true;
			break;
		case 'j':
			singlechar = true;
			break;
		case '?':
			if (slash->optopt == 'c')
				fprintf(stderr, "Option -%c requires an argument.\n", slash->optopt);
			/* FALLTHROUGH */
		default:
			return SLASH_EUSAGE;
		}
	}

	remain = slash->argc - slash->optind;
	index = slash->optind;

	if (index >= slash->argc) {
		return SLASH_EUSAGE;
	}

	if (!strcmp(slash->argv[index], "push") && remain == 3) {
		cmd = CMD_PUSH;
		remotepath = slash->argv[index + 2];
		if (remotepath[strlen(remotepath)-1] == '/') {
			strncpy(pathbuf, remotepath, BLOB_NAME_SIZE - 1);
			pathbuf[BLOB_NAME_SIZE - 1] = '\0';
			strncat(pathbuf, basename(slash->argv[index + 1]), BLOB_NAME_SIZE - 1);
			pathbuf[BLOB_NAME_SIZE - 1] = '\0';
			remotepath = pathbuf;
		}
		localpath = slash->argv[index + 1];
	} else if (!strcmp(slash->argv[index], "pull") && remain >= 2) {
		cmd = CMD_PULL;
		remotepath = slash->argv[index + 1];
		if (remain >= 3) {
			localpath = slash->argv[index + 2];
		} else {
			localpath = basename(slash->argv[index + 1]);
		}
	} else if ((!strcmp(slash->argv[index], "list") || !strcmp(slash->argv[index], "ls")) && remain == 2) {
		cmd = CMD_LIST;
		remotepath = slash->argv[index + 1];
	} else if ((!strcmp(slash->argv[index], "remove") || !strcmp(slash->argv[index], "rm")) && remain == 2) {
		cmd = CMD_RM;
		remotepath = slash->argv[index + 1];
	} else if ((!strcmp(slash->argv[index], "move") || !strcmp(slash->argv[index], "mv")) && remain == 3) {
		cmd = CMD_MV;
		localpath = slash->argv[index + 1];
		remotepath = slash->argv[index + 2];
	} else if ((!strcmp(slash->argv[index], "copy") || !strcmp(slash->argv[index], "cp")) && remain == 3) {
		cmd = CMD_CP;
		localpath = slash->argv[index + 1];
		remotepath = slash->argv[index + 2];
	} else if (!strcmp(slash->argv[index], "shell")) {
		cmd = CMD_SHELL;
		for (i = 0; i < remain - 1; i++) {
			if (i > 0)
				shellp = shell - strncat(shell, " ", SHELL_INP_SIZE - shellp);
			shellp = shell - strncat(shell, slash->argv[index + 1 + i], SHELL_INP_SIZE - shellp);
		}
	} else {
		return SLASH_EUSAGE;
	}

	start = csp_get_ms();

	ctx = btp_client_connect(node, port, timeout, options);
	if (!ctx) {
		fprintf(stderr, "Failed to connect to BTP server\n");
		return SLASH_EIO;
	}

	if (cmd == CMD_PULL) {
		for_each_attempt(tries, attempts) {
			ret = btp_client_download(ctx, backend, remotepath, localpath, bsize, offset, size, timeout, force, paranoid);
			if (is_error(ret, tries, attempts)) {
				fprintf(stderr, "Could not download %s: %s\n", remotepath, btp_error(ret));
				goto discon;
			} else if (ret == BTP_EOK)
				break;
		}

		if (btp_client_finished(ctx)) {
			fprintf(stderr, "File already complete\n");
			goto out;
		}

		while (!btp_client_finished(ctx) && !limit_reached(ctx, limit)) {
			for_each_attempt(tries, attempts) {
				ret = btp_client_send_status(ctx, timeout);
				if (is_error(ret, tries, attempts)) {
					fprintf(stderr, "\nCould not send status: %s\n", btp_error(ret));
					goto discon;
				} else if (pgstate.abort) {
					goto discon;
				} else if (ret == BTP_EOK)
					break;
			}

			pretty_progress(ctx, &pgstate);

			for_each_attempt(tries, attempts) {
				ret = btp_client_get_blocks(ctx, BTP_OFFSET_NEXT, count, timeout, pretty_progress, &pgstate);
				if (is_error(ret, tries, attempts)) {
					fprintf(stderr, "\nCould not get blocks: %s\n", btp_error(ret));
					goto discon;
				} else if (pgstate.abort) {
					goto discon;
				} else if (ret == BTP_EOK)
					break;
			}
		}
	} else if (cmd == CMD_PUSH) {
		for_each_attempt(tries, attempts) {
			ret = btp_client_upload(ctx, backend, remotepath, localpath, bsize, offset, size, timeout, force, paranoid);
			if (is_error(ret, tries, attempts)) {
				fprintf(stderr, "Could not push %s: %s\n", remotepath, btp_error(ret));
				goto discon;
			} else if (ret == BTP_EOK)
				break;
		}
		
		if (btp_client_finished(ctx)) {
			fprintf(stderr, "File already complete\n");
			goto out;
		}

		while (!btp_client_finished(ctx) && !limit_reached(ctx, limit)) {
			for_each_attempt(tries, attempts) {
				ret = btp_client_request_status(ctx, timeout);
				if (is_error(ret, tries, attempts)) {
					fprintf(stderr, "\nCould not get status: %s\n", btp_error(ret));
					goto discon;
				} else if (pgstate.abort) {
					goto discon;
				} else if (ret == BTP_EOK)
					break;
			}

			pretty_progress(ctx, &pgstate);

			for_each_attempt(tries, attempts) {
				ret = btp_client_send_blocks(ctx, BTP_OFFSET_NEXT, count, timeout, pretty_progress, &pgstate);
				if (is_error(ret, tries, attempts)) {
					fprintf(stderr, "\nCould not send blocks: %s\n", btp_error(ret));
					goto discon;
				} else if (pgstate.abort) {
					goto discon;
				} else if (ret == BTP_EOK)
					break;
			}
		}
	} else if (cmd == CMD_LIST) {
		for_each_attempt(tries, attempts) {
			ret = btp_client_list(ctx, count, backend, remotepath, timeout);
			if (is_error(ret, tries, attempts)) {
				fprintf(stderr, "Could not list %s: %s\n", remotepath, btp_error(ret));
				break;
			} else if (ret == BTP_EOK)
				break;
		}
		goto discon;
	} else if (cmd == CMD_RM) {
		for_each_attempt(tries, attempts) {
			ret = btp_client_remove(ctx, backend, remotepath, timeout);
			if (is_error(ret, tries, attempts)) {
				fprintf(stderr, "Could not remove %s: %s\n", remotepath, btp_error(ret));
				break;
			} else if (ret == BTP_EOK)
				break;
		}
		goto discon;
	} else if (cmd == CMD_MV) {
		for_each_attempt(tries, attempts) {
			ret = btp_client_move(ctx, backend, localpath, remotepath, timeout);
			if (is_error(ret, tries, attempts)) {
				fprintf(stderr, "Could not move %s to %s: %s\n", localpath, remotepath, btp_error(ret));
				break;
			} else if (ret == BTP_EOK)
				break;
		}
		goto discon;
	} else if (cmd == CMD_CP) {
		for_each_attempt(tries, attempts) {
			ret = btp_client_copy(ctx, backend, localpath, remotepath, timeout);
			if (is_error(ret, tries, attempts)) {
				fprintf(stderr, "Could not copy %s to %s: %s\n", localpath, remotepath, btp_error(ret));
				break;
			} else if (ret == BTP_EOK)
				break;
		}
		goto discon;
	} else if (cmd == CMD_SHELL) {
		for_each_attempt(tries, attempts) {
			ret = btp_client_shell(ctx, count, shell, singlechar, timeout);
			if (is_error(ret, tries, attempts)) {
				fprintf(stderr, "Could not execute %s: %s\n", shell, btp_error(ret));
				break;
			} else if (ret == BTP_EOK)
				break;
		}
		goto discon;
	} else {
		printf("Unknown command: %d\n", cmd);
		goto discon;
	}

	if (!pgstate.quiet && !pgstate.abort) {
		stop = csp_get_ms();
		pretty_progress(ctx, &pgstate);
		printf("\n%s/s (%"PRIu32" bytes in %.3fs)\n", 
		       pretty_bytes((ctx->progress * 8)/((stop - start)/1000.0)),
		       ctx->progress, (stop - start)/1000.0);
	}

out:
	for_each_attempt(tries, attempts) {
		ret = btp_client_complete(ctx, timeout);
		if (ret == BTP_ESTALE) {
			printf("WARNING: checksum mismatch!\n");
			break;
		} else if (ret == BTP_EOK)
			break;
	}
discon:
	btp_client_disconnect(ctx);

	return SLASH_SUCCESS;
}
slash_command(btp, cmd_btp, "<command> <remote file> [local file]",
	      "Block Transfer Protocol client\n"
	      "\n"
	      "<command> is one of push, pull, list, rm, mv, cp or shell.\n"
	      "\n"
	      "Options:\n"
	      " -n NODE,\tConnect to NODE\n" 
	      " -p PORT,\tConnect to PORT\n"
	      " -b BACKEND,\tUse BACKEND as backend\n"
	      " -s BLOCKSIZE,\tTransfer data in BLOCKSIZE sized blocks\n"
	      " -c COUNT,\tRequest COUNT blocks at a time\n"
	      " -o OFFSET,\tStart OFFSET bytes into file\n"
	      " -z SIZE,\tDownload only SIZE bytes\n"
	      " -t TIMEOUT,\tSet timeout to TIMEOUT ms\n"
	      " -a ATTEMPTS,\tRetry connection up to ATTEMPTS tries\n"
	      " -l LIMIT,\tStop after LIMIT bytes downloaded\n"
	      " -q,\t\tQuiet mode. Do not print progress bar\n"
	      " -r,\t\tEnable RDP for transfer\n"
	      " -f,\t\tDelete data and map file before starting transfer\n"
	      " -d,\t\tParanoid mode. Always sync blocks to disk\n"
	      " -j,\t\tSingle character terminal mode\n");
