#include <stddef.h>
#include <stdint.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>

// known ID table:
// Chip name   | ARM core   | ARM CPUID  | JTAG IDCODE | Jazelle ID
// ------------+------------+------------+-------------+----------- Jazelle DBX:
// ???         | ARM7EJ-S   | ???        | ???         | ???
// Cypress FX3 | ARM926EJ-S | 0x41069265 | 0x07926069  | 0x64100004
// TI Nspire   | ARM926EJ-S?| 0x41069265?| ???         | 0x64100004?
// RPi v1.??   | ARM11???   | ???        | ???         | ???
// Ninty ?3DS  | ARM11MPCore| 0x410FB025?| ???         | 0x74100064?
// Xilinx Zynq7| Cortex-A9MP| 0x413FC090 | 0x4BA00477  | 0xF4100168
// ------------+------------+------------+-------------+----------- Jazelle RCT:
// Xilinx Zynq7| Cortex-A9MP| 0x413FC090 | 0x4BA00477  | 0xF4100168
// TODO: others (BeagleBoard Cortex-A8? Other A9s?)

// TODO: immediate next steps:
// * check if we can override jazelle insn execution with custom handlers
// * check what the control registers actually do
// * enumerate what all insns do

__attribute__((__naked__))
static uint32_t arm_get_id(void) {
	asm volatile(
		"mrc p15, 0, r0, c0, c0, 0\n"
		"bkpt #4\n"
		"bx lr\n"
	);
}

__attribute__((__naked__))
static uint32_t jazelle_get_id(void) {
	asm volatile(
		"mrc p14, 7, r0, c0, c0, 0\n"
		"bkpt #5\n"
		"bx lr\n"
	);
}

__attribute__((__aligned__(1024)))
static struct {
	void* handlers[512];
	uint8_t stack[256];
	uint8_t locals[256];
} jazelle_block;

/*
 *  c0: Jazelle Identity register (read-only)
        Bits 0-11: Subarchitecture-defined bits (reads as 4, meaning unknown)
        Bits 12-19: Subarchitecture (reads as 0, Jazelle V1 according to documentation)
        Bits 20-27: Implementor (reads as 0x41, ARM Limited according to documentation)
        Bits 28-31: Architecture (reads as 6, ARMv5TEJ according to documentation)
    c1: Operating System Control register
        Bit 0: Configuration Disabled (CD) (documented)
        Bit 1: Configuration Valid (CV) (documented)
    c2: Main Configuration register
        Bit 0: Jazelle Enable (JE) (documented)
        Bits 26-28: Unknown
        Bit 29: If set, array object contains its elements directly, otherwise it contains a pointer to its elements
        Bit 31: Disable array instructions if set?
    c3: Array object layout register
        Bits 0-7: Unknown
        Bits 8-11: Offset (in words) within array object of first element or of pointer to first element
        Bits 12-15: Offset (in words) within array object of length
        Bit 16: If set, offset to length is subtracted, otherwise added
        Bits 17-19: Array length shift value (number of elements = stored length >> this)
        Bits 20-21: Disable array instructions if set?
 */

static uint32_t jazelle_exit_save;

__attribute__((__naked__))
static int jazelle_exec_native(const void* bytecode, const void* block) {
	// inline asm parameters seems to be borking in GCC sooooo lets do it this way
	(void)bytecode; (void)block; (void)&jazelle_exit_save;
	asm volatile(
		"push {r4-r12,lr}\n"

		"mov lr, r0\n"

		// init handler table pointer and stack pointer
		"mov r5, r1\n"
		"add r6, r5, #0x800\n"
		"add r7, r5, #0x900\n"
		// "r8: Pointer to constant pool? (haven't checked this yet)"    -Hackspire

		// set configuration valid & jazelle enable bits
		"mov r0, #2\n"
		"mcr p14, 7, r0, c1, c0, 0\n"
		"mov r0, #1\n"
		"mcr p14, 7, r0, c2, c0, 0\n"

		// apparently there's no good way to find the exit point from jazelle,
		// so we're going to hack that into the stuff now
		"ldr r12, =jazelle_exit_save\n"
		"adr r0, .Ljazend\n"
		"str r0, [r12]\n"

		// switch to jazelle mode
		"adr r12, .Lno_jazelle\n"
		"bxj r12\n"

	".Ljazend:\n"
		"mov r0, #0\n"
		"b .Lend\n"

	".Lno_jazelle:\n"
		"bkpt #1\n"
		"mov r0, #1\n"

	".Lend:\n"
		"mov r5, #0\n"
		"mcr p14, 7, %r5, c1, c0, 0\n"
		"mcr p14, 7, %r5, c2, c0, 0\n"

		"pop {r4-r12,lr}\n"
		"bx lr\n"

		".pool\n"
	);
}



__attribute__((__naked__))
static void handler_idiv(void) {
	// r0 = 3
	// r1 = 4
	// r2 = 2
	// editing the above has no effect
	// NOTE: these depend on the stack content, i.e. it's not a moving register
	//       window. pushes happen in the following order: r0,r1,r2,r3
	//       TODO: when the 'register stack cache' is full, what happens? does
	//             it loop or does it act in a FIFO way, moving r0<-r1<-r2<-r3?
	//             how can the fillrate be known??? (libjz says r5 & 3 but i
	//             dont see anything like that, maybe its an ARM11 thing?)
	// "r4: Copy of local variable 0. Only valid when local variable 0 is a
	//  single word (i.e. int, float, or Object; not long or double)"
	//		-Hackspire
	//
	// [r6-4] is stack top (2)
	// [r6-8] is 4
	// etc
	// use the above to manipulate the stack, eg. "iadd" is implemented as:
	//   - add r1, r2         // or equivalently, read from stack i guess
	//   - str r1, [r6, #-8]  // store to the place where it will be read
	//   - sub r6, #4         // pop off & discard stack top element
	// NOTE: this input usage (with r1 and r2) is NOT robust at all, use memory
	//       reads instead!

	asm volatile(
		"bkpt #3\n"
		// FIXME: read out stack contents in a better way
		"add r1, r2\n"
		"str r1, [r6,#-8]\n"
		"sub r6, #4\n"

		// return to jazelle (yes lr has to be incremented otherwise the
		// current instruction keeps getting executed in a loop)
		"add lr, #1\n"
		"bxj r12\n" // FIXME: r12 can be modified by jazelle so it should be restored to something
	);
}

__attribute__((__naked__))
static void handler_ireturn(void) {
	int result;
	asm volatile(
		"bkpt #2\n"
		"ldr %[res], [r6, #-4]!\n"
		:[res]"=r"(result)
	);
	//iprintf("result=%d\r\n", result); // FIXME: save & restore r0-r3 if ret implemented properly

	// get back to original code
	// TODO: later stage: get back to previous bytecode stuff
	asm volatile(
		"ldr r12, %[exsav]\n"
		"bx r12\n"
		:
		:[exsav]"m"(jazelle_exit_save)
		:"r12"
	);
	__builtin_unreachable();
}

static int jazelle_exec(const uint8_t* bytecode) {
	jazelle_block.handlers[0x6C] = handler_idiv;
	jazelle_block.handlers[0xAC] = handler_ireturn;
/*
 *  +000-3FF: Unhandled bytecodes
        The stack is flushed to memory before calling any of these handlers, so they may modify r0-r3 freely
    +400: Null pointer exception
    +404: Array index out of bounds exception
    +40C: Jazelle mode entered with JE = 0
    +410: Configuration invalid (Jazelle mode entered with CV = 0)
        CV is automatically set to 1 on entering this handler
    +414: Prefetch abort occurred in middle of instruction
		-Hackspire
 */
	const void* block = &jazelle_block;

	return jazelle_exec_native(bytecode, block);
}


// https://github.com/SonoSooS/libjz/wiki/Java-instruction-set
static uint8_t bytecode_test1[] = {
	0x06, // iconst_3
	0x07, // iconst_4
	0x05, // iconst_2
	0x6C, // idiv
	0x04, // iconst_1
	0x60, // iadd
	0x60, // iadd
	0xAC, // ireturn
};

void jazelle_main(void) {
	uint32_t aid = arm_get_id();
	uint32_t jid = jazelle_get_id();
	asm volatile("bkpt #0\n");
	//iprintf("hello world! ARM coreID=0x%lx jazelle ID=0x%lx\r\n", aid, jid);
	while (jid == 0) ;

	int r = jazelle_exec(bytecode_test1);
	//iprintf("retcode=%d\r\n", r);

	while (true) ;
}

void main(void) {
	jazelle_main();
}

__attribute__((__naked__, __section__(".entry")))
void _start(void) {
	asm volatile(
		"ldr sp, =0x00010000\n"
		"bl main\n"
	"1:	 b 1b\n"
		".pool\n"
	);
}