#define TINYRISCV_M

#include "tinyriscv.h"
#include "elf.h"
#include "config.h"

#include <errno.h>
#include <stdio.h>
#include <string.h>


#define MAX_STEPS 1024u


int main(){
    const int maxninput = 255;
    printf("Please enter the flag: ");
    fflush(stdout);
    char userinput[maxninput];
    fgets(userinput, sizeof(userinput), stdin);
    // Replace newline with null terminator
    userinput[strcspn(userinput, "\r\n")] = 0;

    uint8_t memory[EMU_MEM_SIZE]; //4kb of internal memory
    tinyriscv_core cpu;   //create the emulation core
    memset(&cpu, 0, sizeof(cpu));
    cpu.mem = memory;  //use the uint8_t array as the memory
    cpu.mem_size = EMU_MEM_SIZE;
    memset(memory, 0, sizeof(memory));
    uint8_t real_len = strnlen(userinput, maxninput);

    // Load user data, first, the string length
    memory[SECTION_USERDATA] = real_len;
    memcpy(memory + SECTION_USERDATA + 1, userinput, real_len);

    // Load elf, keep 512 bytes as safe buffer
    int load_status = load_elf_object_sections(EMU_PROGRAM_PATH, memory, cpu.mem_size - SIZE_USERDATA);
    if (load_status != LOAD_OK) {
        fprintf(stderr, "ELF load failed with code %d for '%s'\n", load_status, EMU_PROGRAM_PATH);
        fflush(stderr);
        printf("ELF load failed with code %d\n", load_status);
        fflush(stdout);
        return 1;
    }

    uint32_t entry_offset = 0u;
    int symbol_status = resolve_elf_symbol_memory_offset(
        EMU_PROGRAM_PATH, EMU_ENTRY_SYMBOL, cpu.mem_size, &entry_offset);
    if (symbol_status != LOAD_OK) {
        fprintf(stderr, "Entry symbol '%s' resolve failed with code %d for '%s'\n",
            EMU_ENTRY_SYMBOL, symbol_status, EMU_PROGRAM_PATH);
        return 2;
    }

    tinyriscv_init(&cpu); //initialise core for execution
    cpu.pc = tinyriscv_MEM_OFFSET + entry_offset;

    cpu.x[10] = tinyriscv_MEM_OFFSET + SECTION_USERDATA;
    cpu.x[1] = tinyriscv_MEM_OFFSET + cpu.mem_size;

    uint32_t steps = 0u;
    while (tinyriscv_valid_step(&cpu) && steps < MAX_STEPS) {
        tinyriscv_step(&cpu);
        ++steps;
    }

    if (steps == MAX_STEPS) {
        fprintf(stderr, "Execution did not halt within %u steps\n", MAX_STEPS);
        return 2;
    }

    uint32_t result = (uint32_t)cpu.x[10];

    if(result == 1) {
        printf("Congrats! That's the flag!");
    } else {
        printf("Nope, sorry, that's incorrect!");
    }

    return 0;
}