/*
 * Copyright (c) 2019 The Fuchsia Authors
 */

#include <common.h>
#include <amlogic/storage.h>
#include <config.h>
#include <fastboot.h>
#include <image-sparse.h>
#include <zircon_uboot/partition.h>

#include "lz4_wrapper.h"

static lbaint_t fb_zircon_sparse_write(struct sparse_storage *info,
				       lbaint_t blk, lbaint_t blkcnt,
				       const void *buffer)
{
	zircon_partition *part = (zircon_partition *)info->priv;
	uint64_t offset = blk * info->blksz;
	size_t size = blkcnt * info->blksz;
	return part->write(part, offset, buffer, size) ? 0 : blkcnt;
}

static lbaint_t fb_zircon_sparse_reserve(struct sparse_storage *info,
					 lbaint_t blk, lbaint_t blkcnt)
{
	return blkcnt;
}

static uint8_t lz4_image_decompress_buffer[4 * 1024 * 1024];

static int write_lz4_compressed_image(zircon_partition *part,
				      void *download_buffer,
				      unsigned int download_bytes)
{
	int ret = 0;
	struct lz4_frame_header lz4_header;
	size_t lz4_data_offset;
	if ((ret = check_and_extract_lz4_frame_header(
		     download_buffer, download_bytes, &lz4_header,
		     &lz4_data_offset))) {
		printf("Failed to extract lz4 frame header\n");
		return -1;
	}

	const uint8_t *compressed_end = download_buffer + download_bytes;
	const uint8_t *compressed_curr =
		(uint8_t *)download_buffer + lz4_data_offset;

	printf("Flashing lz4 compressed image to partition\n");

	size_t write_offset = 0;
	bool has_more_data = true;
	while (has_more_data) {
		size_t decompressed_size = sizeof(lz4_image_decompress_buffer);
		size_t compressed_size = compressed_end - compressed_curr;
		if ((ret = lz4_decompress(
			     &lz4_header, compressed_curr, &compressed_size,
			     lz4_image_decompress_buffer, &decompressed_size,
			     &has_more_data))) {
			printf("Failed to decompress, %d\n", ret);
			return -1;
		}
		compressed_curr += compressed_size;

		if ((ret = part->write(part, write_offset,
				       lz4_image_decompress_buffer,
				       decompressed_size))) {
			printf("Failed to write to storage\n");
			return -1;
		}
		write_offset += decompressed_size;
	}
	return 0;
}

void fb_zircon_flash_write(const char *part_name, void *download_buffer,
			   unsigned int download_bytes)
{
	// `write_sparse_image` sets fastboot_fail() on error, so start with
	// fastboot_okay() to avoid clobbering the return value.
	fastboot_okay(NULL);

	zircon_partition *part = zircon_get_fastboot_partition(part_name);
	if (!part) {
		fastboot_fail("partition not found");
		return;
	}

	if (!part->write) {
		zircon_free_partition(part);
		fastboot_fail("Device must be unlocked to run this command");
		return;
	}

	printf("Flashing '%s'\n", part_name);

	if (is_sparse_image(download_buffer)) {
		struct sparse_storage sparse = {
			.blksz = 1,
			.start = 0,
			.size = part->size,
			.priv = (void *)part,
			.write = fb_zircon_sparse_write,
			.reserve = fb_zircon_sparse_reserve,
		};

		printf("Flashing sparse image to partition %s\n", part_name);

		// sets `fastboot_fail()` on error.
		write_sparse_image(&sparse, part_name, download_buffer,
				   download_bytes);
	} else if (is_data_lz4_compressed(download_buffer, download_bytes)) {
		if ((write_lz4_compressed_image(part, download_buffer,
						download_bytes))) {
			fastboot_fail("failed to write lz4 compressed image");
		}
	} else {
		if (part->write(part, 0, download_buffer, download_bytes)) {
			printf("partition flash failed: %s\n", part_name);
			fastboot_fail("error writing the image");
		}
	}

	printf("Flashing '%s' done\n", part_name);

	zircon_free_partition(part);
}

void fb_zircon_erase(const char *part_name)
{
	zircon_partition *part = zircon_get_fastboot_partition(part_name);
	if (!part) {
		fastboot_fail("partition not found");
		return;
	}

	if (!part->erase) {
		zircon_free_partition(part);
		fastboot_fail("Device must be unlocked to run this command");
		return;
	}

	printf("Erasing '%s'\n", part_name);

	if (part->erase(part)) {
		printf("partition erase failed: %s\n", part_name);
		fastboot_fail("failed erasing from device");
	} else {
		fastboot_okay(NULL);
	}

	zircon_free_partition(part);
}