#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) {
  ttable_t output = calloc(TABLE_SIZE, sizeof(tbucket_t));
  memcpy(table, &output, sizeof(output));
}

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

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

  int count = bucket->items_filled;
  for (int i = 1; i < count; i++) {
    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 + (key & TT_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_empty_ttable() {
  ttable_t table;
  ttable_init(&table);

  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);

  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);

  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);

  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 FAH_test_bucket_collision() {
//   ttable_t table;
//   ttable_init(&table);
//
//
//
//     TranspositionTable tt;
//     tt_init(&tt, 1); // one bucket
//
//     tt_store(&tt, 1, 10, 10, 5, TT_EXACT);
//     tt_store(&tt, 2, 20, 20, 5, TT_EXACT);
//     tt_store(&tt, 3, 30, 30, 5, TT_EXACT);
//     tt_store(&tt, 4, 40, 40, 5, TT_EXACT);
//
//     assert(tt_probe(&tt, 1) != NULL);
//     assert(tt_probe(&tt, 2) != NULL);
//     assert(tt_probe(&tt, 3) != NULL);
//     assert(tt_probe(&tt, 4) != NULL);
//
//     tt_destroy(&tt);
// }

// void test_replace_shallowest(void) {
//     TranspositionTable tt;
//     tt_init(&tt, 1);
//
//     tt_store(&tt, 1, 0, 0, 10, TT_EXACT);
//     tt_store(&tt, 2, 0, 0, 20, TT_EXACT);
//     tt_store(&tt, 3, 0, 0, 30, TT_EXACT);
//     tt_store(&tt, 4, 0, 0, 40, TT_EXACT);
//
//     // Bucket is full.
//
//     tt_store(&tt, 5, 0, 0, 25, TT_EXACT);
//
//     assert(tt_probe(&tt, 1) == NULL);
//
//     assert(tt_probe(&tt, 2) != NULL);
//     assert(tt_probe(&tt, 3) != NULL);
//     assert(tt_probe(&tt, 4) != NULL);
//     assert(tt_probe(&tt, 5) != NULL);
//
//     tt_destroy(&tt);
// }

// void test_many_entries(void)
// {
//     TranspositionTable tt;
//     tt_init(&tt, 1 << 16);
//
//     for (uint64_t i = 0; i < 50000; i++)
//     {
//         tt_store(&tt,
//                  i * 7919,
//                  i,
//                  i,
//                  i % 64,
//                  TT_EXACT);
//     }
//
//     for (uint64_t i = 0; i < 50000; i++)
//     {
//         TTEntry *e = tt_probe(&tt, i * 7919);
//
//         if (e)
//         {
//             assert(e->key == i * 7919);
//         }
//     }
//
//     tt_destroy(&tt);
// }

#endif