#include "idalloc.hpp"
#include <iostream>
#include <cassert>

LuaDefineType(IdGlobalPool);
LuaDefineType(IdPlayerPool);

static int64_t nthbatch(int64_t n) {
    return int64_t(0x0001000000000000) + n*256;
}

static bool ranges_equal(const std::deque<int64_t> &dq, int64_t a, int64_t b, int64_t c) {
    if (dq.size() != 3) return false;
    if (dq[0] != a) return false;
    if (dq[1] != b) return false;
    if (dq[2] != c) return false;
    return true;
}

IdGlobalPool::IdGlobalPool() {
    salvaged_.clear();
    next_batch_ = 0;
    next_id_ = 0;
    queue_fill_ = 10;
}

IdGlobalPool::~IdGlobalPool() {
}

void IdGlobalPool::init_master(int qf) {
    salvaged_.clear();
    next_batch_ = 0x0001000000000000;
    next_id_ = 0x0010000000000000;
    queue_fill_ = qf;
}

void IdGlobalPool::init_synch(int qf) {
    salvaged_.clear();
    next_batch_ = 0;
    next_id_ = 0x001E000000000000;
    queue_fill_ = qf;
}

int64_t IdGlobalPool::get_one() {
    return next_id_++;
}

int64_t IdGlobalPool::get_batch() {
    int64_t batch;
    if (salvaged_.empty()) {
        if (next_batch_ == 0) {
            batch = 0;
        } else {
            batch = next_batch_;
            next_batch_ += 256;
        }
    } else {
        batch = salvaged_.back();
        salvaged_.pop_back();
    }
    return batch;
}

void IdGlobalPool::salvage(int64_t batch) {
    if (batch == 0) return;
    if (next_batch_ == 0) return;
    if ((batch & 0xFF) >= 128) return;
    salvaged_.push_back(batch);
}

void IdGlobalPool::salvage_thread(lua_State *L) {
    salvage(lua_getnextid(L));
    lua_setnextid(L, 0);
}

int64_t IdGlobalPool::alloc_id_for_thread(lua_State *L) {
    int64_t batch = lua_getnextid(L);
    if (batch != 0) {
        int64_t id = batch;
        batch += 1;
        if ((batch & 0xFF) == 0) batch = 0;
        lua_setnextid(L, batch);
        return id;
    } else {
        return get_one();
    }
}

IdPlayerPool::IdPlayerPool(IdGlobalPool *gp) {
    global_ = gp;
}

IdPlayerPool::~IdPlayerPool() {
}

void IdPlayerPool::purge() {
    ranges_.clear();
}

void IdPlayerPool::refill() {
    while (int(ranges_.size()) < global_->queue_fill()) {
        int64_t batch = global_->get_batch();
        if (batch == 0) break;
        ranges_.push_back(batch);
    }
}

void IdPlayerPool::unqueue() {
    while (!ranges_.empty()) {
        global_->salvage(ranges_.front());
        ranges_.pop_front();
    }
}

int64_t IdPlayerPool::get_batch() {
    while (int(ranges_.size()) < global_->queue_fill() + 1) {
        int64_t batch = global_->get_batch();
        if (batch == 0) break;
        ranges_.push_back(batch);
    }
    if (ranges_.empty()) {
        return 0;
    } else {
        int64_t batch = ranges_.front();
        ranges_.pop_front();
        return batch;
    }
}

void IdPlayerPool::salvage_thread(lua_State *L) {
    global_->salvage_thread(L);
}

void IdPlayerPool::prepare_thread(lua_State *L) {
    global_->salvage_thread(L);
    lua_setnextid(L, get_batch());
}

LuaDefine(cunittests_idalloc, "c") {
    IdGlobalPool gp;
    IdPlayerPool pp(&gp);

    // Synchronous pools produce IDs starting at 0x001E000000000000
    gp.init_synch(3);
    assert(gp.get_one() == 0x001E000000000000);
    assert(gp.get_one() == 0x001E000000000001);
    assert(gp.get_one() == 0x001E000000000002);

    // Master pools produce IDs starting at 0x0010000000000000
    gp.init_master(3);
    assert(gp.get_one() == 0x0010000000000000);
    assert(gp.get_one() == 0x0010000000000001);
    assert(gp.get_one() == 0x0010000000000002);

    // Synchronous pools produce only null batches.
    gp.init_synch(3);
    assert(gp.get_batch() == 0);
    assert(gp.get_batch() == 0);
    gp.salvage(nthbatch(5));
    assert(gp.get_batch() == 0);
    
    // Simple fetch batches with a few salvages.
    gp.init_master(3);
    assert(gp.get_batch() == nthbatch(0));
    assert(gp.get_batch() == nthbatch(1));
    assert(gp.get_batch() == nthbatch(2));
    gp.salvage(nthbatch(182));
    gp.salvage(nthbatch(183));
    assert(gp.get_batch() == nthbatch(183));
    assert(gp.get_batch() == nthbatch(182));
    assert(gp.get_batch() == nthbatch(3));
    
    // Salvage of a zero-batch does nothing.
    gp.init_master(3);
    assert(gp.get_batch() == nthbatch(0));
    assert(gp.get_batch() == nthbatch(1));
    gp.salvage(0);
    assert(gp.get_batch() == nthbatch(2));

    // Salvage of a partial batch.
    gp.init_master(3);
    assert(gp.get_batch() == nthbatch(0));
    assert(gp.get_batch() == nthbatch(1));
    gp.salvage(nthbatch(142) + 10);
    assert(gp.get_batch() == nthbatch(142) + 10);
    assert(gp.get_batch() == nthbatch(2));

    // Salvage of a half-empty batch does nothing.
    gp.init_master(3);
    assert(gp.get_batch() == nthbatch(0));
    assert(gp.get_batch() == nthbatch(1));
    gp.salvage(nthbatch(142) + 145);
    assert(gp.get_batch() == nthbatch(2));

    // In the synchronous model, refill should do nothing.
    pp.purge();
    gp.init_synch(3);
    pp.refill();
    assert(pp.size() == 0);
    assert(pp.get_batch() == 0);
    assert(pp.size() == 0);
    assert(pp.get_batch() == 0);

    // Test refill from master.
    pp.purge();
    gp.init_master(3);
    pp.refill();
    assert(ranges_equal(pp.ranges_, nthbatch(0), nthbatch(1), nthbatch(2)));

    // Now test that get_batch keeps the pool filled from master.
    assert(pp.get_batch() == nthbatch(0));
    assert(ranges_equal(pp.ranges_, nthbatch(1), nthbatch(2), nthbatch(3)));

    // Test unqueueing the batches.
    assert(gp.get_batch() == nthbatch(4));
    assert(gp.get_batch() == nthbatch(5));
    pp.unqueue();
    assert(gp.get_batch() == nthbatch(3));
    assert(gp.get_batch() == nthbatch(2));
    assert(gp.get_batch() == nthbatch(1));
    assert(gp.get_batch() == nthbatch(6));

    // Try preparing a thread and salvaging a thread.
    pp.purge();
    gp.init_master(3);
    lua_setnextid(L, 0);
    pp.prepare_thread(L);
    assert(lua_getnextid(L) == nthbatch(0));
    lua_setnextid(L, 0);
    pp.prepare_thread(L);
    assert(lua_getnextid(L) == nthbatch(1));

    // Try salvaging the pool from the thread.
    pp.salvage_thread(L);
    assert(lua_getnextid(L) == 0);
    assert(gp.get_batch() == nthbatch(1));

    // Allocate IDs from inside a thread.
    lua_setnextid(L, 0xFD);
    gp.init_master(3);
    assert(gp.alloc_id_for_thread(L) == 0xFD);
    assert(gp.alloc_id_for_thread(L) == 0xFE);
    assert(gp.alloc_id_for_thread(L) == 0xFF);
    assert(gp.alloc_id_for_thread(L) == 0x0010000000000000);
    assert(lua_getnextid(L) == 0);
    
    return 0;
}