/*
 * QEMU System Emulator
 *
 * Copyright (c) 2003-2020 Fabrice Bellard
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

#include "qemu/osdep.h"
#include "qemu-main.h"
#include "sysemu/runstate.h"
#include "sysemu/sysemu.h"
#include "migration/snapshot.h"
#include <math.h>

#ifdef CONFIG_SDL
#include <SDL.h>
#endif

int snapshot_save(const char *name);
int snapshot_load(const char *name);

int snapshot_save(const char *name)
{
    Error *err = NULL;
    save_snapshot(name, true, NULL, false, NULL, &err);
    return err == 0;
}

int snapshot_load(const char *name)
{
    Error *err = NULL;
    load_snapshot(name, NULL, false, NULL, &err);
    return err == 0;
}

int qemu_default_main(void)
{
    int status;

    status = qemu_main_loop();
    qemu_cleanup(status);

    return status;
}

int (*qemu_main)(void) = qemu_default_main;

#ifndef AS_LIB
//========= Instrumentation start
#include <stdio.h>
#include <stdlib.h>
#include "exec/cpu-common.h"
void libafl_qemu_set_native_breakpoint(vaddr);
void libafl_qemu_remove_native_breakpoint(vaddr);
int libafl_qemu_write_reg(CPUState *cpu, int reg, uint8_t *val);
int libafl_qemu_read_reg(CPUState *cpu, int reg, uint8_t *val);
CPUState *libafl_qemu_current_cpu(void);
int libafl_qemu_num_regs(CPUState *cpu);
int libafl_qemu_num_cpus(void);
CPUState *libafl_qemu_get_cpu(int cpu_index);
int64_t icount_get_raw(void);
//========= Instrumentation end
int main(int argc, char **argv)
{

    int input_size;
    printf("argc: %d\n", argc);
    //========= Instrumentation start
    // read addr and input to load
    if (argc < 3)
    {
        fprintf(stderr, "Need address and input file argument\n");
        exit(1);
    }
    hwaddr prep = (hwaddr)strtoll(argv[1], NULL, 16);
    hwaddr start = (hwaddr)strtoll(argv[2], NULL, 16);
    hwaddr end = (hwaddr)strtoll(argv[3], NULL, 16);
    input_size = atoi(argv[4]);
    char *output_path = argv[5];
    unsigned int num_tasks = atoi(argv[6]);
    // hwaddr target_addr = (hwaddr) strtoll(argv[1], NULL, 16);
    // vm_start();
    // fix arguments for qemu
    argv[6] = argv[0];
    argv = &argv[6];

    argc -= 6;

    int full_input_room = (int)pow(input_size, num_tasks);
    printf("Full input room: %d\n", full_input_room);
    unsigned long *deltas = malloc(full_input_room * sizeof(unsigned long));
    if (deltas == NULL)
    {
        fprintf(stderr, "Memory allocation failed\n");
        exit(1);
    }
    u_int32_t *inputs = malloc(full_input_room * sizeof(u_int32_t));
    if (inputs == NULL)
    {
        fprintf(stderr, "Memory allocation failed\n");
        exit(1);
    }
    //========= Instrumentation end
    qemu_init(argc, argv);
    //========= Instrumentation start

    libafl_qemu_set_native_breakpoint(prep);
    // set int in in the vm to i
    vm_start();
    qemu_main_loop();

    // Now execution is halted at the start of the task we want to measure in order to write the input to a register
    libafl_qemu_remove_native_breakpoint(prep);

    snapshot_save("base");

    uint8_t register_in_32b[4];
    uint8_t reg_tmp_val[num_tasks][4];
    // load input
    // cpu_physical_memory_rw(target_addr, buffer, read_len, true);

    u_int32_t task_inputs[num_tasks];

    fclose(fopen(output_path, "w"));

    FILE *fptr = fopen(output_path, "a");

    // input of all tasks combined
    for (long i = 0; i < full_input_room; i++)
    {
        printf("Input: %lu\n", i);
        for (int j = 0; j < num_tasks; j++)
        {
            // from the "global input" i, extract the input bits for the task j
            task_inputs[j] = (i >> (j * (__builtin_popcount(input_size - 1)))) & (input_size - 1);
            printf("Task %d: %d\n", j + 1, task_inputs[j]);
        }

        // load the system in the halted state at the beginning of the task; Write input to register
        snapshot_load("base");
        CPUState *cpu = libafl_qemu_get_cpu(0);
        if (cpu == NULL)
        {
            printf("Error: CPU is NULL.\n");
        }
        // printf("reg count: %d\n", libafl_qemu_num_regs(cpu));
        for (int j = 0; j < num_tasks; j++)
        {
            // Initialize register_in_32b array
            memset(register_in_32b, 0, sizeof(register_in_32b));

            // Write i to register format
            register_in_32b[0] = task_inputs[j] & 0xFF; // Least significant byte
            register_in_32b[1] = (task_inputs[j] >> 8) & 0xFF;
            register_in_32b[2] = (task_inputs[j] >> 16) & 0xFF;
            register_in_32b[3] = (task_inputs[j] >> 24) & 0xFF; // Most significant byte

            int length = libafl_qemu_read_reg(cpu, j + 1, reg_tmp_val[j]);
            if (length != 4)
            {
                printf("Error: Could not read register\n");
            }
            libafl_qemu_write_reg(cpu, j + 1, register_in_32b);
        }

        libafl_qemu_set_native_breakpoint(start);

        vm_start();
        qemu_main_loop();

        // Now execution is halted at the beginning of the snipped we want to measure. The written input value has been read from the register
        libafl_qemu_remove_native_breakpoint(start);
        libafl_qemu_set_native_breakpoint(end);

        // Write back the original value to the register
        for (int j = 0; j < num_tasks; j++)
        {
            libafl_qemu_write_reg(cpu, j + 1, reg_tmp_val[j]);
        }

        unsigned long start_count = icount_get_raw();
        // printf("Start: %lu\n", start_count);
        vm_start();
        qemu_main_loop();
        libafl_qemu_remove_native_breakpoint(end);

        unsigned long end_count = icount_get_raw();
        // printf("End: %lu\n", end_count);
        inputs[i] = i;
        deltas[i] = end_count - start_count;

        printf("Delta: %lu\n", deltas[i]);

        fprintf(fptr, "%d", inputs[i]);
        for (int j = 0; j < num_tasks; j++)
        {
            fprintf(fptr, ",%d", task_inputs[j]);
        }
        fprintf(fptr, ",%lu\n", deltas[i]);
    }

    // // Write to serial port
    // qemu_chr_fe_write(serial_chr, data, length);

    fclose(fptr);
    free(deltas);
    free(inputs);
    // // Write some text to the file
    // fprintf(fptr, "%lu",delta);
    return 0;
    //========= Instrumentation end
    return qemu_main();
}
#endif