#pragma once

#include <forstio/io_codec/rpc.hpp>
#include <forstio/codec/data.hpp>
#include <forstio/codec/id_map.hpp>

#include <AdaptiveCpp/CL/sycl.hpp>

namespace saw {
namespace rmt {
struct Sycl {};
}

template<>
class remote<rmt::Sycl>;

/**
 * Remote data class for the Sycl backend.
 */
template<typename T, typename Encoding, typename Storage>
class remote_data<T, Encoding, Storage, rmt::Sycl> {
private:
	/**
	 * Id representing the remote data
	 */
	id<T> id_;
	/**
	 * Storage for the
	 */
	id_map<T,Encoding,rmt::Sycl>* map_;
public:
	/**
	 * Main constructor
	 */
	remote_data(const id<T>& id, id_map<T, Encoding, rmt::Sycl>& map, cl::sycl::queue& queue__):
		id_{id},
		map_{&map}
	{}

	/**
	 * Wait for the data
	 */
	error_or<data<T,Encoding,Storage>> wait(){
		auto eov = map_->find(id_);
		if(eov.is_error()){
			auto& err = eov.get_error();
			return std::move(err);
		}
		auto& val = eov.get_value();
		std::cout<<"Values Sycl in Map: "<<val->size()<<std::endl;

		{
			auto eocop = val->template copy_to_host<storage::Default>();
			if(eocop.is_error()){
				return eocop;
			}
			return eocop.get_value();
		}
	}

	/**
	 * Request data asynchronously
	 */
	conveyor<data<T,Encoding,Storage>> on_receive(); /// Stopped here
};

/**
 * 
template<typename T, uint64_t N>
class data<schema::Primitive<T,N>, encode::Native, rmt::Sycl> {
public:
	using Schema = schema::Primitive<T,N>;
	using NativeType = typename native_data_type<Schema>::type;
private:
	NativeType val_;
public:
	data(NativeType val__):
		val_{val__}
	{}

	NativeType get(){
		return val_;
	}
};
 */

template<typename T, uint64_t D>
class data<schema::Array<T,D>, encode::Native, rmt::Sycl> {
public:
	using Schema = schema::Array<T,D>;
private:
	uint64_t total_length_;
	data<T,encode::Native,storage::Default>* device_data_;
	cl::sycl::queue* queue_;

	static_assert(is_primitive<T>::value, "Only supports primitives for now");
	static_assert(D==1u, "For now we only support 1D Arrays");
public:
	data(uint64_t size, cl::sycl::queue& q__):
		total_length_{size},
		device_data_{cl::sycl::malloc_device<data<T,encode::Native,storage::Default>>(size, q__)},
		queue_{&q__}
	{
		if(!device_data_){
			total_length_ = 0u;
			return;
		}
		queue_->wait();
	}

	template<typename Encoding, typename Storage>
	data(const data<Schema, Encoding, Storage>& from, cl::sycl::queue& q__):
		total_length_{from.size()},
		device_data_{cl::sycl::malloc_device<data<T,encode::Native,storage::Default>>(from.size(), q__)},
		queue_{&q__}
	{
		if(!device_data_){
			total_length_ = 0u;
			return;
		}
		queue_->template copy<data<T,encode::Native,storage::Default>>(&from.at(0), device_data_, total_length_);
		queue_->wait();
	}
	
	data(const data<Schema, encode::Native, rmt::Sycl>& from):
		total_length_{from.size()},
		device_data_{nullptr},
		queue_{from.queue_}
	{
		if(total_length_ == 0u || !queue_){
			return;
		}
		device_data_ = cl::sycl::malloc_device<data<T,encode::Native,storage::Default>>(from.size(), *queue_);
		// device_data_ = cl::sycl::malloc_device<typename native_data_type<T>::type>(from.size(), *queue_);
		if(!device_data_){
			total_length_ = 0u;
			return;
		}

		queue_->template copy<data<T,encode::Native,storage::Default>>(from.device_data_, device_data_, total_length_);
	}
	
	data(data<Schema, encode::Native, rmt::Sycl>&& rhs):
		total_length_{rhs.total_length_},
		device_data_{rhs.device_data_},
		queue_{rhs.queue_}
	{
		rhs.total_length_ = 0u;
		rhs.device_data_ = nullptr;
		rhs.queue_ = nullptr;
	}

	data<Schema, encode::Native, rmt::Sycl>& operator=(data<Schema, encode::Native, rmt::Sycl>&& rhs){
		total_length_ = rhs.total_length_;
		device_data_ = rhs.device_data_;
		queue_ = rhs.queue_;
		rhs.total_length_ = 0u;
		rhs.device_data_ = nullptr;
		rhs.queue_ = nullptr;
		return *this;
	}

	~data(){
		// free data
		if(device_data_){
			/// SYCL FREE
			cl::sycl::free(device_data_, *queue_);
		}
	}

	/**
	 * Allocate appropriate meta data and then copy to host
	 */
	template<typename Storage>
	error_or<data<Schema, encode::Native, Storage>> copy_to_host() const {
		data<Schema,encode::Native, Storage> data_{total_length_};

		/// TODO Check success
		queue_->template copy<data<T,encode::Native,storage::Default>>(device_data_, &data_.at(0), total_length_);
		queue_->wait();
		return data_;
	}

	data<T, encode::Native, storage::Default>& at(uint64_t i){
		return device_data_[i];
	}

	uint64_t size() const {
		return total_length_;
	}
};

namespace impl {
template<typename Iface, typename Encoding, typename Storage>
struct rpc_id_map_helper {
	static_assert(always_false<Iface, Encoding,Storage>, "Only supports Interface schema types.");
};

template<typename... Members, typename Encoding, typename Storage>
struct rpc_id_map_helper<schema::Interface<Members...>, Encoding, Storage> {
	std::tuple<id_map<typename Members::ValueType::ResponseT, Encoding, Storage>...> maps;
};
}

template<typename T>
class device;

/**
 * Represents a remote Sycl device.
 *
 */
template<>
class device<rmt::Sycl> {
private:
	cl::sycl::queue cmd_queue_;
public:
	/**
	 * Copy data to device
	 */
	template<typename Schema, typename Encoding, typename Storage, typename EncodingHost, typename StorageHost>
	error_or<data<Schema, Encoding, Storage>> copy_to_device(const data<Schema, EncodingHost, StorageHost>>& host_data){
		(void) host_data;
		return make_error<err::not_implemented>();
	}
	
	/**
	 * Copy data to host
	 */
	template<typename Schema, typename Encoding, typename Storage, typename EncodingDevice, typename StorageDevice>
	error_or<data<Schema, Encoding, Storage>> copy_to_host(const data<Schema, EncodingDevice, StorageDevice>>& device_data){
		(void) device_data;
		return make_error<err::not_implemented>();
	}
};

/**
 * Device data transport
 */


/**
 * Rpc Client class for the Sycl backend.
 */
template<typename Iface, typename Encoding, typename Storage>
class rpc_client<Iface, Encoding, Storage, rmt::Sycl> {
public:
private:
	/**
	 * Server this client is tied to
	 */
	rpc_server<Iface, Encoding, rmt::Sycl>* srv_;

	/**
	 * Generated some sort of id for the request.
	 */
public:
	rpc_client(rpc_server<Iface, Encoding, rmt::Sycl>& srv):
		srv_{&srv}
	{}

	/**
	 * Rpc call
	 */
	template<string_literal Name>
	error_or<
		id<
			typename schema_member_type<Name, Iface>::type::ResponseT
		>
	> call(const data_or_id<typename schema_member_type<Name, Iface>::type::RequestT, Encoding, Storage>& input){
		auto next_free_id = srv_->template next_free_id<typename schema_member_type<Name, Iface>::type::ResponseT>();
		return srv_->template call<Name, Storage>(input, next_free_id);
	}

};

/**
 * Rpc Server class for the Sycl backend.
 */
template<typename Iface, typename Encoding>
class rpc_server<Iface, Encoding, rmt::Sycl> {
public:
	using InterfaceCtxT = cl::sycl::queue*;
	using InterfaceT = interface<Iface, Encoding, rmt::Sycl, InterfaceCtxT>;
private:
	/**
	 * Device instance enabling the use of the remote device.
	 */
	device<rmt::Sycl>* device_;

	/**
	 * The interface including the relevant context class.
	 */
	interface<Iface, Encoding, rmt::Sycl, InterfaceCtxT> cl_interface_;

	/**
	 * Basic storage for response data.
	 */
	impl::rpc_id_map_helper<Iface, Encoding, rmt::Sycl> storage_;
public:

	/**
	 * Main constructor
	 */
	rpc_server(device<rmt::Sycl>& dev__, interface<Iface, Encoding, rmt::Sycl, InterfaceCtxT> cl_iface):
		device_{&dev__},
		cl_interface_{std::move(cl_iface)},
		storage_{}
	{}
	
	/**
	 * Ask which id the server prefers as the next one. Only available for fast requests on no roundtrip setups.
	 */
	template<typename T>
	id<T> next_free_id() const {
		return std::get<id_map<T,Encoding,rmt::Sycl>>(storage_.maps).next_free_id();
	}

	template<typename IdT, typename Storage>
	remote_data<IdT, Encoding, Storage, rmt::Sycl> request_data(id<IdT> dat_id){
		return {dat_id, std::get<id_map<IdT,Encoding,rmt::Sycl>>(storage_.maps), device_.cmd_queue_};
	}

	/**
	 * Rpc call based on the name
	 */
	template<string_literal Name, typename ClientAllocation>
	error_or<
		id<
			typename schema_member_type<Name, Iface>::type::ResponseT
		>
	> call(data_or_id<typename schema_member_type<Name, Iface>::type::RequestT, Encoding, ClientAllocation> input, id<typename schema_member_type<Name,Iface>::type::ResponseT> rpc_id){
		using FuncT = typename schema_member_type<Name, Iface>::type;

		/**
		 * Object needed if and only if the provided data type is not an id
		 */
		own<data<typename FuncT::RequestT, Encoding, rmt::Sycl>> dev_tmp_inp = nullptr;
		/**
		 * First check if it's data or an id.
		 * If it's an id, check if it's registered within the storage and retrieve it.
		 */
		auto eoinp = [&,this]() -> error_or<data<typename FuncT::RequestT, Encoding, rmt::Sycl>* > {
			if(input.is_id()){
				// storage_.maps
				auto& inner_map = std::get<id_map<typename FuncT::RequestT, Encoding, rmt::Sycl>> (storage_.maps);
				auto eov = inner_map.find(input.get_id());
				if(eov.is_error()){
					return std::move(eov.get_error());
				}
				return eov.get_value();
			} else {
				auto& client_data = input.get_data();
				dev_tmp_inp = heap<data<typename FuncT::RequestT, Encoding, rmt::Sycl>>(client_data, cmd_queue_);
				cmd_queue_.wait();
				return dev_tmp_inp.get();
			}
		}();
		if(eoinp.is_error()){
			return std::move(eoinp.get_error());
		}
		auto& inp = *(eoinp.get_value());

		auto eod = cl_interface_.template call<Name>(inp, &cmd_queue_);

		if(eod.is_error()){
			return std::move(eod.get_error());
		}

		auto& val = eod.get_value();
		/**
		 * Store returned data in rpc storage
		 */
		auto& inner_map = std::get<id_map<typename schema_member_type<Name, Iface>::type::RequestT, Encoding,rmt::Sycl>> (storage_.maps);
		auto eoid = inner_map.insert_as(std::move(val), rpc_id);
		if(eoid.is_error()){
			return std::move(eoid.get_error());
		}
		return rpc_id;
	}
};


template<>
struct remote_address<rmt::Sycl> {
private:
	remote<rmt::Sycl>* ctx_;

	SAW_FORBID_COPY(remote_address);
	SAW_FORBID_MOVE(remote_address);
public:
	remote_address(remote<rmt::Sycl>& r_ctx):
		ctx_{&r_ctx}
	{}
};

template<>
class remote<rmt::Sycl> {
private:
	SAW_FORBID_COPY(remote);
	SAW_FORBID_MOVE(remote);

public:
	/**
	 * Default constructor
	 */
	remote(){}

	/**
	 * For now we don't need to specify the location since
	 * we just create a default.
	 */
	conveyor<own<remote_address<rmt::Sycl>>> resolve_address(){
		return heap<remote_address<rmt::Sycl>>(*this);
	}

	/**
	 * Connect to a device
	 */
	device<rmt::Sycl> connect_device(const remote_address<rmt::Sycl>&){
		return {};
	}

	/**
	 * Spin up a rpc server
	 */
	template<typename Iface, typename Encoding>
	rpc_server<Iface, Encoding, rmt::Sycl> listen(const device<rmt::Sycl>& dev, typename rpc_server<Iface, Encoding, rmt::Sycl>::InterfaceT iface){
		using RpcServerT = rpc_server<Iface, Encoding, rmt::Sycl>;
		using InterfaceT = typename RpcServerT::InterfaceT;
		return {dev, std::move(iface)};
	}
};

}