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

#include "engine/ttable.h"
#include "random.h"

void zobrist_init() {
  for (int i = 0; i < PIECE_TYPE_COUNT; i++) {
    for (int j = 0; j < 64; j++) {
      zobrist_keys.piece[i][j] = random_u64();
    }
  }
  for (int i = 0; i < CASTLE_COUNT; i++) {
    zobrist_keys.castle[i] = random_u64();
  }
  for (int i = 0; i < TURN_COUNT; i++) {
    zobrist_keys.turn[i] = random_u64();
  }
  for (int i = 0; i < 64; i++) {
    zobrist_keys.en_passant[i] = random_u64();
  }
}

u64 zobrist_hash(position_t position) {
  u64 output = 0;
  for (int i = 0; i < 64; i++) {
    int piece_type = find_piece_on_square(&position, i);
    output ^= zobrist_keys.piece[piece_type][i];
  }

  if (position.castling & WHITE_SHORT_CASTLE) {
    output ^= zobrist_keys.castle[0];
  }
  if (position.castling & WHITE_LONG_CASTLE) {
    output ^= zobrist_keys.castle[1];
  }
  if (position.castling & BLACK_SHORT_CASTLE) {
    output ^= zobrist_keys.castle[2];
  }
  if (position.castling & BLACK_LONG_CASTLE) {
    output ^= zobrist_keys.castle[3];
  }

  output ^= zobrist_keys.turn[position.turn];
  output ^= zobrist_keys.en_passant[position.passantable_file];

  return output;
}

void ttable_init(ttable_t *table, size_t size_in_binary_log) {
  assert(size_in_binary_log < 64);
  int size = 1ULL << size_in_binary_log;
  table->buckets = calloc(size, sizeof(tbucket_t));
  table->mask = size - 1;
}

void ttable_deinit(ttable_t table) {
  free(table.buckets);
}

void __ttable_insert(ttable_t table, struct ttable_insert_args args) {
  tbucket_t *bucket = table.buckets + (args.key & table.mask);
  tentry_t *best = bucket->entries;

  int ci = 0;

  for (int i = 0; i < TT_BUCKET_SIZE; i++) {
    if (atomic_load(&bucket->entries[i].key) == args.key) {
      best = bucket->entries + i;
      break;
    }
    if (bucket->entries[i].depth >= best->depth) continue;
    best = bucket->entries + i;
  }

  tentry_t object;
  object.best_move = args.move;
  object.score = args.score;
  object.eval = args.eval;
  object.depth = args.depth;
  object.flag = args.flag;
  object.generation = args.generation;

  memcpy(best, &object, sizeof(object));
  atomic_store(&best->key, args.key);
}

tentry_t *ttable_find(ttable_t table, u64 key) {
  tbucket_t *bucket = table.buckets + (key & table.mask);

  for (int i = 0; i < TT_BUCKET_SIZE; i++) {
    if (atomic_load(&(bucket->entries + i)->key) != key) continue;
    return bucket->entries + i;
  }

  return NULL;
}

#ifdef TEST_MOD

bool test_ttable_empty() {
  ttable_t table;
  ttable_init(&table, 4);

  if (ttable_find(table, 123456789ULL) != NULL) return false;

  ttable_deinit(table);
  return true;
}

bool test_ttable_insert_and_find() {
  ttable_t table;
  ttable_init(&table, 4);

  u64 key = 0x123456789ABCDEF0ULL;
  ttable_insert(table, .key = key, .depth = 8, .flag = TT_EXACT);

  tentry_t *e = ttable_find(table, key);

  assert(e != NULL);
  assert(e->key == key);
  assert(e->depth == 8);
  assert(e->flag == TT_EXACT);

  ttable_deinit(table);

  return true;
}

bool test_ttable_wrong_key() {
  ttable_t table;
  ttable_init(&table, 4);

  u64 key = 0x123456789ABCDEF0ULL;
  ttable_insert(table, .key = key, .depth = 8, .flag = TT_EXACT);
  assert(ttable_find(table, --key) == NULL);

  ttable_deinit(table);
  return true;
}

bool test_ttable_update_existing() {
  ttable_t table;
  ttable_init(&table, 4);

  u64 key = 987654321ULL;
  ttable_insert(table, .key = key, .depth = 10, .flag = TT_LOWERBOUND);
  ttable_insert(table, .key = key, .depth = 11, .flag = TT_EXACT);

  tentry_t *e = ttable_find(table, key);
  assert(e != NULL);
  assert(e->key == key);
  assert(e->depth == 11);
  assert(e->flag == TT_EXACT);

  ttable_deinit(table);
  return true;
}

bool test_bucket_collision() {
  ttable_t table;
  ttable_init(&table, 2);

  ttable_insert(table, .key = 1, .depth = 10, .flag = TT_EXACT);
  ttable_insert(table, .key = 2, .depth = 20, .flag = TT_EXACT);
  ttable_insert(table, .key = 3, .depth = 30, .flag = TT_EXACT);
  ttable_insert(table, .key = 4, .depth = 40, .flag = TT_EXACT);

  if (ttable_find(table, 1) == NULL) return false;
  if (ttable_find(table, 2) == NULL) return false;
  if (ttable_find(table, 3) == NULL) return false;
  if (ttable_find(table, 4) == NULL) return false;

  ttable_deinit(table);
  return true;
}

bool test_ttable_replace_shallowest() {
  ttable_t table;
  ttable_init(&table, 0);

  ttable_insert(table, .key = 1, .depth = 10, .flag = TT_EXACT);
  ttable_insert(table, .key = 2, .depth = 20, .flag = TT_EXACT);
  ttable_insert(table, .key = 3, .depth = 30, .flag = TT_EXACT);
  ttable_insert(table, .key = 4, .depth = 40, .flag = TT_EXACT);

  // Bucket is full.
  ttable_insert(table, .key = 5, .depth = 25, .flag = TT_EXACT);

  if (ttable_find(table, 1) != NULL) return false;

  if (ttable_find(table, 2) == NULL) return false;
  if (ttable_find(table, 3) == NULL) return false;
  if (ttable_find(table, 4) == NULL) return false;
  if (ttable_find(table, 5) == NULL) return false;

  ttable_deinit(table);
  return true;
}

bool test_ttable_many_entries() {
  ttable_t table;
  ttable_init(&table, 16);

  for (u64 i = 0; i < 50000; i++) {
    ttable_insert(table, .key = i * 7919, .depth = i % 64, .flag = TT_EXACT);
  }

  for (u64 i = 0; i < 50000; i++) {
    tentry_t *e = ttable_find(table, i * 7919);

    if (e == NULL) continue;
    if (e->key != i * 7919) return false;
  }

  ttable_deinit(table);
  return true;
}

#endif