#define _POSIX_C_SOURCE 200112L
#define _XOPEN_SOURCE 600

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <assert.h>
#include <string.h>
#include <argp.h>
#include <inttypes.h>
#include <err.h>
#include <unistd.h>

#include "version.h"
#include "libutils.h"
#include "libdevs.h"
#include "libflow.h"

/* Argp's global variables. */
const char *argp_program_version = "F3 BREW " F3_STR_VERSION;

/* Arguments. */
static char adoc[] = "<DISK_DEV>";

/* The capital "E" in "REad" in the string below is not a typo.
 * It shows from where the name B-RE-W comes.
 */
static char doc[] = "F3 Block REad and Write -- assess the media of "
	"a block device writing blocks, resetting the drive, and "
	"reading the blocks back";

static struct argp_option options[] = {
	{"debug",		'd',	NULL,		OPTION_HIDDEN,
		"Enable debugging; only needed if none --debug-* option used",
		1},
	{"debug-real-size",	'a',	"SIZE_BYTE",	OPTION_HIDDEN,
		"Real size of the emulated drive",	0},
	{"debug-fake-size",	'f',	"SIZE_BYTE",	OPTION_HIDDEN,
		"Fake size of the emulated drive",	0},
	{"debug-wrap",		'm',	"N",		OPTION_HIDDEN,
		"Wrap parameter of the emulated drive",	0},
	{"debug-block-order",	'b',	"ORDER",	OPTION_HIDDEN,
		"Block size of the emulated drive is 2^ORDER Bytes",	0},
	{"debug-cache-order",	'c',	"ORDER",	OPTION_HIDDEN,
		"Cache size of the emulated drive is 2^ORDER blocks",	0},
	{"debug-strict-cache",	'o',	NULL,		OPTION_HIDDEN,
		"Force the cache to be strict",				0},
	{"debug-keep-file",	'k',	NULL,		OPTION_HIDDEN,
		"Don't remove file used for emulating the drive",	0},
	{"reset-type",		's',	"TYPE",		0,
		"Reset method to use during the probe",		2},
	{"start-at",		'h',	"BLOCK",	0,
		"Where test begins; the default is block zero",	0},
	{"end-at",		'e',	"BLOCK",	0,
		"Where test ends; the default is the very last block",	0},
	{"fix-cmd",		'x',	NULL,		0,
		"Show how to call f3fix to fix the drive",	0},
	{"max-read-rate",	'r',	"KB/s",		0,
		"Maximum read rate",					0},
	{"max-write-rate",	'w',	"KB/s",		0,
		"Maximum write rate",					0},
	{"show-progress",	'p',	"NUM",		0,
		"Show progress if NUM is not zero",			0},
	{"do-not-write",	'W',	NULL,		0,
		"Do not write blocks",				0},
	{"do-not-read",		'R',	NULL,		0,
		"Do not read blocks",				0},
	{ 0 }
};

struct args {
	char		*filename;

	/* Debugging options. */
	bool		debug;
	bool		keep_file;

	/* Behavior options. */
	enum reset_type	reset_type;
	bool test_write;
	bool test_read;
	bool fix_cmd;

	/* Geometry. */
	uint64_t	real_size_byte;
	uint64_t	fake_size_byte;
	int		wrap;
	unsigned int	block_order;
	int		cache_order;
	int		strict_cache;

	/* What to do. */
	uint64_t	first_block;
	uint64_t	last_block;

	/* Flow control and reporting. */
	long		max_read_rate;
	long		max_write_rate;
	int		show_progress;
};

static error_t parse_opt(int key, char *arg, struct argp_state *state)
{
	struct args *args = state->input;
	long long ll;

	switch (key) {
	case 'd':
		args->debug = true;
		break;

	case 'a':
		ll = arg_to_ll_bytes(state, arg);
		if (ll < 0)
			argp_error(state,
				"Real size must be greater or equal to zero");
		args->real_size_byte = ll;
		args->debug = true;
		break;

	case 'f':
		ll = arg_to_ll_bytes(state, arg);
		if (ll < 0)
			argp_error(state,
				"Fake size must be greater or equal to zero");
		args->fake_size_byte = ll;
		args->debug = true;
		break;

	case 'm':
		ll = arg_to_ll_bytes(state, arg);
		if (ll < 0 || ll >= 64)
			argp_error(state,
				"Wrap must be in the interval [0, 63]");
		args->wrap = ll;
		args->debug = true;
		break;

	case 'b':
		ll = arg_to_ll_bytes(state, arg);
		if (ll != 0 && (ll < SECTOR_ORDER || ll > MEGABYTE_ORDER))
			argp_error(state,
				"Block order must be in the interval [%i, %i] or be zero",
				SECTOR_ORDER, MEGABYTE_ORDER);
		args->block_order = ll;
		args->debug = true;
		break;

	case 'c':
		ll = arg_to_ll_bytes(state, arg);
		if (ll < -1 || ll > 64)
			argp_error(state,
				"Cache order must be in the interval [-1, 64]");
		args->cache_order = ll;
		args->debug = true;
		break;

	case 'o':
		args->strict_cache = true;
		args->debug = true;
		break;

	case 'k':
		args->keep_file = true;
		args->debug = true;
		break;

	case 's':
		ll = arg_to_ll_bytes(state, arg);
		if (ll < 0 || ll >= RT_MAX)
			argp_error(state,
				"Reset type must be in the interval [0, %i]",
				RT_MAX - 1);
		args->reset_type = ll;
		break;

	case 'h':
		ll = arg_to_ll_bytes(state, arg);
		if (ll < 0)
			argp_error(state,
				"The first block must be greater or equal to zero");
		args->first_block = ll;
		break;

	case 'e':
		ll = arg_to_ll_bytes(state, arg);
		if (ll < 0)
			argp_error(state,
				"The last block must be greater or equal to zero");
		args->last_block = ll;
		break;

	case 'r':
		ll = arg_to_ll_bytes(state, arg);
		if (ll <= 0)
			argp_error(state,
				"KB/s must be greater than zero");
		args->max_read_rate = ll;
		break;

	case 'w':
		ll = arg_to_ll_bytes(state, arg);
		if (ll <= 0)
			argp_error(state,
				"KB/s must be greater than zero");
		args->max_write_rate = ll;
		break;

	case 'p':
		args->show_progress = !!arg_to_ll_bytes(state, arg);
		break;

	case 'W':
		args->test_write = false;
		break;

	case 'R':
		args->test_read = false;
		break;

	case 'x':
		args->fix_cmd = true;
		break;

	case ARGP_KEY_INIT:
		args->filename = NULL;
		break;

	case ARGP_KEY_ARG:
		if (args->filename)
			argp_error(state,
				"Wrong number of arguments; only one is allowed");
		args->filename = arg;
		break;

	case ARGP_KEY_END:
		if (!args->filename)
			argp_error(state,
				"The disk device was not specified");
		if (args->debug &&
			!dev_param_valid(args->real_size_byte,
				args->fake_size_byte, args->wrap,
				args->block_order))
			argp_error(state,
				"The debugging parameters are not valid");

		if (args->first_block > args->last_block)
			argp_error(state,
				"The first block parameter must be less or equal to the last block parameter. They are now: first_block=%"
				PRIu64 " > last_block=%" PRIu64,
				args->first_block, args->last_block);

		break;

	default:
		return ARGP_ERR_UNKNOWN;
	}
	return 0;
}

static struct argp argp = {options, parse_opt, adoc, doc, NULL, NULL, NULL};

static void write_blocks(struct device *dev, struct flow *fw,
	uint64_t first_block, uint64_t last_block)
{
	const unsigned int block_size = dev_get_block_size(dev);
	const unsigned int block_order = dev_get_block_order(dev);
	uint64_t offset = first_block << block_order;
	uint64_t first_pos = first_block;
	struct dynamic_buffer dbuf;

	dbuf_init(&dbuf);

	start_measurement(fw);
	while (first_pos <= last_block) {
		const uint64_t max_blocks_to_write = last_block - first_pos + 1;
		uint64_t blocks_to_write =
			MIN(get_rem_chunk_blocks(fw), max_blocks_to_write);
		size_t buf_len = blocks_to_write << block_order;
		char *buffer, *stamp_blk;
		uint64_t pos, next_pos;

		buffer = dbuf_get_buf(&dbuf, block_order, &buf_len);
		blocks_to_write = buf_len >> block_order;
		assert(blocks_to_write > 0);
		next_pos = first_pos + blocks_to_write;

		stamp_blk = buffer;
		for (pos = first_pos; pos < next_pos; pos++) {
			fill_buffer_with_block(stamp_blk, block_order, offset, 0);
			stamp_blk += block_size;
			offset += block_size;
		}

		if (dev_write_blocks(dev, buffer, first_pos, next_pos - 1)) {
			clear_progress(fw);
			warn("Failed to write blocks from 0x%" PRIx64
				" to 0x%" PRIx64, first_pos, next_pos - 1);
		}

		measure(fw, blocks_to_write, NULL);
		first_pos = next_pos;
	}
	end_measurement(fw);
	dbuf_free(&dbuf);
}

/* XXX Properly handle return errors. */
static void test_write_blocks(struct device *dev,
	uint64_t first_block, uint64_t last_block,
	long max_write_rate, int show_progress)
{
	const unsigned int block_order = dev_get_block_order(dev);
	const uint64_t total_blocks = last_block - first_block + 1;
	struct flow fw;

	printf("Writing block%s from 0x%" PRIx64 " to 0x%" PRIx64 "... ",
		first_block != last_block ? "s" : "", first_block, last_block);
	fflush(stdout);

	init_flow(&fw, block_order, total_blocks, max_write_rate,
		FW_MAX_BLOCKS_PER_DELAY_NONE,
		show_progress ? printf_flush_cb : dummy_cb, 0);

	write_blocks(dev, &fw, first_block, last_block);

	printf("Done\n");
	print_avg_seq_speed(&fw, "write", false);
	printf("\n");
}

struct block_range {
	enum block_state	state;
	unsigned int		block_order;
	uint64_t		start_block;
	uint64_t		end_block;

	/* Only used by state bs_overwritten. */
	uint64_t		found_block;
};

#define INIT_UNKNOWN_RANGE(block_order) {	\
	.state = bs_unknown,			\
	.block_order = (block_order),		\
	.start_block = 0,			\
	.end_block = UINT64_MAX,		\
	.found_block = 0			\
}

static void print_block_range(const struct block_range *range)
{
	printf("[%s] from block 0x%" PRIx64 " to 0x%" PRIx64,
		block_state_to_str(range->state),
		range->start_block, range->end_block);

	switch (range->state) {
	case bs_good:
	case bs_bad:
	case bs_changed:
		break;

	case bs_overwritten:
		printf(", found block 0x%" PRIx64, range->found_block);
		break;

	default:
		assert(0);
		break;
	}
	printf("\n");
}

static inline bool is_block(uint64_t offset, unsigned int block_order)
{
	return !(((1ULL << block_order) - 1) & offset);
}

static void update_good_range(struct block_range *good_range,
	const struct block_range *range)
{
	uint64_t good_range_blocks, range_blocks;

	if (range->state != bs_good)
		return;

	good_range_blocks =
		good_range->end_block - good_range->start_block + 1;
	range_blocks = range->end_block - range->start_block + 1;
	if (range_blocks > good_range_blocks)
		*good_range = *range;
}

static void validate_block(struct flow *fw, uint64_t block,
	const char *probe_blk, unsigned int block_order,
	struct block_range *range, struct block_range *good_range,
	struct block_stats *stats)
{
	const uint64_t expected_offset = block << block_order;
	uint64_t found_offset;
	enum block_state state = validate_block_update_stats(probe_blk,
		block_order, expected_offset, &found_offset, 0, stats);
	const uint64_t found_block = found_offset >> block_order;
	bool push_range;

	assert(is_block(found_offset, block_order));

	push_range = (range->state != state) || (
			state == bs_overwritten
			&& (
				(block - range->start_block)
				!=
				(found_block - range->found_block)
			)
		);

	if (push_range) {
		if (range->state != bs_unknown) {
			clear_progress(fw);
			print_block_range(range);
			update_good_range(good_range, range);
		}
		range->state = state;
		range->start_block = block;
		range->end_block = block;
		range->found_block = found_block;
	} else {
		range->end_block = block;
	}
}

static void read_blocks(struct device *dev, struct flow *fw,
	uint64_t first_block, uint64_t last_block, struct block_stats *stats,
	struct block_range *good_range)
{
	const unsigned int block_size = dev_get_block_size(dev);
	const unsigned int block_order = dev_get_block_order(dev);
	uint64_t first_pos = first_block;
	struct block_range range = INIT_UNKNOWN_RANGE(block_order);
	struct dynamic_buffer dbuf;

	dbuf_init(&dbuf);

	start_measurement(fw);
	while (first_pos <= last_block) {
		const uint64_t max_blocks_to_read = last_block - first_pos + 1;
		uint64_t blocks_to_read =
			MIN(get_rem_chunk_blocks(fw), max_blocks_to_read);
		size_t buf_len = blocks_to_read << block_order;
		char *buffer, *probe_blk;
		uint64_t pos, next_pos;

		buffer = dbuf_get_buf(&dbuf, block_order, &buf_len);
		blocks_to_read = buf_len >> block_order;
		assert(blocks_to_read > 0);
		next_pos = first_pos + blocks_to_read;

		if (dev_read_blocks(dev, buffer, first_pos, next_pos - 1)) {
			clear_progress(fw);
			warn("Failed to read blocks from 0x%" PRIx64
				" to 0x%" PRIx64, first_pos, next_pos - 1);
		}

		probe_blk = buffer;
		for (pos = first_pos; pos < next_pos; pos++) {
			validate_block(fw, pos, probe_blk, block_order,
				&range, good_range, stats);
			probe_blk += block_size;
		}

		measure(fw, blocks_to_read, NULL);
		first_pos = next_pos;
	}
	end_measurement(fw);
	dbuf_free(&dbuf);

	if (range.state != bs_unknown) {
		print_block_range(&range);
		update_good_range(good_range, &range);
	} else {
		assert(first_block > last_block);
	}
}

static void print_fix_cmd(const char *filename,
	const struct block_range *good_range)
{
	const uint64_t first_1MB_block =
		MEGABYTE_SIZE >> good_range->block_order;
	const unsigned int shift = good_range->block_order - SECTOR_ORDER;
	const uint64_t first_good_sector = good_range->start_block << shift;
	const uint64_t last_good_sector =
		((good_range->end_block + 1) << shift) - 1;
	uint64_t size_byte;

	assert(MEGABYTE_ORDER >= good_range->block_order);
	assert(good_range->block_order >= SECTOR_ORDER);

	/* Set size_byte. */
	if (good_range->state != bs_good) {
		size_byte = 0;
	} else if (good_range->end_block >= first_1MB_block) {
		const uint64_t start_block =
			good_range->start_block <= first_1MB_block
				? first_1MB_block
				: good_range->start_block;
		size_byte = (good_range->end_block - start_block + 1) <<
			good_range->block_order;
	} else {
		size_byte = 0;
	}

	if (size_byte < MEGABYTE_SIZE) {
		printf("There is no good region large enough to \"fix\" this device.\n\n");
		return;
	}

	printf("You can \"fix\" this device using the following command:\n");
	if (good_range->start_block < first_1MB_block) {
		printf("f3fix --last-sec=%" PRIu64 " %s\n\n",
			last_good_sector, filename);
		return;
	}
	printf("f3fix --first-sec=%" PRIu64 " --last-sec=%" PRIu64 " %s\n\n",
		first_good_sector, last_good_sector, filename);
}

/* XXX Properly handle return errors. */
static void test_read_blocks(struct device *dev,
	uint64_t first_block, uint64_t last_block,
	long max_read_rate, int show_progress, bool fix_cmd)
{
	const unsigned int block_order = dev_get_block_order(dev);
	const uint64_t total_blocks = last_block - first_block + 1;
	struct flow fw;
	struct block_stats stats = { 0, 0, 0, 0 };
	struct block_range good_range = INIT_UNKNOWN_RANGE(block_order);

	printf("Reading block%s from 0x%" PRIx64 " to 0x%" PRIx64 ":\n",
		first_block != last_block ? "s" : "", first_block, last_block);

	init_flow(&fw, block_order, total_blocks, max_read_rate,
		FW_MAX_BLOCKS_PER_DELAY_NONE,
		show_progress ? printf_flush_cb : dummy_cb, 0);

	read_blocks(dev, &fw, first_block, last_block, &stats, &good_range);

	print_stats(&stats, block_order, "block");
	print_avg_seq_speed(&fw, "read", false);
	printf("\n");

	if (fix_cmd)
		print_fix_cmd(dev_get_filename(dev), &good_range);
}

int main(int argc, char **argv)
{
	struct args args = {
		/* Defaults. */
		.debug		= false,
		.keep_file	= false,
		.reset_type	= RT_MANUAL_USB,
		.test_write	= true,
		.test_read	= true,
		.fix_cmd	= false,
		.real_size_byte	= 2 * GIGABYTE_SIZE,
		.fake_size_byte	= 16 * GIGABYTE_SIZE,
		.wrap		= 31,
		.block_order	= SECTOR_ORDER,
		.cache_order	= -1,
		.strict_cache	= false,
		.first_block	= 0,
		.last_block	= -1ULL,
		.max_read_rate	= FW_MAX_PROCESS_RATE_NONE,
		.max_write_rate = FW_MAX_PROCESS_RATE_NONE,
		/* If stdout isn't a terminal, suppress progress. */
		.show_progress	= isatty(STDOUT_FILENO),
	};
	struct device *dev;
	unsigned int block_order;
	uint64_t very_last_block;

	/* Read parameters. */
	argp_parse(&argp, argc, argv, 0, NULL, &args);
	print_header(stdout, "brew");

	dev = args.debug
		? create_file_device(args.filename, args.real_size_byte,
			args.fake_size_byte, args.wrap, args.block_order,
			args.cache_order, args.strict_cache, args.keep_file)
		: create_block_device(args.filename, args.reset_type);
	if (!dev) {
		fprintf(stderr, "\nApplication cannot continue, finishing...\n");
		exit(1);
	}

	block_order = dev_get_block_order(dev);
	printf("Physical block size: 2^%i Byte%s\n\n",
		block_order, block_order != 0 ? "s" : "");

	very_last_block = (dev_get_size_byte(dev) >> block_order) - 1;
	if (args.first_block > very_last_block)
		args.first_block = very_last_block;
	if (args.last_block > very_last_block)
		args.last_block = very_last_block;

	if (args.test_write)
		test_write_blocks(dev, args.first_block, args.last_block,
			args.max_write_rate, args.show_progress);

	if (args.test_write && args.test_read) {
		const char *final_dev_filename;

		assert(!dev_reset(dev));
		final_dev_filename = dev_get_filename(dev);
		if (strcmp(args.filename, final_dev_filename))
			printf("\nWARNING: device `%s' moved to `%s' due to the reset\n\n",
				args.filename, final_dev_filename);
	}

	if (args.test_read)
		test_read_blocks(dev, args.first_block, args.last_block,
			args.max_read_rate, args.show_progress, args.fix_cmd);

	free_device(dev);
	return 0;
}