#define pr_fmt(fmt) "[avm partparse][%s] " fmt, __func__

#include <linux/slab.h>
#include <uapi/linux/magic.h>
#include <linux/mtd/mtd.h>
#include <linux/mtd/partitions.h>
#include <linux/stddef.h>
#include <linux/env.h>
#include <linux/jffs2.h>
#include <linux/of.h>

#define IFX_MTD_JFFS2_MIN_SIZE  6
#define IFX_MTD_JFFS2_MAX_SIZE  16

#define MAX_RAM_MTDS 2

static int itsmagic(uint32_t word)
{
	/* avm squashfs may be using cpu byte-order, so check both */
	if(word == SQUASHFS_MAGIC || swab32(word) == SQUASHFS_MAGIC) {
		return SQUASHFS_MAGIC;
	} else {
		return false;
	}
}

/* returns the offset of a superblock or mtd->size on error */
static size_t find_rootfs(struct mtd_info *mtd, uint32_t start_offs)
{
	size_t pos, readlen;
	uint32_t word;
	for(pos = (start_offs & ~0xFF); pos < 0x80000000; pos += 256) {
		if(mtd_read(mtd, pos, sizeof(word), &readlen, (u_char *)&word) ==
				-EINVAL) {
			return mtd->size;
		}
		if(itsmagic(word)) {
			pr_debug("found some superblock at mtd offset %d.\n",
					pos);
			break;
		}
	}

	return pos;
}

/* TARPLUGIN:
 * Used for manufacturing process or recovery.
 * The TARPLUGIN is stored in ramfs mtd behind the squashfs.
 */

#define TAR_ALIGN_SHIFT 9
#define TAR_ALIGN			(1 << TAR_ALIGN_SHIFT)
#define TAR_MAGIC_VALUE			"TAR_PLUGIN_AVM"
#define TAR_MAGIC_VALUE_LEN		(sizeof(TAR_MAGIC_VALUE) - 1)
#define TAR_MAGIC_START_POS		(TAR_ALIGN - TAR_MAGIC_VALUE_LEN)

/* returns the offset of tar-plugin or mtd->size on error */
static size_t find_tarplugin(struct mtd_info *mtd, uint32_t startpos)
{
	u_char buf[TAR_MAGIC_VALUE_LEN];
	size_t pos, readlen;

	for (pos = startpos + TAR_MAGIC_START_POS; pos < mtd->size;
			pos += TAR_ALIGN) {
		if (mtd_read(mtd, pos, TAR_MAGIC_VALUE_LEN, &readlen, buf) ==
				-EINVAL) {
			return mtd->size;
		}
		if ((readlen == TAR_MAGIC_VALUE_LEN) &&
		    (memcmp(buf, TAR_MAGIC_VALUE, TAR_MAGIC_VALUE_LEN) == 0)) {
			pos += TAR_MAGIC_VALUE_LEN;
			pr_debug("found tar-plugin-avm at mtd offset 0x%x.\n",
			       pos);
			break;
		}
	}
	return pos;
}

/* Devicetree MTD parser */
extern int parse_ofpart_partitions(struct mtd_info *master,
				   struct mtd_partition **pparts,
				   struct mtd_part_parser_data *data);

#if IS_ENABLED(CONFIG_MTD_AVM_MTDRAM)
#  define AVM_MTDRAM1_NAME "update-image.0"
#else
#  define AVM_MTDRAM1_NAME "ram-filesystem"
#endif

static int mtd_parser_of_nor_spi(struct mtd_info *mtd,
                         struct mtd_partition **p_mtd_pat,
                         struct mtd_part_parser_data *parse_data)
{
	struct mtd_partition *parts = 0;
	struct mtd_partition *tmp_ofparts = 0;
	int parts_n = 0;

	pr_debug("entered, name: %s, idx: %u; size: %u (%uMB)\n", mtd->name,
			mtd->index,
			(unsigned int)mtd->size,
			(unsigned int)(mtd->size >> 20));

	if (strcmp(mtd->name, AVM_MTDRAM1_NAME) == 0) {
		unsigned long int offset;
		int tar_plugin_enabled;

		tar_plugin_enabled = of_property_read_bool(of_root,
				"avm,tarplugin");
		parts = kzalloc(sizeof(*parts) *
				(tar_plugin_enabled ? MAX_RAM_MTDS + 1 : MAX_RAM_MTDS),
				GFP_KERNEL);
		BUG_ON(!parts);
		parts[0].name = "rootfs_ram";
		parts[0].mask_flags = 0;
		parts[1].name = "kernel_ram";
		parts[1].mask_flags = 0;
		/* There is no kernel header in ram-filesystem
		 * --> search RAM for squashfs magic */
		offset = find_rootfs(mtd, 0);
		if(offset >= mtd->size) {
			/* something is wrong, as we couldn't find a
			 * filesystem */
			parts_n = 0;
			pr_debug("no fs found\n");
		} else if(offset == 0) {
			parts_n = 1;
			pr_debug("fs found at offset 0x%08x\n",
					(unsigned int)parts[1].offset);
		} else {
			/* we found a fs behind something that we assume to
			 * be the kernel */
			parts[0].offset = offset;
			parts[0].size = mtd->size - offset;
			parts[1].size = offset;
			parts_n = 2;
			pr_debug("fs found at offset 0x%08x\n",
					(unsigned int)parts[1].offset);
		}
		if (offset >= mtd->size)
			offset = 0;
		/* tar plugin may be attached to the end here. */
		offset = tar_plugin_enabled ?
			find_tarplugin(mtd, offset) : mtd->size + 1;
		if (offset < mtd->size) {
			if (parts_n > 0)
				parts[0].size = offset - parts[0].offset;

			parts[parts_n++] = (struct mtd_partition) {
				.name   = "tarplugin_ram",
				.offset = offset,
				.size   = mtd->size - offset,
			};
		} else if (tar_plugin_enabled) {
			pr_info("tarplugin not found.\n");
		}

	} else if (strcmp(mtd->name, "spi0.0") == 0) {
		/* Layout for this mtd is expexted to hold some partition
		 * definitions in device tree. parse_ofpart_partitions is used
		 * to handle this part. The remaining space in the mtd is
		 * expected to contain a construct like this:
		 * | Kernel Header | Kernel | root fs (squash fs) | jffs2 (user fs) |
		 *
		 * Kernel header is:
		 *   int tag: TAG_WANTED | TAG_PRE_IPTYPE | TAG_COMPRESSED
		 *            (we ignore this for now)
		 *   int len: Length of this header block; Kernel start address is
		 *       header_addr + ((len + 3) & ~3) + 6 * sizeof(int)
		 *
		 * root fs starts like this:
		 *   int SQASHFS_MAGIC
		 *   int bytes_used
		 *
		 * jffs2 starts with a special token
		 *
		 * If JUST contents from DT should be considered for mtd
		 * partitioning, use a different name than "spi0.0"!
		 */
		/* TODO: Die Namensänderung für SPI müsste auch im Urlader nachgezogen werden! */
		int n, next_n, i, next_i;
		unsigned int offset, next_offs, jffs2_offs_1st = 0, jffs2_offs_last = 0, jffs2_align;
		struct mtd_partition tmp_part;
		bool found_jffs = false, found_rootfs = false;

		parts_n = parse_ofpart_partitions(mtd, &tmp_ofparts, parse_data);
		BUG_ON(parts_n <= 0);
		pr_debug("%d Partitions found in device tree.\n", parts_n);
		if (mtd->name[6] == 'n') {
			*p_mtd_pat = tmp_ofparts;
			return parts_n;
		}

		parts = kzalloc(sizeof(*parts) * (3 + parts_n), GFP_KERNEL);
		BUG_ON(!parts);

		parts[parts_n + 0].mask_flags = 0;
		parts[parts_n + 1].mask_flags = 0;

		memcpy(parts, tmp_ofparts, sizeof(*tmp_ofparts) * parts_n);
		kfree(tmp_ofparts);
		/* Use bubble sort on list of partitions. */
		for (n = parts_n; n; n = next_n) {
			next_n = 0;
			for (i = 1; i < n; ++i) {
				if (parts[i - 1].offset > parts[i].offset) {
					tmp_part = parts[i - 1];
					parts[i - 1] = parts[i];
					parts[i] = tmp_part;
					next_n = i;
				}
			}
		}

		/* TODO: tffs_mtd array was updated here; This is "not so good" -- better do this in tffs itself when
		 * we are done with all mtds and partitions.
		 * Needs an update of TFFS! */
		// TKL: NOR-TFFS Partitionen muessen "tffs (1/2)" heissen, damit Scripte sie in /proc/mtd
		//      finden koennen.
		for (i = 0; i < parts_n; ++i) {
			if (strcmp(parts[i].name, "tffs1") == 0
					|| (strcmp(parts[i].name, "tffs (1)") == 0))
			{
				parts[i].name = "tffs (1)";
			} else if (strcmp(parts[i].name, "tffs2") == 0
					|| (strcmp(parts[i].name, "tffs (2)") == 0))
			{
				parts[i].name = "tffs (2)";
			}
		}

		/* At least 64k align for JFFS2 */
		jffs2_align = mtd->erasesize;
		if (jffs2_align < 0x10000) jffs2_align = 0x10000;
		/* Search whole mtd for Kernel/Rootfs/JFFS2 */
		for (i = 0, offset = 0; offset < mtd->size; i = next_i, offset = next_offs) {
			unsigned int tmp_word, mtd_start;

			pr_debug("i=%d; offs=0x%08x\n", i, offset);

			if (i < parts_n) {
				next_i = i + 1;
				next_offs = min(offset + parts[i].size, mtd->size);
				/* Prüfen, ob der Bereich bereits im DT deklariert wurde */
				if (parts[i].offset == offset) {
					pr_debug("mtd part offs match\n");
					continue;
				}
			} else {
				next_i = 0;
				next_offs = mtd->size;
			}

			/* ...oder zu Klein für was Sinnvolles ist */
			if (next_offs - offset < 0x40000) {
				pr_debug("small gap\n");
				continue;
			}

			/* Hier wird dann wohl der Kernel-Header drin stecken, gefolgt von Squashfs und JFFS2 */
			if (mtd_read(mtd, offset + sizeof(int),
						sizeof(tmp_word), &n,
						(u_char *)&tmp_word) ==
					-EINVAL ||
					n != sizeof(tmp_word)) {
				pr_err("mtd read error @ offs=0x%08x for kernel header\n",
						offset + sizeof(int));
				break;
			}
			tmp_word = ((tmp_word + 3) & ~3) + (2 * 3 * sizeof(unsigned int));  /* 4byte align + header */
			if (tmp_word + offset >= next_offs || tmp_word < 0x10000 || tmp_word > mtd->size) {
				pr_info("Invalid kernel header 0x%08x @ offs 0x%08x\n",
						tmp_word, offset + sizeof(int));
				offset = (offset + 0x10000) & ~(jffs2_align - 1);
				if (!found_jffs && offset < next_offs &&
						next_offs - offset >= IFX_MTD_JFFS2_MIN_SIZE * jffs2_align &&
						next_offs - offset > jffs2_offs_last - jffs2_offs_1st) {
					pr_info("Saving 0x%08x -- 0x%08x as empty for jffs2 if no existing jffs2 is found otherwise.\n",
							offset, next_offs);
					jffs2_offs_1st = offset;
					jffs2_offs_last = next_offs;
				}
				continue;
			}

			mtd_start = offset;
			offset = (offset + tmp_word + 256) & ~0xFF;  /* 256 byte align */

			/* squashfs magic suchen ... zumindest bis 1 JFFS block vor Ende */
			tmp_word = 0;
			while (offset < ((next_offs - 0x10000) & ~0xFFFF)) {
				if (mtd_read(mtd, offset, sizeof(tmp_word), &n, (u_char *)&tmp_word) == -EINVAL || n != sizeof(tmp_word)) break;
				if (itsmagic(tmp_word)) {
					/* So weit so gut: Squashfs-Größe? */
					offset += sizeof(unsigned int) * 2;
					if (mtd_read(mtd, offset, sizeof(tmp_word), &n, (u_char *)&tmp_word) == -EINVAL || n != sizeof(tmp_word)) {
						pr_err("mtd read error @ offs=0x%08x for squashfs size\n",
								offset);
						tmp_word = 0;
						break;
					}
					if (tmp_word < 0x10000) {
						pr_warn("weird small squashfs size 0x%08x - smells fishy!\n",
								tmp_word);
					} else {
						found_rootfs = true;
						/*pr_denbug*/pr_info("fond_rootfs\n");
						break;
					}
				}
				offset = (offset + 256) & ~0xFF;
			}

			if (!found_rootfs) {
				pr_info("Did not find Squashfs after kernel\n");
				continue;
			}

			/* sieht gut aus: Kernel / rootfs eintragen */
			parts[parts_n].offset = mtd_start;
			parts[parts_n].size = next_offs - mtd_start;  /* Erst mal alles zuschlagen */
			parts[parts_n].name = "rootfs_kernel_spi";
			pr_debug("rootfs found @offs 0x%08x --> mtd part offset: 0x%08x; size: 0x%08x; end_offs: 0x%08x; next_offs: 0x%08x\n",
					offset & ~ 0xFF,
					(unsigned int)parts[parts_n].offset, (unsigned int)parts[parts_n].size,
					(unsigned int)parts[parts_n].offset + (unsigned int)parts[parts_n].size, next_offs);
			pr_debug("rootfs bytes used: 0x%08x\n", tmp_word);
			parts_n++;

			offset = (offset + tmp_word + 0x10000) & ~(jffs2_align - 1);
#ifdef CONFIG_JFFS2_FS
			jffs2_offs_1st = offset;
			jffs2_offs_last = next_offs & ~(jffs2_align - 1);
			pr_debug("Start search for JFFS2 @0x%08x\n", offset);
			while (offset < next_offs) {
				if (mtd_read (mtd, offset, sizeof(tmp_word), &n, (u_char *)&tmp_word) == -EINVAL || n != sizeof(tmp_word)) {
					pr_err("mtd read error @ offs=0x%08x for JFFS2 token\n", offset);
					break;
				}

#ifdef __LITTLE_ENDIAN
				if ((tmp_word & 0xFFFF) == JFFS2_MAGIC_BITMASK) 
#else
                if ((tmp_word >> 16) == JFFS2_MAGIC_BITMASK) 
#endif
                {
						pr_debug("JFFS2 token 0x%x found @0x%08x\n",
								tmp_word,
								offset);
						/* JFFS2 gefunden: eintragen */
						parts[parts_n].offset = offset;
						parts[parts_n].size = next_offs - offset;
						parts[parts_n].name = "jffs2";
						/* JFFS2 von kernel/rootfs mtd size wieder abziehen */
						/*--- parts[parts_n - 1].size = offset - mtd_start; ---*/
						parts_n++;
						// Wir sind hier fertig
						offset = mtd->size;
						found_jffs = true;
						break;
                }

                offset += 1 << 16;  /* next 64k block */
            }
            break;

        }

			if (!found_jffs) {
				/* JFFS2 nicht gefunden */
				unsigned int jffs2_size_urlader, jffs2_size, jffs2_offs;
				bool jffs2_size_changed = false;
				char *p;

				/* size in 64k-blöcken und aligned */
				jffs2_size = (jffs2_offs_last - jffs2_offs_1st) >> 16;
				jffs2_offs = jffs2_offs_1st;
				pr_debug("JFFS2 not found. jffs2_size: %d (64k blocks); jffs2_offs_1st: 0x%08x; jffs2_offs_last: 0x%08x\n",
						jffs2_size, jffs2_offs_1st, jffs2_offs_last);
				p = prom_getenv("jffs2_size");
				if (p) {
					jffs2_size_urlader = simple_strtoul(p, NULL, 10);
				} else {
					jffs2_size_changed = true;
					jffs2_size_urlader = 0;
					pr_err("jffs2_size not set\n");
				}

				if (jffs2_size < (IFX_MTD_JFFS2_MIN_SIZE * jffs2_align >> 16)) {
					pr_warn("not enough space for JFFS2!\n");
				} else {
					struct erase_info instr;
					/* Genug Platz wäre vorhanden... */
					if (jffs2_size_urlader == 0) {
						/* Keine size im Urlader:
						 * auf 1 MB begrenzen und nach ganz hinten
						 * oder wenn weniger frei ist, dann das nutzen
						 */
						if (jffs2_size > 1 << (20 - 16)) {
							jffs2_offs = jffs2_offs_last - (1 << 20);
							jffs2_size = 1 << (20 - 16);
							jffs2_size_changed = true;
							pr_info("using last 1MB for JFFS2: offs: 0x%08x; size= 0x%08x\n", jffs2_offs, jffs2_size);
						}
					} else {
						/* Size aus Urlader verwenden */
						if (jffs2_size > jffs2_size_urlader) {
							/* Größe beschränken: von hinten abziehen */
							if (jffs2_size_urlader < (IFX_MTD_JFFS2_MIN_SIZE * jffs2_align >> 16)) {
								jffs2_size_urlader = (IFX_MTD_JFFS2_MIN_SIZE * jffs2_align >> 16);
								jffs2_size_changed = true;
								printk(KERN_WARNING "[%s]: jffs2_size set too small, use %d instead\n", __func__, jffs2_size_urlader);
							}
							jffs2_offs = jffs2_offs_last - (jffs2_size_urlader << 16);
							jffs2_size = jffs2_size_urlader;
						} else if (jffs2_size < jffs2_size_urlader) {
							pr_warn("jffs2_size is set greater than space left on flash for JFFS2 MTD: %u vs %u (64k blocks)\n",
									jffs2_size_urlader, jffs2_size);
							jffs2_size_changed = true;
						}
						/* Wenn zu wenig oder gerade richtig Platz da ist, hier fertig. */
					}
					/* mtdpart eintragen und löschen --> wird beim mounten neu erzeugt. */
					parts[parts_n].offset = jffs2_offs;
					parts[parts_n].size = jffs2_size << 16;
					parts[parts_n].name = "jffs2";
					/*--- parts[parts_n - 1].size = parts[parts_n].offset - parts[parts_n - 1].offset;  // kernel/rootfs mtdpart verkleinern ---*/
					parts_n++;

					if (jffs2_size_changed) {
						pr_info("Erase JFFS2 mtd @0x%08x, size 0x%08x (%u 64k blocks)\n",
								(unsigned int)parts[parts_n - 1].offset,
								(unsigned int)(jffs2_size << 16), jffs2_size);

						memset(&instr, 0, sizeof(instr));
						instr.mtd = mtd;
						instr.addr = jffs2_offs;
						instr.len = jffs2_size << 16;
						instr.callback = NULL;
						instr.fail_addr = 0xffffffff;

						i = mtd_erase(mtd, &instr);
						if (i) {
							printk(KERN_ERR "jffs2 mtd erase failed %d\n", i);
						}
						pr_info("Erase done.\n");
					}
				}
			}
#endif

			pr_debug("done.\n");

		} else {
			/* nand not handled yet */
			/* TODO: Decide on a way to store partitioning for NAND! (And a layout) */
		}

		if (parts_n > 0)
			*p_mtd_pat = parts;
		else
			kfree(parts);

		return parts_n;
}

struct mtd_part_parser avm_mtd_parser = {
	.owner = THIS_MODULE,
	.parse_fn = mtd_parser_of_nor_spi,
	.name = "avmpart_of_nor_spi",
};

static int __init avm_mtd_parser_init(void)
{
	register_mtd_parser(&avm_mtd_parser);

	return 0;
}
subsys_initcall(avm_mtd_parser_init);