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


struct BlockHeader {
    BlockHeader *prev_;
    BlockHeader *next_;
    int32_t sanity_; // 0x12345678 only when linked.
    size_t size_;
    bool in_use_;
    void *snapshot_;
};

size_t blocksize(size_t base) {
    return base + sizeof(BlockHeader);
}

BlockHeader *pointer_to_header(void *ptr) {
    return (BlockHeader*)(((char *)ptr) - sizeof(BlockHeader));
}

void *header_to_pointer(BlockHeader *blk) {
    return (void *)(((char *)blk) + sizeof(BlockHeader));
}

class LuaSnapData {
public:
    lua_State *state_;
    bool have_snapshot_;
    BlockHeader sentinel_;

    LuaSnapData();
    ~LuaSnapData();
    lua_State *state() const { return state_; }
    bool have_snapshot() const { return have_snapshot_; }
    void snapshot();
    void rollback();
    void link(BlockHeader *blk);
    void unlink(BlockHeader *blk);
    void *allocf(void *ptr, size_t osize, size_t nsize);
};

void *lsd_alloc_f(void *lsd, void *ptr, size_t osize, size_t nsize) {
    return ((LuaSnapData*)lsd)->allocf(ptr, osize, nsize);
}

LuaSnapData::LuaSnapData() {
    sentinel_.prev_ = &sentinel_;
    sentinel_.next_ = &sentinel_;
    sentinel_.size_ = 0;
    sentinel_.in_use_ = false;
    sentinel_.snapshot_ = 0;
    have_snapshot_ = false;
    state_ = lua_newstate(lsd_alloc_f, (void*)this);
}

LuaSnapData::~LuaSnapData() {
    std::cerr << "LuaSnapData destructor not implemented." << std::endl;
    exit(1);
}

void LuaSnapData::snapshot() {
    assert(!have_snapshot_);
    for (BlockHeader *blk = sentinel_.next_; blk != &sentinel_; blk = blk->next_) {
        assert(blk->in_use_);
        blk->snapshot_ = malloc(blk->size_);
        memcpy(blk->snapshot_, header_to_pointer(blk), blk->size_);
    }
    have_snapshot_ = true;
}

void LuaSnapData::rollback() {
    assert(have_snapshot_);
    for (BlockHeader *blk = sentinel_.next_; blk != &sentinel_; ) {
        BlockHeader *next = blk->next_;
        if (blk->snapshot_) {
            memcpy(header_to_pointer(blk), blk->snapshot_, blk->size_);
            free(blk->snapshot_);
            blk->snapshot_ = nullptr;
            blk->in_use_ = true;
        } else {
            assert(blk->in_use_);
            unlink(blk);
            free(blk);
        }
        blk = next;
    }
    have_snapshot_ = false;
}

void LuaSnapData::link(BlockHeader *blk) {
    blk->prev_ = sentinel_.prev_;
    blk->next_ = &sentinel_;
    blk->prev_->next_ = blk;
    blk->next_->prev_ = blk;
    blk->sanity_ = 0x12345678;
}

void LuaSnapData::unlink(BlockHeader *blk) {
    blk->prev_->next_ = blk->next_;
    blk->next_->prev_ = blk->prev_;
    blk->prev_ = nullptr;
    blk->next_ = nullptr;
    blk->sanity_ = 0;
}

void *LuaSnapData::allocf(void *ptr, size_t osize, size_t nsize) {
    if (nsize == 0) {
        if (ptr == 0) return nullptr;
        BlockHeader *blk = pointer_to_header(ptr);
        assert(blk->sanity_ = 0x12345678);
        assert(blk->in_use_);
        if (blk->snapshot_ == nullptr) {
            unlink(blk);
            free(blk);
        } else {
            blk->in_use_ = false;
        }
        return nullptr;
    } else if (ptr == NULL) {
        BlockHeader *blk = (BlockHeader*)malloc(blocksize(nsize));
        link(blk);
        blk->size_ = nsize;
        blk->in_use_ = true;
        blk->snapshot_ = nullptr;
        return header_to_pointer(blk);
    } else {
        BlockHeader *blk = pointer_to_header(ptr);
        assert(blk->sanity_ = 0x12345678);
        assert(blk->in_use_);
        if (blk->snapshot_ == nullptr) {
            unlink(blk);
            blk = (BlockHeader*)realloc(blk, blocksize(nsize));
            blk->size_ = nsize;
            link(blk);
            return header_to_pointer(blk);
        } else {
            BlockHeader *nblk = (BlockHeader*)malloc(blocksize(nsize));
            memcpy(nblk, blk, (blk->size_ < nsize) ? blk->size_ : nsize);
            blk->in_use_ = false;
            link(nblk);
            nblk->size_ = nsize;
            nblk->in_use_ = true;
            nblk->snapshot_ = nullptr;
            return header_to_pointer(nblk);
        }
    }
}


LuaSnap::LuaSnap() {
    data_ = new LuaSnapData;
    state_ = data_->state();
}

LuaSnap::~LuaSnap() {
    delete data_;
}

bool LuaSnap::have_snapshot() const {
    return data_->have_snapshot();
}

void LuaSnap::snapshot() {
    data_->snapshot();
}

void LuaSnap::rollback() {
    data_->rollback();
}