/*
 * Copyright (C) 2016 Paul Kocialkowski <contact@paulk.fr>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/stat.h>

#include "lg-downloader.h"
#include "gpt.h"
#include "download.h"
#include "usb.h"

static int usage_print(void)
{
	printf("Usage: lg-downloader [OPTIONS] [OPERATION]\n\n"
		"Options:\n"
		" -h                             help\n"
		" -v                             verbose\n"
		" -w                             wait for device\n\n"
		"Operations:\n"
		" reboot                         reboot device\n"
		" erase [partition] [filename]   erase a partition\n"
		" flash [partition] [filename]   flash a file to partition\n"
		" dump [partition] [filename]    dump a partition to file\n"
		" partitions                     list partitions\n");
}

static int arguments_parse(struct context *context, int argc, char *argv[])
{
	int i;

	if (context == NULL || argc < 2)
		return -1;

	for (i = 1; i < argc; i++) {
		if (strcmp(argv[i], "reboot") == 0) {
			if (context->operation)
				return -1;

			context->operation = 'r';
		} else if (strcmp(argv[i], "erase") == 0) {
			if (context->operation || argc <= (i + 1))
				return -1;

			context->operation = 'e';
			context->partition = argv[++i];
			context->filename = argv[++i];
		} else if (strcmp(argv[i], "flash") == 0) {
			if (context->operation || argc <= (i + 2))
				return -1;

			context->operation = 'f';
			context->partition = argv[++i];
			context->filename = argv[++i];
		} else if (strcmp(argv[i], "dump") == 0) {
			if (context->operation || argc <= (i + 2))
				return -1;

			context->operation = 'd';
			context->partition = argv[++i];
			context->filename = argv[++i];
		} else if (strcmp(argv[i], "partitions") == 0) {
			if (context->operation)
				return -1;

			context->operation = 'p';
		} else if (strcmp(argv[i], "-h") == 0) {
			if (context->operation)
				return -1;

			context->operation = 'h';
		} else if (strcmp(argv[i], "-v") == 0) {
			context->verbose = 1;
		} else if (strcmp(argv[i], "-w") == 0) {
			context->wait = 1;
		} else {
			return -1;
		}
	}

	if (!context->operation)
		return -1;

	return 0;
}

int erase(struct context *context)
{
	int rc;

	rc = gpt_partition_find(context);
	if (rc < 0) {
		fprintf(stderr, "Finding GPT partition failed\n");
		return -1;
	}

	if (context->verbose)
		printf("Erasing partition with address %d blocks and length %d bytes\n", (int) context->address, (int) context->length);

	rc = download_erase(context, context->address, context->length);
	if (rc < 0)
		return -1;

	return 0;
}

int flash(struct context *context)
{
	void *buffer = NULL;
	off_t address;
	size_t length;
	size_t count;
	size_t chunk;
	uint8_t *p;
	size_t i;
	struct stat st;
	int fd = -1;
	int rc;

	if (context == NULL || context->filename == NULL)
		return -1;

	rc = gpt_partition_find(context);
	if (rc < 0) {
		fprintf(stderr, "Finding GPT partition failed\n");
		goto error;
	}

	rc = stat(context->filename, &st);
	if (rc < 0) {
		fprintf(stderr, "Stating file failed\n");
		goto error;
	}

	length = st.st_size;

	if (length > context->length) {
		fprintf(stderr, "File size (%d bytes) is too big for partition (%d bytes)\n", (int) length, (int) context->length);
		goto error;
	}

	fd = open(context->filename, O_RDONLY);
	if (fd < 0) {
		fprintf(stderr, "Opening file failed\n");
		goto error;
	}

	if (context->verbose)
		printf("Writing to partition with address %d blocks and length %d bytes\n", (int) context->address, (int) length);

	buffer = calloc(1, GPT_BLOCK_SIZE);

	rc = download_start_flash_write(context);
	if (rc < 0)
		goto error;

	address = context->address;
	count = 0;

	rc = download_initialize_partition(context, address, GPT_BLOCK_SIZE);
	if (rc < 0)
		goto error;

	while (count < length) {
		chunk = (length - count) < GPT_BLOCK_SIZE ? (length - count) : GPT_BLOCK_SIZE;

		/* Writes are always block-aligned: fill the rest with zeros */
		if (chunk < GPT_BLOCK_SIZE)
			memset(buffer, 0, GPT_BLOCK_SIZE);

		p = (uint8_t *) buffer;
		i = 0;

		while (i < chunk) {
			rc = read(fd, p, chunk - i);
			if (rc <= 0)
				goto error;

			i += rc;
			p += rc;
		}

		rc = download_write_async(context, buffer, GPT_BLOCK_SIZE);
		if (rc < 0)
			goto error;

		count += chunk;
		address++;
	}

	rc = 0;
	goto complete;

error:
	rc = -1;

complete:
	if (buffer != NULL)
		free(buffer);

	if (fd >= 0)
		close(fd);

	return rc;
}

int dump(struct context *context)
{
	void *buffer = NULL;
	off_t address;
	size_t length;
	size_t count;
	size_t chunk;
	uint8_t *p;
	size_t c;
	int fd = -1;
	int rc;

	if (context == NULL || context->filename == NULL)
		return -1;

	rc = gpt_partition_find(context);
	if (rc < 0) {
		fprintf(stderr, "Finding GPT partition failed\n");
		goto error;
	}

	fd = open(context->filename, O_WRONLY | O_CREAT | O_TRUNC, 644);
	if (fd < 0) {
		fprintf(stderr, "Opening file failed\n");
		goto error;
	}

	if (context->verbose)
		printf("Dumping partition with address %d blocks and length %d bytes\n", (int) context->address, (int) context->length);

	buffer = calloc(1, GPT_BLOCK_SIZE);

	address = context->address;
	length = context->length;
	count = 0;

	rc = download_ready_read(context, address, GPT_BLOCK_SIZE);
	if (rc < 0)
		goto error;

	while (count < length) {
		chunk = (length - count) < GPT_BLOCK_SIZE ? (length - count) : GPT_BLOCK_SIZE;

		if (chunk != GPT_BLOCK_SIZE) {
			rc = download_ready_read(context, address, chunk);
			if (rc < 0)
				goto error;
		}

		rc = download_read(context, buffer, chunk);
		if (rc < 0)
			goto error;

		c = 0;
		p = (uint8_t *) buffer;

		while (c < chunk) {
			rc = write(fd, p, chunk - c);
			if (rc <= 0)
				goto error;

			c += rc;
			p += rc;
		}

		count += chunk;
		address++;
	}

	rc = 0;
	goto complete;

error:
	rc = -1;

complete:
	if (buffer != NULL)
		free(buffer);

	if (fd >= 0)
		close(fd);

	return rc;
}

static int partitions(struct context *context)
{
	int rc;

	rc = gpt_partitions_print(context);
	if (rc < 0)
		return -1;

	return 0;
}

int main(int argc, char *argv[])
{
	struct context context = { 0 };
	int rc;

	rc = arguments_parse(&context, argc, argv);
	if (rc < 0) {
		usage_print();
		return 1;
	}

	if (context.operation == 'h') {
		usage_print();
		return 0;
	}

	printf("Finding and opening USB device\n");

	rc = usb_open(&context);
	if (rc < 0)
		goto error;

	if (context.verbose)
		printf("Notifying download\n");

	rc = download_notify_start_dl(&context);
	if (rc < 0)
		goto error;

	switch (context.operation) {
		case 'r':
			printf("Rebooting device...\n");

			rc = download_reset(&context);
			if (rc < 0) {
				fprintf(stderr, "Rebooting device failed\n");
				goto error;
			}

			break;
		case 'e':
			printf("Erasing partition %s...\n", context.partition);

			/* Give the user a chance to hit Ctrl+C */
			sleep(1);

			rc = erase(&context);
			if (rc < 0) {
				fprintf(stderr, "Erasing partition failed\n");
				goto error;
			}

			break;
		case 'f':
			printf("Flashing %s to partition %s...\n", context.filename, context.partition);

			rc = flash(&context);
			if (rc < 0) {
				fprintf(stderr, "Flashing partition failed\n");
				goto error;
			}

			break;
		case 'd':
			printf("Dumping partition %s to %s...\n", context.partition, context.filename);

			rc = dump(&context);
			if (rc < 0) {
				fprintf(stderr, "Dumping partition failed\n");
				goto error;
			}

			break;
		case 'p':
			printf("Reading partitions table...\n");

			rc = partitions(&context);
			if (rc < 0) {
				fprintf(stderr, "Reading partitions table failed\n");
				goto error;
			}

			break;
		default:
			goto error;
	}

	printf("Done!\n");

	rc = 0;
	goto complete;

error:
	rc = 1;

complete:
	usb_close(&context);

	return rc;
}