#include "netboot.h"
#include "bba.h"
#include <stdio.h>
#include <string.h>
#include <malloc.h>
#include <ogc/lwp_watchdog.h>
#include <ogc/cache.h>

/* ── frame layout ────────────────────────────────────────────────────────
 *  0  dst MAC (6)
 *  6  src MAC (6)
 * 12  ethertype (2) = NB_ETHERTYPE
 * 14  magic (4) = "GCUP"
 * 18  opcode (1)
 * 19  reserved (1)
 * 20  seq (2, big endian)
 * 22  nblocks (2, big endian)
 * 24  len (2, big endian)
 * 26  data (len)
 */
#define NB_HDR      26u
#define NB_MIN      60u          /* ethernet minimum, padded */

static u8 s_our_mac[6];
static u8 s_peer_mac[6];

static u16 rd16(const u8 *p) { return ((u16)p[0] << 8) | p[1]; }
static void wr16(u8 *p, u16 v) { p[0] = (u8)(v >> 8); p[1] = (u8)v; }

/* ── RX: pull one frame out of the SPRAM ring ────────────────────────────
 * The ring is pages 1-15 with a 4-byte descriptor at each page start whose
 * length field INCLUDES the descriptor.  RRP is ours to advance; the
 * gateware does not act on it. */
static int ring_next(u8 *buf, u32 cap)
{
	u8 rp[2], wp[2], hdr[4];
	u32 rpage, wpage, dlen, flen, addr;

	if (bba_read(REG_RRP_LO, rp, 2) < 0) return -1;
	if (bba_read(REG_RWP_LO, wp, 2) < 0) return -1;
	rpage = ((u32)rp[1] << 8) | rp[0];
	wpage = ((u32)wp[1] << 8) | wp[0];
	if (rpage == wpage) return 0;                 /* empty */
	if (rpage < 1 || rpage > 15) return -2;       /* pointer corrupt */

	addr = rpage << 8;
	if (bba_stream((u16)addr, hdr, 4) < 0) return -1;
	dlen = ((u32)hdr[2] << 8) | hdr[3];
	if (dlen <= 4 || dlen > 1518 + 4) return -3;  /* descriptor implausible */
	flen = dlen - 4;
	if (flen > cap) flen = cap;

	/* Descriptor + frame are contiguous from the page start.  Read both and
	 * drop the descriptor rather than streaming from a mid-page address, so
	 * one DMA read covers it. */
	{
		static u8 tmp[1536 + 4];
		u32 want = dlen;
		if (want > sizeof(tmp)) want = sizeof(tmp);
		if (bba_stream((u16)addr, tmp, want) < 0) return -1;
		memcpy(buf, tmp + 4, flen);
	}

	/* Advance RRP by the pages this frame occupied, wrapping 1..15. */
	{
		u32 pages = (dlen + 255) >> 8;
		u32 next = ((rpage - 1 + pages) % 15) + 1;
		if (bba_w8(REG_RRP_LO, (u8)next) < 0) return -1;
	}
	return (int)flen;
}

static int send_ack(u16 seq)
{
	u8 f[NB_MIN];
	memset(f, 0, sizeof(f));
	memcpy(f + 0, s_peer_mac, 6);
	memcpy(f + 6, s_our_mac, 6);
	wr16(f + 12, NB_ETHERTYPE);
	f[14] = 'G'; f[15] = 'C'; f[16] = 'U'; f[17] = 'P';
	f[18] = NB_OP_ACK;
	wr16(f + 20, seq);
	wr16(f + 22, 0);
	wr16(f + 24, 0);
	return bba_tx_bytes(f, sizeof(f));
}

int netboot_receive(void *dst, u32 maxlen, void (*status)(const char *msg))
{
	static u8 frame[1536];
	u8 *out = (u8 *)dst;
	u32 got = 0, nblocks = 0;
	u16 expect = 0;
	char line[80];
	u64 t0;
	int have_peer = 0;

	if (bba_read(REG_PAR0, s_our_mac, 6) < 0) return -1;

	/* Make sure RX is actually enabled -- a prior test may have left the
	 * NCRA SR bit clear, and then this would wait forever for frames the
	 * gateware is not storing. */
	{
		int ncra = bba_r8(REG_NCRA);
		if (ncra >= 0 && !(ncra & NCRA_SR))
			bba_w8(REG_NCRA, (u8)(ncra | NCRA_SR));
	}

	if (status) status("waiting for host (send-dol.py)...");
	t0 = gettime();

	for (;;) {
		int n = ring_next(frame, sizeof(frame));
		if (n < 0) {
			if (status) { snprintf(line, sizeof(line), "ring error %d", n);
			              status(line); }
			return -2;
		}
		if (n == 0) {
			/* 30 s with nothing at all, or 5 s mid-transfer, is a bail-out
			 * rather than a hang -- a wedged update must not require a
			 * console power cycle. */
			u32 idle = (u32)ticks_to_millisecs(diff_ticks(t0, gettime()));
			if (idle > (got ? 5000u : 30000u)) {
				if (status) status(got ? "timeout mid-transfer" : "timeout: no frames");
				return -3;
			}
			continue;
		}
		if ((u32)n < NB_HDR) continue;
		if (rd16(frame + 12) != NB_ETHERTYPE) continue;
		if (frame[14] != 'G' || frame[15] != 'C' ||
		    frame[16] != 'U' || frame[17] != 'P') continue;

		if (!have_peer) { memcpy(s_peer_mac, frame + 6, 6); have_peer = 1; }
		t0 = gettime();

		if (frame[18] == NB_OP_DONE) {
			if (status) { snprintf(line, sizeof(line),
			                       "received %u bytes in %u blocks",
			                       (unsigned)got, (unsigned)nblocks);
			              status(line); }
			return (int)got;
		}
		if (frame[18] != NB_OP_DATA) continue;

		{
			u16 seq = rd16(frame + 20);
			u16 len = rd16(frame + 24);
			nblocks = rd16(frame + 22);
			if (len > NB_BLOCK_MAX || NB_HDR + len > (u32)n) continue;

			if (seq == expect) {
				if (got + len > maxlen) {
					if (status) status("image larger than buffer");
					return -4;
				}
				memcpy(out + got, frame + NB_HDR, len);
				got += len;
				expect++;
				if ((expect & 0x1F) == 0 && status) {
					snprintf(line, sizeof(line), "block %u/%u (%u KB)",
					         (unsigned)expect, (unsigned)nblocks,
					         (unsigned)(got >> 10));
					status(line);
				}
			}
			/* ACK even a duplicate: it means our previous ACK was lost, and
			 * re-ACKing is what lets the host move on. */
			send_ack(seq);
		}
	}
}

/* ── DOL loader ──────────────────────────────────────────────────────────
 * 7 text + 11 data sections, each with a file offset, load address and size,
 * then BSS and the entry point.  All big endian. */
typedef struct {
	u32 text_off[7];
	u32 data_off[11];
	u32 text_addr[7];
	u32 data_addr[11];
	u32 text_len[7];
	u32 data_len[11];
	u32 bss_addr;
	u32 bss_len;
	u32 entry;
	u32 pad[7];
} dolhdr_t;

int netboot_launch(const void *dol)
{
	const dolhdr_t *h = (const dolhdr_t *)dol;
	const u8 *base = (const u8 *)dol;
	void (*entry)(void);
	int i;

	/* Sanity: the entry point and at least one text section must live in
	 * MEM1.  Jumping to a mis-parsed header hangs the console hard enough
	 * to need a power cycle, so refuse rather than guess. */
	if ((h->entry >> 24) != 0x80 || h->text_len[0] == 0 ||
	    (h->text_addr[0] >> 24) != 0x80)
		return -1;

	/* No interrupts from here on: we are about to overwrite the memory this
	 * program is running from, and with interrupts off the scheduler cannot
	 * preempt us into code that no longer exists. */
	IRQ_Disable();

	for (i = 0; i < 7; i++) {
		if (!h->text_len[i]) continue;
		memmove((void *)h->text_addr[i], base + h->text_off[i], h->text_len[i]);
		DCFlushRange((void *)h->text_addr[i], h->text_len[i]);
		ICInvalidateRange((void *)h->text_addr[i], h->text_len[i]);
	}
	for (i = 0; i < 11; i++) {
		if (!h->data_len[i]) continue;
		memmove((void *)h->data_addr[i], base + h->data_off[i], h->data_len[i]);
		DCFlushRange((void *)h->data_addr[i], h->data_len[i]);
	}
	if (h->bss_len) {
		memset((void *)h->bss_addr, 0, h->bss_len);
		DCFlushRange((void *)h->bss_addr, h->bss_len);
	}

	entry = (void (*)(void))h->entry;
	entry();
	return 0;               /* not reached */
}