#pragma once

#include "data.hpp"

namespace saw {
namespace encode {
struct NativeRaw {
	static constexpr string_literal name = "encode::NativeRaw";
};
}

template<typename T>
struct raw_native_array_type_helper {
	using type = data<T,encode::NativeRaw>;
};

template<typename T, uint64_t N>
struct raw_native_array_type_helper<schema::Primitive<T,N>> {
	using Type = typename native_data_type<schema::Primitive<T,N>>::type;
};

template<typename T, typename S>
struct raw_native_array_type_helper<schema::MixedPrecision<T,S>> {
	using Type = typename native_data_type<typename schema::MixedPrecision<T,S>::StorageSchema>::type;
};

template<>
class data<schema::Void, encode::NativeRaw> {
public:
	using Schema = schema::Void;
	using MetaSchema = schema::Void;
};

template<>
class data<schema::Bool, encode::NativeRaw> {
public:
	using Schema = schema::Bool;
	using MetaSchema = schema::Void;
private:
	bool value_;
public:
	data():value_{false}{}
	data(data<MetaSchema, encode::NativeRaw>):value_{false}{}

	SAW_DEFAULT_COPY(data);
	SAW_DEFAULT_MOVE(data);

	void set(bool val){
		value_ = val;
	}

	bool get() const {
		return value_;
	}
};

template<typename T, size_t N>
class data<schema::Primitive<T,N>, encode::NativeRaw> {
public:
	using Schema = schema::Primitive<T,N>;
	using MetaSchema = typename meta_schema<Schema>::MetaSchema;
private:
	typename native_data_type<Schema>::type value_;
public:
	data():value_{}{}
	data(data<MetaSchema, encode::NativeRaw>):value_{}{}

	SAW_DEFAULT_COPY(data);
	SAW_DEFAULT_MOVE(data);

	data(typename native_data_type<Schema>::type value__):
		value_{std::move(value__)}{}

	void set(typename native_data_type<Schema>::type val){
		value_ = val;
	}

	typename native_data_type<Schema>::type get() const {return value_;}

	data<Schema, encode::NativeRaw> operator*(const data<Schema, encode::NativeRaw>& rhs)const{
		return {get() * rhs.get()};
	}
	
	data<Schema, encode::NativeRaw> operator/(const data<Schema, encode::NativeRaw>& rhs)const{
		return {get() / rhs.get()};
	}
	
	data<Schema, encode::NativeRaw> operator+(const data<Schema, encode::NativeRaw>& rhs)const{
		return {get() + rhs.get()};
	}
	
	data<Schema, encode::NativeRaw>& operator+=(const data<Schema, encode::NativeRaw>& rhs)const{
		value_ += rhs.get();
		return *this;
	}
	
	data<Schema, encode::NativeRaw>& operator-=(const data<Schema, encode::NativeRaw>& rhs)const{
		value_ -= rhs.get();
		return *this;
	}

	data<Schema, encode::NativeRaw> operator-(const data<Schema, encode::NativeRaw>& rhs)const{
		return {get() - rhs.get()};
	}

	template<typename Enc>
	bool operator==(const data<Schema, Enc>& rhs)const{
		return get() == rhs.get();
	}

	template<typename Enc>
	bool operator<(const data<Schema, Enc>& rhs) const {
		return get() < rhs.get();
	}

	/**
	 * Casts 
	 */
	template<typename Target>
	data<Target, encode::NativeRaw> cast_to() const {
		auto raw_to = static_cast<typename saw::native_data_type<Target>::type>(value_);
		return {raw_to};
	}
};

template<typename T, uint64_t N>
class data<schema::Ref<schema::Primitive<T,N>>, encode::NativeRaw> {
public:
	using ReferencedSchema = schema::Primitive<T,N>;
	using Schema = schema::Ref<ReferencedSchema>;
	using MetaSchema = typename meta_schema<Schema>::MetaSchema;
private:
	ref<typename native_data_type<ReferencedSchema>::type > value_;
public:
	data(ref<typename native_data_type<ReferencedSchema >::type > value__):
		value_{value__}
	{}

	SAW_DEFAULT_COPY(data);
	SAW_FORBID_MOVE(data);

	void set(typename native_data_type<ReferencedSchema>::type val){
		value_() = val;
	}

	typename native_data_type<ReferencedSchema>::type get() const {
		return value_();
	}
	
	data<ReferencedSchema, encode::NativeRaw> operator*(const data<Schema, encode::NativeRaw>& rhs)const{
		return {get() * rhs.get()};
	}
	
	data<ReferencedSchema, encode::NativeRaw> operator/(const data<Schema, encode::NativeRaw>& rhs)const{
		return {get() / rhs.get()};
	}
	
	data<ReferencedSchema, encode::NativeRaw> operator+(const data<Schema, encode::NativeRaw>& rhs)const{
		return {get() + rhs.get()};
	}
	
	data<Schema, encode::NativeRaw>& operator+=(const data<Schema, encode::NativeRaw>& rhs)const{
		value_ += rhs.get();
		return *this;
	}
	
	data<Schema, encode::NativeRaw>& operator-=(const data<Schema, encode::NativeRaw>& rhs)const{
		value_ -= rhs.get();
		return *this;
	}

	data<ReferencedSchema, encode::NativeRaw> operator-(const data<Schema, encode::NativeRaw>& rhs)const{
		return {get() - rhs.get()};
	}

	template<typename Enc>
	bool operator==(const data<Schema, Enc>& rhs)const{
		return get() == rhs.get();
	}
	
	template<typename Enc>
	bool operator==(const data<ReferencedSchema, Enc>& rhs)const{
		return get() == rhs.get();
	}

	template<typename Enc>
	bool operator<(const data<Schema, Enc>& rhs) const {
		return get() < rhs.get();
	}
	
	template<typename Enc>
	bool operator<(const data<ReferencedSchema, Enc>& rhs) const {
		return get() < rhs.get();
	}
};

/**
 * Mixed precision class for native formats
 */
template<typename TA, uint64_t NA, typename TB, uint64_t NB> 
class data<schema::MixedPrecision<schema::Primitive<TA,NA>, schema::Primitive<TB,NB>>, encode::NativeRaw>{
public:
	using Schema = schema::MixedPrecision<schema::Primitive<TA,NA>, schema::Primitive<TB,NB>>;
	using MetaSchema = typename meta_schema<Schema>::MetaSchema;
private:
	data<typename Schema::StorageSchema, encode::NativeRaw> value_;
public:
	data():value_{}{}
	data(data<MetaSchema, encode::NativeRaw>):value_{}{}

	data(typename saw::native_data_type<typename Schema::InterfaceSchema>::type val__):value_{static_cast<typename saw::native_data_type<typename Schema::StorageSchema>::type>(val__)}{}

	typename saw::native_data_type<typename Schema::InterfaceSchema>::type get() const {
		return value_.template cast_to<typename Schema::InterfaceSchema>().get();
	}

	data(const saw::data<typename Schema::InterfaceSchema, encode::NativeRaw>& val){
		value_ = val.template cast_to<typename Schema::StorageSchema>();
	}
	
	void set(typename saw::native_data_type<typename Schema::InterfaceSchema>::type val){
		value_.set(static_cast<typename Schema::StorageSchema>(val));
	}
	
	data<Schema, encode::NativeRaw> operator*(const data<Schema, encode::NativeRaw>& rhs) const {
		using CalcType = typename native_data_type<typename Schema::InterfaceSchema>::type;
		CalcType left = static_cast<CalcType>(value_.get());
		CalcType right = static_cast<CalcType>(rhs.get());
		return {left * right};
	}
	
	data<typename Schema::InterfaceSchema, encode::NativeRaw> operator*(const data<typename Schema::InterfaceSchema, encode::NativeRaw>& rhs) const {
		using CalcType = typename native_data_type<typename Schema::InterfaceSchema>::type;
		CalcType left = static_cast<CalcType>(value_.get());
		CalcType right = rhs.get();
		return {left * right};
	}
	
	data<typename Schema::InterfaceSchema, encode::NativeRaw> operator/(const data<Schema, encode::NativeRaw>& rhs)const{
		using CalcType = typename native_data_type<typename Schema::InterfaceSchema>::type;
		CalcType left = static_cast<CalcType>(value_.get());
		CalcType right = static_cast<CalcType>(rhs.get());
		return {left / right};
	}
	
	data<typename Schema::InterfaceSchema, encode::NativeRaw> operator+(const data<Schema, encode::NativeRaw>& rhs)const{
		using CalcType = typename native_data_type<typename Schema::InterfaceSchema>::type;
		CalcType left = static_cast<CalcType>(value_.get());
		CalcType right = static_cast<CalcType>(rhs.get());
		return {left + right};
	}
	
	data<Schema, encode::NativeRaw>& operator+=(const data<Schema, encode::NativeRaw>& rhs)const{
		*this = *this + rhs.get();
		return *this;
	}

	data<typename Schema::InterfaceSchema, encode::NativeRaw> operator-(const data<Schema, encode::NativeRaw>& rhs)const{
		using CalcType = typename native_data_type<typename Schema::InterfaceSchema>::type;
		CalcType left = static_cast<CalcType>(value_.get());
		CalcType right = static_cast<CalcType>(rhs.get());
		return {left - right};
	}
	
	data<Schema, encode::NativeRaw>& operator-=(const data<Schema, encode::NativeRaw>& rhs) const {
		*this = *this - rhs.get();
		return *this;
	}

	template<typename Enc>
	bool operator==(const data<Schema, Enc>& rhs)const{
		return get() == rhs.get();
	}

	template<typename Enc>
	bool operator<(const data<Schema, Enc>& rhs) const {
		return get() < rhs.get();
	}
};

/**
 * Data class which represents a struct in the native format.
 */
template<typename... T, string_literal... literals>
class data<schema::Struct<schema::Member<T, literals>...>, encode::NativeRaw> {
public:
	using Schema = schema::Struct<schema::Member<T,literals>...>;
	using MetaSchema = typename meta_schema<Schema>::MetaSchema;
private:
	/**
	 * Tuple storing the member values.
	 */
	std::tuple<data<T,encode::NativeRaw>...> value_;
	
	template<uint64_t i = 0u>
	struct init_helper {
		static void apply(std::tuple<data<T,encode::NativeRaw>...>& val, data<Schema, encode::NativeRaw>&& init){
			auto& val_i = val.template get<i>();
			auto& init_i = init.template get<i>();

			val_i = {std::move(init_i)};

			if constexpr ( (i+1u) < sizeof...(T) ){
				init_helper<i+1u>::apply(val, init);
			}
		}
	};
public:
	/**
	 * Default constructor.
	 */
	data() = default;
	data(data<MetaSchema, encode::NativeRaw> init){
		if constexpr ( 0u < sizeof...(T) ){
			init_helper<0u>::apply(value_, std::move(init));
		}
	}

	SAW_DEFAULT_COPY(data);
	SAW_DEFAULT_MOVE(data);
	
	/**
	 * Get the member value based on the string_literal.
	 */
	template<string_literal literal>
	data<
		typename parameter_pack_type<
			parameter_key_pack_index<
				literal, literals...
			>::value
		, T...>::type
	, encode::NativeRaw>& get(){
		return std::get<parameter_key_pack_index<literal, literals...>::value>(value_);
	}
	
	/**
	 * Get the member value based on the string_literal.
	 */
	template<string_literal literal>
	const data<
		typename parameter_pack_type<
			parameter_key_pack_index<
				literal, literals...
			>::value
		, T...>::type
	, encode::NativeRaw>& get() const {
		return std::get<parameter_key_pack_index<literal, literals...>::value>(value_);
	}

	/**
	 * Return the amount of members.
	 */
	constexpr size_t size() const {
		return sizeof...(T);
	}
};

template<typename... T>
class data<schema::Tuple<T...>, encode::NativeRaw> {
public:
	using Schema = schema::Tuple<T...>;
	using MetaSchema = typename meta_schema<Schema>::MetaSchema;
private:
	std::tuple<data<T,encode::NativeRaw>...> value_;

	template<uint64_t i = 0u>
	struct init_helper {
		static void apply(std::tuple<data<T,encode::NativeRaw>...>& val, data<MetaSchema, encode::NativeRaw>&& init){
			auto& val_i = val.template get<i>();
			auto& init_i = init.template get<i>();

			val_i = {std::move(init_i)};

			if constexpr ( (i+1u) < sizeof...(T) ){
				init_helper<i+1u>::apply(val, init);
			}
		}
	};
public:
	data() = default;
	data(data<MetaSchema, encode::NativeRaw> init){
		if constexpr ( 0u < sizeof...(T) ){
			init_helper<0u>::apply(value_, std::move(init));
		}
	}

	SAW_DEFAULT_COPY(data);
	SAW_DEFAULT_MOVE(data);
	
	template<size_t i>
	data<typename parameter_pack_type<i,T...>::type, encode::NativeRaw>& get(){
		return std::get<i>(value_);
	}
	
	template<size_t i>
	const data<typename parameter_pack_type<i,T...>::type, encode::NativeRaw>& get() const{
		return std::get<i>(value_);
	}

	constexpr size_t size() const {
		return sizeof...(T);
	}
};

template<typename T, size_t Dim>
class data<schema::Array<T,Dim>, encode::NativeRaw> {
	public:
		using Schema = schema::Array<T,Dim>;
		using MetaSchema = typename meta_schema<Schema>::MetaSchema;

	private:
		// rawr
		// data<schema::FixedArray<schema::UInt64, Dim>, encode::NativeRaw> dims_;
		std::array<uint64_t, Dim> dims_;
		std::vector<typename raw_native_array_type_helper<T>::Type> value_;

		uint64_t get_full_size() const {
			uint64_t s = 1;

			for(uint64_t iter = 0; iter < Dim; ++iter){
				assert(dims_.at(iter) > 0);
				s *= dims_.at(iter);
			}

			return s;
		}
	public:
		data():
			value_{}
		{
			for(auto& iter : dims_){
				iter = 0u;
			}
		}

		/**
		 * Purely internal function for working C interfacing
		 */
		typename raw_native_array_type_helper<T>::Type* get_raw_data() {
			if(value_.size() == 0u){
				return nullptr;
			}
			return &(value_[0]);
		}

		SAW_DEFAULT_COPY(data);
		SAW_DEFAULT_MOVE(data);

		data(const std::array<uint64_t, Dim>& i):
			dims_{i},
			value_{}
		{
			value_.resize(get_full_size());
		}

		data(data<MetaSchema, encode::NativeRaw> init)
		{
			for(uint64_t i = 0; i < Dim; ++i){
				dims_.at(i) = init.at(i).get();
			}
			value_.resize(get_full_size());
		}

		template<size_t i = 0>
		error_or<void> add(saw::data<T,encode::NativeRaw> data){
			/** @todo
			 * Generally the last dimension can always accept a element so to say.
			 * Changing the others would require moving data due to the stride changing.
			 * Since the last dimension doesn't affect the stride, we don't need reordering there.
			 * But I want a quick solution for one dimension so here we are.
			 *
			 * I can always ignore strides and use a stacked std::vector
			 * std::vector<std::vector<...>> and so on.
			 * But for now I'm keeping the strides. Smaller chunks of memory aren't to bad either
			 * though.
			 * I'll probably change it to the smaller chunks
			 */
			static_assert(Dim == 1, "Currently can't deal with higher dims");
			static_assert(i < Dim, "Can't add to dimension. Index i larger than dimension size");

			try {
				value_.emplace_back(std::move(data));
			}catch(const std::exception& e){
				(void) e;
				return make_error<err::out_of_memory>();
			}

			++dims_.at(i);

			return void_t{};
		}

		template<std::integral... Dims>
		data(Dims... size_):
			data{{static_cast<uint64_t>(size_)...}}
		{
			static_assert(sizeof...(Dims)==Dim, "Argument size must be equal to the Dimension");
		}

		data<schema::Ref<T>, encode::NativeRaw> at(const std::array<uint64_t, Dim>& ind){
			return {value_.at(this->get_flat_index(ind))};
		}
		
		const data<schema::Ref<T>, encode::NativeRaw> at(const std::array<uint64_t, Dim>& ind) const {
			return {value_.at(this->get_flat_index(ind))};
		}

		template<std::integral... Dims>
		data<schema::Ref<T>, encode::NativeRaw> at(Dims... i){
			return {value_.at(this->get_flat_index(std::array<uint64_t, Dim>{static_cast<uint64_t>(i)...}))};
		}
		
		template<std::integral... Dims>
		const data<schema::Ref<T>, encode::NativeRaw> at(Dims... i) const {
			return {value_.at(this->get_flat_index(std::array<uint64_t, Dim>{static_cast<uint64_t>(i)...}))};
		}

		template<typename Encoding>
		data<schema::Ref<T>,encode::NativeRaw> at(const data<schema::FixedArray<schema::UInt64,Dim>, Encoding>& i){
			return {value_.at(this->get_flat_index(i))};
		}
		
		template<typename Encoding>
		const data<schema::Ref<T>,encode::NativeRaw> at(const data<schema::FixedArray<schema::UInt64,Dim>, Encoding>& i)const{
			return {value_.at(this->get_flat_index(i))};
		}
		
		std::size_t get_dim_size(uint64_t i) const {
			return dims_.at(i);
		}

		size_t size() const { return value_.size();}

		data<schema::FixedArray<schema::UInt64, Dim>, encode::NativeRaw> get_dims() const {
			return {dims_};
		}

private:
		template<typename U>
		uint64_t get_flat_index(const U& i) const {
			assert(value_.size() == get_full_size());
			uint64_t s = 0;

			uint64_t stride = 1;

			for(uint64_t iter = 0; iter < Dim; ++iter){
				uint64_t ind = [](auto val) -> uint64_t {
					using V = std::decay_t<decltype(val)>;
					if constexpr (std::is_same_v<V,data<schema::UInt64>>){
						return val.get();
					}else if constexpr (std::is_same_v<V, uint64_t>){
						return val;
					}else{
						static_assert(always_false<V>, "Cases exhausted");
					}
				}(i.at(iter));
				assert(ind < dims_.at(iter));
				s += ind * stride;
				stride *= dims_.at(iter);
			}

			return s;
		}
};

template<typename T, uint64_t... D>
class data<schema::FixedArray<T,D...>, encode::NativeRaw> {
public:
	using Schema = schema::FixedArray<T,D...>;
	using MetaSchema = typename meta_schema<Schema>::MetaSchema;
private:
	//using inner_type = std::array<data<T, encode::Native>, multiply_helper<Dims...>::value>;
	//std::unique_ptr<inner_type> value_;
	using ArrayT = std::array<typename raw_native_array_type_helper<T>::Type, ct_multiply<uint64_t, D...>::value>;
	ArrayT value_;
	
public:
	data() = default;
	data(data<MetaSchema, encode::NativeRaw>){}
	data(const std::array<data<T, encode::NativeRaw>, ct_multiply<uint64_t,D...>::value>& value__):
		value_{value__}
	{}

	data<schema::Ref<T>, encode::NativeRaw> at(const std::array<uint64_t, sizeof...(D)>& ind){
		return {value_.at(this->get_flat_index(ind))};
	}
		
	const data<schema::Ref<T>, encode::Native> at(const std::array<uint64_t, sizeof...(D)>& ind) const {
		return {value_.at(this->get_flat_index(ind))};
	}

	template<std::integral... Dims>
	data<schema::Ref<T>, encode::Native> at(Dims... i) {
		return {value_.at(this->get_flat_index({i...}))};
	}

	template<std::integral... Dims>
	const data<schema::Ref<T>, encode::Native> at(Dims... i) const {
		return {value_.at(this->get_flat_index({i...}))};
	}

	data<schema::Ref<T>, encode::Native> at(const data<schema::FixedArray<schema::UInt64, sizeof...(D)>, encode::NativeRaw>& i){
		return {value_.at(this->get_flat_index(i))};
	}
	
	const data<schema::Ref<T>, encode::Native> at(const data<schema::FixedArray<schema::UInt64, sizeof...(D)>, encode::NativeRaw>& i)const{
		return {value_.at(this->get_flat_index(i))};
	}

	template<uint64_t i>
	uint64_t get_dim_size() const {
		return parameter_pack_value<i, uint64_t, D...>::value;
	}

	data<schema::FixedArray<schema::UInt64, sizeof...(D)>, encode::NativeRaw> get_dims() const {
		return {std::array<uint64_t, sizeof...(D)>{D...}};
	}
private:
	uint64_t get_flat_index(const std::array<uint64_t, sizeof...(D)>& i) const {
		uint64_t s = 0;

		uint64_t stride = 1;

		constexpr static std::array<uint64_t, sizeof...(D)> dims_{D...};

		for(uint64_t iter = 0; iter < sizeof...(D); ++iter){
			assert(i.at(iter) < dims_.at(iter));
			s += i.at(iter) * stride;
			stride *= dims_.at(iter);
		}

		return s;
	}
};


}