#pragma once

#include <iostream>

#include <atomic>
#include <memory>
#include <vector>

#include <dml/dml.hpp>

#include "util/dml-helper.hpp"

namespace dsacache {
    class Cache;

    // the cache task structure will be used to submit and
    // control a cache element, while providing source pointer
    // and size in bytes for submission
    //
    // then the submitting thread may wait on the atomic "result"
    // which will be notified by the cache worker upon processing
    // after which the atomic-bool-ptr active will also become valid
    class CacheData {
    public:
        using dml_handler = dml::handler<dml::mem_copy_operation, std::allocator<uint8_t>>;

    private:
        // data source and size of the block
        uint8_t* src_;
        size_t size_;

        // global reference counting object
        std::atomic<int32_t>* active_;

        // global cache-location pointer
        std::atomic<uint8_t*>* cache_;

        // object-local incomplete cache location pointer
        // which is only available in the first instance
        uint8_t* incomplete_cache_;

        // dml handler vector pointer which is only
        // available in the first instance
        std::unique_ptr<std::vector<dml_handler>> handlers_;

        // deallocates the global cache-location
        // and invalidates it
        void Deallocate();

        // checks whether there are at least two
        // valid references to this object which
        // is done as the cache always has one
        // internally to any living instance
        bool Active() const;

        friend Cache;
    public:
        CacheData(uint8_t* data, const size_t size);
        CacheData(const CacheData& other);
        ~CacheData();

        // waits on completion of caching operations
        // for this task and is safe to be called in
        // any state of the object
        void WaitOnCompletion();

        // returns the cache data location for this
        // instance which is valid as long as the
        // instance is alive
        uint8_t* GetDataLocation() const;
    };
}

inline void dsacache::CacheData::WaitOnCompletion() {
    // the cache data entry can be in two states
    // either it is the original one which has not
    // been waited for in which case the handlers
    // are non-null or it is not

    if (handlers_ == nullptr) {
        std::cout << "[-] Waiting on cache-var-update for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;

        // when no handlers are attached to this cache entry we wait on a
        // value change for the cache structure from nullptr to non-null
        // which will either go through immediately if the cache is valid
        // already or wait until the handler-owning thread notifies us

        cache_->wait(nullptr);

        std::cout << "[+] Finished waiting on cache-var-update for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;
    }
    else {
        // when the handlers are non-null there are some DSA task handlers
        // available on which we must wait here

        std::cout << "[-] Waiting on handlers for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;

        // abort is set if any operation encountered an error

        bool abort = false;

        for (auto& handler : *handlers_) {
            auto result = handler.get();

            if (result.status != dml::status_code::ok) {
                std::cerr << "[x] Encountered bad status code for operation: " << dml::StatusCodeToString(result.status) << std::endl;

                // if one of the copy tasks failed we abort the whole task
                // after all operations are completed on it

                abort = true;
            }
        }

        // the handlers are cleared after all have completed

        handlers_ = nullptr;

        // now we act depending on whether an abort has been
        // called for which signals operation incomplete

        if (abort) {
            // store nullptr in the cache location

            cache_->store(nullptr);

            // then free the now incomplete cache

            // TODO: it would be possible to salvage the
            // TODO: operation at this point but this
            // TODO: is quite complicated so we just abort

            numa_free(incomplete_cache_, size_);
        }
        else {
            std::cout << "[+] Finished waiting on handlers for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;

            // incomplete cache is now safe to use and therefore we
            // swap it with the global cache state of this entry
            // and notify potentially waiting threads

            cache_->store(incomplete_cache_);
        }

        // as a last step all waiting threads must
        // be notified (copies of this will wait on value
        // change of the cache) and the incomplete cache
        // is cleared to nullptr as it is not incomplete

        cache_->notify_all();
        incomplete_cache_ = nullptr;
    }
}

dsacache::CacheData::CacheData(uint8_t* data, const size_t size) {
    std::cout << "[-] New CacheData 0x" << std::hex << (uint64_t)data << std::dec << std::endl;

    src_ = data;
    size_ = size;
    active_ = new std::atomic<int32_t>(1);
    cache_ = new std::atomic<uint8_t*>();
    incomplete_cache_ = nullptr;
    handlers_ = std::make_unique<std::vector<dml_handler>>();
}

dsacache::CacheData::CacheData(const dsacache::CacheData& other) {
    std::cout << "[-] Copy Created for CacheData 0x" << std::hex << (uint64_t)other.src_ << std::dec << std::endl;

    // we copy the ptr to the global atomic reference counter
    // and increase the amount of active references

    active_ = other.active_;
    const int current_active = active_->fetch_add(1);

    // source and size will be copied too
    // as well as the reference to the global
    // atomic cache pointer

    src_ = other.src_;
    size_ = other.size_;
    cache_ = other.cache_;

    // incomplete cache and handlers will not
    // be copied because only the first instance
    // will wait on the completion of handlers

    incomplete_cache_ = nullptr;
    handlers_ = nullptr;
}

dsacache::CacheData::~CacheData() {
    std::cout << "[-] Destructor for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;

    // if this is the first instance of this cache structure
    // and it has not been waited on and is now being destroyed
    // we must wait on completion here to ensure the cache
    // remains in a valid state

    if (handlers_ != nullptr) {
        WaitOnCompletion();
    }

    // due to fetch_sub returning the preivously held value
    // we must subtract one locally to get the current value

    const int32_t v = active_->fetch_sub(1) - 1;

    // if the returned value is zero or lower
    // then we must execute proper deletion
    // as this was the last reference

    if (v <= 0) {
        std::cout << "[!] Full Destructor for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;

        Deallocate();

        delete active_;
        delete cache_;
    }
}

void dsacache::CacheData::Deallocate() {
    std::cout << "[!] Deallocating for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;

    // although deallocate should only be called from
    // a safe context to do so, it can not hurt to
    // defensively perform the operation atomically

    uint8_t* cache_local = cache_->exchange(nullptr);
    if (cache_local != nullptr) numa_free(cache_local, size_);
}

uint8_t* dsacache::CacheData::GetDataLocation() const {
    return cache_->load();
}

bool dsacache::CacheData::Active() const {
    // this entry is active if more than one
    // reference exists to it, as the Cache
    // will always keep one internally until
    // the entry is cleared from cache

    return active_->load() > 1;
}