/*
 * Copyright (C) 2022 Denis 'GNUtoo' Carikli <GNUtoo@cyberdimension.org>
 *
 * 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 <https://www.gnu.org/licenses/>.
 */

#include <endian.h>
#include <errno.h>
#include <fcntl.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <sys/mman.h>

struct pit_partition_v2 {
	uint32_t binary_type;
	uint32_t device_type;
	uint32_t partition_id;
	uint32_t partition_type;
	uint32_t file_system;
	uint32_t start_block;
	uint32_t nr_blocks;
	uint32_t file_offset;
	uint32_t file_size;
	char name[32];
	char file_name[32];
	char fota_name[32];
} __attribute__((__packed__));

int hexdump(void *buffer, size_t buffer_len, size_t offset, size_t len)
{
	char* b = buffer + offset;

	if (len == 1)
		printf("0x%x\n", *b);
}


size_t find_partition(void *buffer, size_t buffer_len,
		      char *name)
{
	off_t offset = 0;
	int found;

	/* Entries begin after first 28 bytes[1].
	 * [1]https://samsung-loki.github.io/samsung-docs/docs/PIT/
	 */
	offset += 28;
	if (buffer_len < offset)
		return -1;

	do {
		struct pit_partition_v2 *partition;

		if (buffer_len < offset + sizeof(struct pit_partition_v2))
			return -1;
		partition = buffer + offset;
		printf("%s: found %s\n", __func__, partition->name);
		found = !strncmp(name, partition->name, strlen(name));
		if (found)
			return offset;
		else
			offset += sizeof(struct pit_partition_v2);
	} while(!found);

	return -1;
}

size_t find_next_partition(void *buffer,
			   size_t buffer_size,
			   size_t offset)
{
	struct pit_partition_v2 *partition;

	size_t start_offset = 28;
	size_t part_size = sizeof(struct pit_partition_v2);

	if (offset > 28) {
		uint32_t nr_prev_partitions = ((offset - 28) / part_size);
		if ((offset - 28) % part_size)
			start_offset += ((nr_prev_partitions + 1) * part_size);
		else
			start_offset += (nr_prev_partitions * part_size);
	}

	if (buffer_size < (start_offset + part_size))
		return -1;

	return start_offset + part_size;
}

int grow_partition(void *buffer, size_t buffer_size, char *name,
		   size_t extra_bytes)
{
	struct pit_partition_v2 *partition;
	size_t offset;
	size_t partition_nr_blocks;
	size_t nr_extra_blocks;

	if (extra_bytes % 512)
		nr_extra_blocks = extra_bytes / 512;
	else
		nr_extra_blocks = (extra_bytes / 512) + 1;

	offset = find_partition(buffer, buffer_size, name);
	if (offset == -1) {
		printf("%s: \"%s\" partition not found\n",
		       __func__, name);
		return -1;
	}
	partition = buffer + offset;

	printf("%s: growing \"%s\" from %d to %d blocks\n",
	       __func__, partition->name,
	       le32toh(partition->nr_blocks),
	       le32toh(partition->nr_blocks) + nr_extra_blocks);
	partition_nr_blocks = le32toh(partition->nr_blocks) + nr_extra_blocks;
	partition->nr_blocks = htole32(partition_nr_blocks);

	do {
		size_t new_offset = 0;

		offset = find_next_partition(buffer, buffer_size, offset);
		if (offset == -1)
			break;

		partition = buffer + offset;

		if (partition->nr_blocks == 0)
			continue;

		new_offset = le32toh(partition->start_block) + nr_extra_blocks;

		printf("%s: moving \"%s\" from 0x%x to 0x%x\n",
		       __func__,
		       partition->name,
		       le32toh(partition->start_block),
		       new_offset);
		partition->start_block = new_offset;
	}
	while (true);

	return 0;
}

int main()
{
	int fd;
	int rc;

	void *orig;
	off_t orig_size;
	off_t offset;

	fd = open("16G.pit", O_RDWR);
	if (fd == -1) {
		rc = errno;
		printf("%s: open failed with error %d: %s\n",
		       __func__, rc, strerror(rc));
		return rc;
	}

	offset = lseek(fd, 0, SEEK_END);
	if (offset == -1) {
		rc = errno;
		printf("%s: lseek to SEEK_END failed with error %d: %s\n",
		       __func__, rc, strerror(rc));
		return rc;
	}

	orig_size = offset;
	offset = lseek(fd, 0, SEEK_SET);
	if (offset == -1) {
		rc = errno;
		printf("%s: lseek to 0 failed with error %d: %s\n",
		       __func__, rc, strerror(rc));
		return rc;
	}

	orig = mmap(NULL, orig_size, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0);
	if (orig == MAP_FAILED) {
		rc = errno;
		printf("%s: mmap failed with error %d: %s\n",
		       __func__, rc, strerror(rc));
		return rc;
	}

	rc = grow_partition(orig, orig_size, "FACTORYFS", 512*1024*1024);
	if (rc == -1) {
		printf("%s: modify_partitions failed with error %d\n",
		       __func__, rc);
		return rc;
	}

	rc = munmap(orig, orig_size);
	if (rc == -1) {
		rc = errno;
		printf("%s: munmap failed with error %d: %s\n",
		       __func__, rc, strerror(rc));
		return rc;
	}

	rc = close(fd);
	if (rc == -1) {
		rc = errno;
		printf("%s: close failed with error %d: %s\n",
		       __func__, rc, strerror(rc));
		return rc;
	}
	printf("DONE\n");
}