#pragma once

#include <memory>

namespace ast {

	class Node {
	public:
		virtual ~Node() = default;

		enum class ref_type {
			strong,
			weak
		};

		void increment(ref_type ref) {
			switch (ref) 			{
				case ref_type::strong: strong_ref++; break;
				case ref_type::weak  : weak_ref  ++; break;
			}
		}

		void decrement(ref_type ref) {
			switch (ref) 			{
				case ref_type::strong: strong_ref--; break;
				case ref_type::weak  : weak_ref  --; break;
			}

			if(!strong_ref && !weak_ref) {
				delete this;
			}
		}

	private:
		size_t strong_ref = 0;
		size_t weak_ref = 0;
	};

	template< typename node_t, enum  Node::ref_type ref_t>
	class ptr_base {
	public:
		ptr_base() : node(nullptr) {}
		ptr_base( node_t * n ) : node(n) { if( !node ) node->increment(ref_t); }
		~ptr_base() { if( node ) node->decrement(ref_t); }

		template< enum  Node::ref_type o_ref_t >
		ptr_base( const ptr_base<node_t, o_ref_t> & o ) : node(o.node) {
			if( !node ) return;
			node->increment(ref_t);
		}

		template< enum  Node::ref_type o_ref_t >
		ptr_base( ptr_base<node_t, o_ref_t> && o ) : node(o.node) {
			if( node ) node->increment(ref_t);
			if( node ) node->decrement(o_ref_t);
		}

		template< enum  Node::ref_type o_ref_t >
		ptr_base & operator=( const ptr_base<node_t, o_ref_t> & o ) {
			assign(o.node);
			return *this;
		}

		template< enum  Node::ref_type o_ref_t >
		ptr_base & operator=( ptr_base<node_t, o_ref_t> && o ) {
			if(o.node == node) return *this;
			assign(o.node);
			if( node ) node->decrement(o_ref_t);
			return *this;
		}

		const node_t * get() const { return  node; }
		const node_t * operator->() const { return  node; }
		const node_t & operator* () const { return *node; }
		operator bool() const { return node; }

	private:
		void assign(node_t * other ) {
			if( other ) other->increment(ref_t);
			if( node  ) node ->decrement(ref_t);
			node = other;
		}

	protected:
		node_t * node;
	};

	template< typename node_t >
	class ptr : public ptr_base< node_t, Node::ref_type::strong > {
	public:
		typedef ptr_base< node_t, Node::ref_type::strong > base_t;

		ptr() = default;
		ptr( node_t node ) : base_t( node ) {}
		~ptr() = default;

		template< enum  Node::ref_type ref_t >
		ptr( const ptr_base<node_t, ref_t> & o ) : base_t(o) {}

		template< enum  Node::ref_type ref_t >
		ptr( ptr_base<node_t, ref_t> && o ) : base_t( std::move(o) ) {}

		template< enum  Node::ref_type o_ref_t >
		ptr & operator=( const ptr_base<node_t, o_ref_t> & o ) {
			base_t::operator=(o);
			return *this;
		}

		template< enum  Node::ref_type o_ref_t >
		ptr & operator=( ptr_base<node_t, o_ref_t> && o ) {
			base_t::operator=(std::move(o));
			return *this;
		}

		node_t * mutate() {
			using base_t::node;
			assert(node->strong_count);
			if (node->strong_count == 1) {
				return node;
			}

			assertf(node->weak_count == 0, "Error: mutating node with weak references to it will invalided some references");
			auto n = new node_t(*node);
			assign(n);
			return node;
		}
	};

	template< typename node_t >
	class readonly : public ptr_base< node_t, Node::ref_type::weak > {
	public:
		typedef ptr_base< node_t, Node::ref_type::strong > base_t;

		readonly() = default;
		readonly( node_t node ) : base_t( node ) {}
		~readonly() = default;

		template< enum  Node::ref_type ref_t >
		readonly( const ptr_base<node_t, ref_t> & o ) : base_t(o) {}

		template< enum  Node::ref_type ref_t >
		readonly( ptr_base<node_t, ref_t> && o ) : base_t( std::move(o) ) {}

		template< enum  Node::ref_type o_ref_t >
		readonly & operator=( const ptr_base<node_t, o_ref_t> & o ) {
			base_t::operator=(o);
			return *this;
		}

		template< enum  Node::ref_type o_ref_t >
		readonly & operator=( ptr_base<node_t, o_ref_t> && o ) {
			base_t::operator=(std::move(o));
			return *this;
		}
	};
}