//
// Created by martin on 13.11.20.
//

#ifndef STHSTRY_H
#define STHSTRY_H

#ifdef __cplusplus
extern "C" {
#endif

#include <malloc.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdlib.h>
#include <string.h>

typedef struct sthstry sthstry;

sthstry* sthstry_create();
void sthstry_destroy(sthstry* sh);

void sthstry_register(sthstry* sh, void* ptr, size_t size);
void sthstry_store(sthstry* sh);
void sthstry_restore(sthstry* sh, int index);
int sthstry_get_num_states(sthstry* sh);
void sthstry_reset_storage(sthstry* sh);

#define STHSTRY_BUFFER_SIZE 1024 * 1024 * 100

#ifdef STHSTRY_IMPLEMENTATION

typedef struct stentry {
  void* ptr;
  size_t size;
  size_t offset;
} stentry;

typedef struct sthstry {
  char* storage;
  size_t storage_size;
  size_t state_block_size;
  int num_state_blocks;

  stentry* registry;
  int registry_capacity;
  int num_registry_entries;
  bool registry_locked;
} sthstry;

// Forward declarations
bool sthstry_resize_registry(sthstry* sh, int count);

sthstry* sthstry_create() {
  sthstry* result = (sthstry*)calloc(1, sizeof(sthstry));
  result->storage = (char*)calloc(1, STHSTRY_BUFFER_SIZE);
  result->storage_size = STHSTRY_BUFFER_SIZE;
  result->state_block_size = 0;
  result->num_state_blocks = 0;

  result->registry_capacity = 0;
  result->num_registry_entries = 0;
  result->registry_locked = false;
  sthstry_resize_registry(result, 10);
}

void sthstry_destroy(sthstry* sh) {
  free(sh->storage);
  free(sh->registry);
  free(sh);
}

bool sthstry_resize_registry(sthstry* sh, int count) {
  if (count == 0) {
    free(sh->registry);
    sh->registry = NULL;
    sh->registry_capacity = 0;
    return true;
  }

  stentry* result = (stentry*)calloc(count, sizeof(stentry));
  if (result == NULL) {
    fprintf(stderr, "Error: could not allocate memory for registry!\n");
  }
  memcpy(result, sh->registry, sizeof(stentry) * sh->num_registry_entries);
  free(sh->registry);
  sh->registry = result;
  sh->registry_capacity = count;
}

void sthstry_register(sthstry* sh, void* ptr, size_t size) {
  if (sh->registry_locked) {
    fprintf(
        stderr,
        "Error: cannot register entry to locked state history registry!\n");
    abort();
  }

  if (sh->registry_capacity == sh->num_registry_entries) {
    sthstry_resize_registry(sh, sh->registry_capacity * 2);
  }

  stentry* entry = &sh->registry[sh->num_registry_entries++];
  entry->offset = sh->state_block_size;
  entry->ptr = ptr;
  entry->size = size;
  sh->state_block_size += entry->size;
}

void sthstry_store(sthstry* sh) {
  sh->registry_locked = true;

  size_t needed_size = (sh->num_state_blocks + 1) * sh->state_block_size;
  if (needed_size > sh->storage_size) {
    fprintf(
        stderr,
        "Error: cannot store another block. Storage exceeded by %ld bytes.\n",
        needed_size - sh->storage_size);
    abort();
  }

  char* state_block = sh->storage + sh->num_state_blocks * sh->state_block_size;
  for (int i = 0; i < sh->num_registry_entries; i++) {
    stentry* entry = &sh->registry[i];
    memcpy((void*)(state_block + entry->offset), entry->ptr, entry->size);
  }

  sh->num_state_blocks++;
}

int sthstry_get_num_states(sthstry* sh) { return sh->num_state_blocks; }

void sthstry_restore(sthstry* sh, int index) {
  if (index >= sh->num_state_blocks) {
    fprintf(
        stderr,
        "Error: cannot restore state block %d only have %d blocks stored.\n",
        index,
        sh->num_state_blocks);
  }
  char* state_block = sh->storage + index * sh->state_block_size;
  for (int i = 0; i < sh->num_registry_entries; i++) {
    stentry* entry = &sh->registry[i];
    memcpy(entry->ptr, (void*)(state_block + entry->offset), entry->size);
  }
}

void sthstry_reset_storage(sthstry* sh) { sh->num_state_blocks = 0; }

#endif  // STHSTRY_IMPLEMENTATION

#ifdef __cplusplus
}
#endif

#endif  // STHSTRY_H