#define TINYRISCV_M

#include "elf.h"
#include "hwinfo.h"
#include "util.h"

#include "aes.h"
#include "tinyriscv.h"
#include "config.h"

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

#define MESSAGE_UNAUTHORIZED "You are not authorized to use this software!\n"

#define MAX_STEPS 1024u


static const char* emu_program_paths[] = EMU_PROGRAM_PATHS;

int main(){
    hw_info_t hwinfo = {0};
    if(get_hw_info(&hwinfo) != 0){
        return 1;
    }

    uint8_t* key = hwinfo_to_key(hwinfo);

    if(key == NULL) {
        return 1;
    }

    // fprintf(stderr, "HWINFO: %s\n", hw_info_flatten(hwinfo));

    fprintf(stderr, "[Info] Unique hardware key: ");
    for (int i = 0; i < 16; i++)
        fprintf(stderr, "%02X", key[i]);
    fprintf(stderr, "\n");



    struct AES_ctx ctx;
    uint8_t nonce[16] = {0};

    FILE* license;
    for(size_t i = 0; i < sizeof(emu_program_paths); i++) {
        license = fopen(emu_program_paths[i], "rb");
        if (license != NULL) {
            break;
        }
    }

    if (license == NULL) {
        fprintf(stderr, "[Error] Could not find license file\n");
        return 1;
    }

    fread(nonce, 16, 1, license);
    long data_start = ftell(license);
    fseek(license, 0, SEEK_END);
    int license_size = ftell(license) - data_start;
    fseek(license, data_start, SEEK_SET);
    uint8_t* buf = malloc(license_size);
    fread(buf, 1, license_size, license);


    AES_init_ctx(&ctx, key);
    AES_init_ctx_iv(&ctx, key, nonce);
    AES_CTR_xcrypt_buffer(&ctx, buf, license_size);

    FILE* decrypted_license = tmpfile();
    fwrite(buf, 1, license_size, decrypted_license);
    rewind(decrypted_license);

    char begin[8] = {0};
    char expected[] = {0x7F, 'E', 'L', 'F', 1, 1, 1, 0};
    fread(begin, 1, sizeof(begin), decrypted_license);
    if(memcmp(begin, expected, sizeof(begin)) != 0) {
        // This happens when decryption fails
        printf(MESSAGE_UNAUTHORIZED);
        fflush(stdout);
        return 1;
    }

    const int maxninput = 255;
    printf("Please enter your authorization code: ");
    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(decrypted_license, memory, cpu.mem_size - SIZE_USERDATA);
    if (load_status != LOAD_OK) {
        fprintf(stderr, "[Error] ELF load failed with code %d for decrypted license\n", load_status);
        fflush(stderr);
        printf("[Error] 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(
        decrypted_license, EMU_ENTRY_SYMBOL, cpu.mem_size, &entry_offset);
    if (symbol_status != LOAD_OK) {
        fprintf(stderr, "[Error] Entry symbol '%s' resolve failed with code %d\n",
            EMU_ENTRY_SYMBOL, symbol_status);
        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, "[Error] Execution did not halt within %u steps\n", MAX_STEPS);
        return 2;
    }

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

    if(result == 1) {
        printf("License activated successfully!\n");
    } else {
        printf(MESSAGE_UNAUTHORIZED);
    }

    return 0;
}