#include <stdio.h>
#include <libfdt.h>
#include <fdt.h>
#include <sys/stat.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <arpa/inet.h>

#if 0
#define debug(...) printf(__VA_ARGS__)
#else
#define debug(...)
#endif

static int reads_counter = 0;
static int fit_offset = 0;

static int firmware_fdt_read_chunk(const char *filename, char *buf, int offset, int len)
{
	int fd = 0;	/* assume stdin */
	int ret = 0;

	fd = open(filename, O_RDONLY);
	if (fd < 0)
		return errno;

	offset += fit_offset;

	struct stat st;
	stat(filename, &st);
	size_t file_size = st.st_size;

	if (file_size < offset || offset < 0)
	{
		printf("Bad offset: %d/%ld\n", offset, file_size);
		return -1;
	}
	if (file_size < offset + len)
		len = file_size - offset;

	if (offset)
		lseek(fd, offset, SEEK_SET);
	/* Loop until we have read everything */
	size_t readed = 0;
	do {
		ret = read(fd, buf + readed, len - readed);
		if (ret < 0)
		{
			ret = errno;
			break;
		}
		readed += ret;
	} while (ret != 0 && readed < len);

	/* Clean up, including closing stdin; return errno on error */
	close(fd);

	debug("Read %d bytes from %x (%x)\n", readed, offset - fit_offset, offset);

	memset(buf + readed, 0, len - readed);
	reads_counter++;
	return readed;
}

static inline const void *fdt_offset_ptr_(const void *fdt, const void *buf, int buf_offset)
{
	debug("fdt_offset_ptr_: %p + %x\n", buf, buf_offset);
	return (const char *)buf + buf_offset;
}

static inline const int fdt_offset_to_file_offset(const void *fdt, int buf_offset)
{
	debug("fdt_offset_ptr_: %x\n", buf_offset);
	return buf_offset + fdt_off_dt_struct(fdt);
}

#define FDT_ALIGN(x, a)		(((x) + (a) - 1) & ~((a) - 1))
#define FDT_TAGALIGN(x)		(FDT_ALIGN((x), FDT_TAGSIZE))

/* Do not touch any forward data, caller will validate it later */
int guess_next_tag(char * fdt, int offset, char * buf, int buf_offset)
{
	const unsigned * val = fdt_offset_ptr_(fdt, buf, buf_offset);
	unsigned tag = fdt32_to_cpu(*val);
	char * ptr = NULL;
	int offset_adjust = 0;

	switch (tag) {
		case FDT_BEGIN_NODE:
			ptr = (char *)val;
			ptr += FDT_TAGSIZE;
			debug("Starts: '%s'\n", ptr);
			while (*ptr != '\0')
				ptr++;
			offset_adjust += ptr - (char *)val + 1;
			debug("skipped %d chars\n", offset_adjust - FDT_TAGSIZE);
			break;
		case FDT_PROP:
			ptr = (char *)val;
			ptr += FDT_TAGSIZE;
			int len = fdt32_to_cpu(*((int *)ptr));
			debug("It's prop: len %d\n", len);
			offset_adjust += sizeof(struct fdt_property) + len;
			break;
		case FDT_END:
		case FDT_END_NODE:
		case FDT_NOP:
			offset_adjust += FDT_TAGSIZE;
			break;
		default:
			printf("Unknown label: %x\n", tag);
			return -1;
	}


	offset += offset_adjust;
	offset_adjust += FDT_TAGALIGN(offset) - offset;
	debug("Next tag offset: %x\n", offset);
	return offset_adjust;
}

int is_offset_valid(char * fdt, int offset, int len)
{
	unsigned int uoffset = offset;
	unsigned int absoffset = offset + fdt_off_dt_struct(fdt);

	if (offset < 0)
		return 0;

	if ((absoffset < uoffset)
	    || ((absoffset + len) < absoffset)
	    || (absoffset + len) > fdt_totalsize(fdt))
		return 0;

	if (fdt_version(fdt) >= 0x11)
		if (((uoffset + len) < uoffset)
		    || ((offset + len) > fdt_size_dt_struct(fdt)))
			return 0;

	return 1;
}

#define WINDOW_SIZE (1024 * 2)

int buf_offset_left_after_tag(int offset)
{
#define GUESS_NODE_NAME_LEN 64
	/* In edge case we have:
	 * |----|----------------|
	 * |tag |node name      0|
	 *                      ^-- last readed byte
	 */
	return WINDOW_SIZE - offset - FDT_TAGSIZE - GUESS_NODE_NAME_LEN;
}

static int fdt_node_get_prop(char * filename, char * fdt, char * strings, int offset, const char * prop_name, int * data_offset, int * data_len)
{
	char * buf = malloc(WINDOW_SIZE);
	int readed = firmware_fdt_read_chunk(filename, buf, fdt_offset_to_file_offset(fdt, offset),
	                                     WINDOW_SIZE);
	if (readed <= 0)
		return -1;

	*data_offset = 0;
	*data_len = 0;
	int buf_offset = 0;
	do {
		int offset_adjust = guess_next_tag(fdt, offset, buf, buf_offset);
		int valid = is_offset_valid(fdt, offset, offset_adjust);
		if (!valid)
		{
			printf("DUMP: next tag is not valid! (%d + %d)\n", offset, offset_adjust);
			break;
		}
		offset += offset_adjust;
		buf_offset += offset_adjust;
		debug("DUMP: Offset %x(%x) is valid\n", offset, buf_offset);

		/* Check if we exceed the window */
		if (buf_offset_left_after_tag(buf_offset) < 0)
		{
			readed = firmware_fdt_read_chunk(filename, buf,
			                                 fdt_offset_to_file_offset(fdt, offset),
			                                 WINDOW_SIZE);
			debug("DUMP: Read additional %d bytes from offset %x\n", readed, offset);
			buf_offset = 0;
		}

		int tag = fdt32_to_cpu(*(int *)(fdt_offset_ptr_(fdt, buf, buf_offset)));
		if (FDT_END_NODE == tag)
			break;
		else if (FDT_PROP == tag)
		{
			const struct fdt_property *prop = fdt_offset_ptr_(fdt, buf, buf_offset);
			debug("prop->nameoff: %x (%x)\n", fdt32_to_cpu(prop->nameoff), fdt_off_dt_strings(fdt));
			char * name = strings + fdt32_to_cpu(prop->nameoff);
			debug("Property name: '%s'\n", name);

			if (strcmp(name, prop_name))
				continue;
			*data_offset = prop->data - (char *)prop + offset;
			*data_len = fdt32_to_cpu(prop->len);
			debug("Property \"%s\" at offset %x len %d\n", prop_name, data_offset, data_len);
			return 0;
		}
	} while (FDT_END != fdt32_to_cpu(*(int *)(fdt_offset_ptr_(fdt, buf, buf_offset))));

	return -1;
}

static int fdt_dump_node_attribute(char * fdt, char * strings, const char *attribute_name,
                                   const char *outname, const char *filename, int offset)
{
	int data_offset = 0;
	int data_len = 0;
	if (fdt_node_get_prop(filename, fdt, strings, offset, attribute_name,
	                      &data_offset, &data_len))
		return -1;

	char * data = malloc(data_len);
	int readed = firmware_fdt_read_chunk(filename, data,
	                                     fdt_off_dt_struct(fdt) + data_offset, data_len);
	if (readed <= 0)
		return -1;

#if 0
	if (fdt_node_get_prop(filename, fdt, strings, offset, "commit", &data_offset, &data_len))
		return -1;


	char * buf = malloc(WINDOW_SIZE);
	readed = firmware_fdt_read_chunk(filename, buf,
	                                 fdt_off_dt_struct(fdt) + data_offset, data_len);
	char commit[42] = {0};
	strncpy(commit, buf, sizeof(commit));
	printf("Commit(len: %d): '%s'\n", data_len, commit);
#endif
#if 1
	int fd = 0;	/* assume stdin */
	int ret = 0;

	char out_filename[128] = "";
	snprintf(out_filename, 128, "./%s", outname);
	*strchr(out_filename, '@') = '_';
	debug("Open file '%s'\n", out_filename);
	fd = open(out_filename, O_WRONLY | O_CREAT | O_TRUNC, 0666);
	if (fd < 0)
	{
		printf("Error: %d\n", errno);
		return -errno;
	}

	size_t writed = 0;
	do {
		ret = write(fd, data + writed, data_len - writed);
		if (ret < 0)
		{
			ret = -errno;
			break;
		}
		writed += ret;
	} while (ret != 0 && writed < data_len);

	/* Clean up, including closing stdin; return errno on error */
	close(fd);

	if (ret < 0)
		return ret;

	return writed;
#else
	return 0;
#endif
}

#define IH_NMLEN		32	/* Image Name Length		*/
#define IH_MAGIC	0x27051956	/* Image Magic Number		*/

/*
 * all data in network byte order (aka natural aka bigendian)
 */

typedef struct image_header {
	uint32_t	ih_magic;	/* Image Header Magic Number	*/
	uint32_t	ih_hcrc;	/* Image Header CRC Checksum	*/
	uint32_t	ih_time;	/* Image Creation Timestamp	*/
	uint32_t	ih_size;	/* Image Data Size		*/
	uint32_t	ih_load;	/* Data	 Load  Address		*/
	uint32_t	ih_ep;		/* Entry Point Address		*/
	uint32_t	ih_dcrc;	/* Image Data CRC Checksum	*/
	uint8_t		ih_os;		/* Operating System		*/
	uint8_t		ih_arch;	/* CPU architecture		*/
	uint8_t		ih_type;	/* Image Type			*/
	uint8_t		ih_comp;	/* Compression Type		*/
	uint8_t		ih_name[IH_NMLEN];	/* Image Name		*/
} image_header_t;


int b(int argc, const char * argv[])
{
	int len = 0;
	if (argc < 2)
	{
		printf("Gimme path to FIT!\n");
		return -1;
	}

	if (argc < 3)
	{
		printf("Gimme name of attribute to dump from FIT!\n");
		return -1;
	}

	if (argc < 4)
	{
		printf("Gimme FIT node name to dump it's attribute!\n");
		return -1;
	}

	int argi = 1;
	const char * filename = argv[argi++];
	const char * attribute_name = argv[argi++];

	int fdt_start = 0;
	char * milf = malloc(WINDOW_SIZE);
	if (!milf)
	{
		printf("<%s> malloc failed", __func__);
		return -1;
	}

	int readed = firmware_fdt_read_chunk(filename, milf, 0, WINDOW_SIZE);
	if (readed < 0)
	{
		free(milf);
		return -1;
	}

	fit_offset = 0;
	if (IH_MAGIC != ntohl(*(uint32_t *)milf))
		goto skip_milf;

	fdt_start += sizeof(image_header_t);
	char * subimage_len_ptr = milf + sizeof(image_header_t);
	int subimage_len = 0;
	int subimage_idx = 0;
	do
	{
		subimage_len = ntohl(*(uint32_t *)subimage_len_ptr);
		subimage_len_ptr += sizeof(uint32_t);
		debug("milf%d: `%x`\n", subimage_idx, subimage_len);
		fdt_start += sizeof(uint32_t);
		if (subimage_idx++ < 3)
			/* Images is aligned for 4 bytes, round up to get offset */
			fdt_start += (subimage_len + 3) & ~0b11;
	} while (subimage_len);

	/* No FIT in firmware */
	if (3 > subimage_idx)
	{
		/* Clean up useless */
		printf("Not enough idxs\n");
		free(milf);
		return -1;
	}

	fit_offset = fdt_start;
	debug("Calculated fit_offset: %x\n", fit_offset);

	char *fdt;
skip_milf:
	fdt = malloc(WINDOW_SIZE);
	readed = firmware_fdt_read_chunk(filename, fdt, 0, WINDOW_SIZE);
	if (readed <= 0)
		return -1;

	char *buf = malloc(WINDOW_SIZE);
	readed = firmware_fdt_read_chunk(filename, buf, fdt_offset_to_file_offset(fdt, 0),
	                                     WINDOW_SIZE);
	if (readed <= 0)
		return -1;

	int offset = 0;
	struct { const char * name; int offset; } node_offsets[10] = { 0 };

	int node_idx = 0;
	while (argi < argc)
	{
		node_offsets[node_idx].name = argv[argi];
		node_offsets[node_idx].offset = -1;
		debug("Remember argv[%d]: '%s'\n", 2 + node_idx, argv[argi]);
		argi++;
		node_idx++;
	}
	debug("Search these nodes:\n");
	for (int i = 0; node_offsets[i].name; i++)
	{
		debug("%s\n", node_offsets[i].name);
	}

	int buf_offset = 0;
	do {
		int offset_adjust = guess_next_tag(fdt, offset, buf, buf_offset);
		int valid = is_offset_valid(fdt, offset, offset_adjust);
		if (!valid)
		{
			printf("READ: next tag is not valid! (%d + %d)\n", offset, offset_adjust);
			break;
		}
		offset += offset_adjust;
		buf_offset += offset_adjust;
		debug("READ: Offset %x(%x) is valid\n", offset, buf_offset);

		/* Check if we exceed the window */
		if (buf_offset_left_after_tag(buf_offset) < 0)
		{
			readed = firmware_fdt_read_chunk(filename, buf,
			                                 fdt_offset_to_file_offset(fdt, offset),
			                                 WINDOW_SIZE);
			debug("READ: Read additional %d bytes from offset %x\n", readed, offset);
			buf_offset = 0;
		}

		if (FDT_BEGIN_NODE == fdt32_to_cpu(*(int *)(fdt_offset_ptr_(fdt, buf, buf_offset))))
		{
			debug("Search starts\n");
			for (int i = 0; node_offsets[i].name; i++)
			{
				if (!strcmp(node_offsets[i].name,
				            (char *)(fdt_offset_ptr_(fdt, buf,
				                                     buf_offset + FDT_TAGSIZE))))
				{
					node_offsets[i].offset = offset;
				}

			}
			debug("Search ends\n");
		}
	} while (FDT_END != fdt32_to_cpu(*(int *)(fdt_offset_ptr_(fdt, buf, buf_offset))));

	debug("Yay\n");

	char * strings = malloc(WINDOW_SIZE);
	readed = firmware_fdt_read_chunk(filename, strings, fdt_off_dt_strings(fdt), WINDOW_SIZE);
	if (readed <= 0)
		return -1;

	for (int i = 0; node_offsets[i].name; i++)
	{
		if (node_offsets[i].offset < 0)
			continue;
		printf("%s offset: 0x%x\n", node_offsets[i].name, node_offsets[i].offset);
		int writed = fdt_dump_node_attribute(fdt, strings, attribute_name, node_offsets[i].name,
		                                     filename, node_offsets[i].offset);
		if (writed)
			printf("Written %d bytes\n", writed);
	}

	return 0;
}


int main(int argc, char *argv[])
{
	b(argc, argv);
	debug("Counter: %d\n", reads_counter);
	return 0;
}