#pragma once

#include "BsCorePrerequisites.h"
#include "BsModule.h"
#include "BsBlendState.h"
#include "BsRasterizerState.h"
#include "BsDepthStencilState.h"
#include "BsSamplerState.h"

namespace BansheeEngine
{
	/**
	 * @brief	Handles creation of various render states.
	 */
	class BS_CORE_EXPORT RenderStateManager : public Module <RenderStateManager>
	{
	public:
		/**
		 * @brief	Creates and initializes a new SamplerState.
		 */
		SamplerStatePtr createSamplerState(const SAMPLER_STATE_DESC& desc) const;

		/**
		 * @brief	Creates and initializes a new DepthStencilState.
		 */
		DepthStencilStatePtr createDepthStencilState(const DEPTH_STENCIL_STATE_DESC& desc) const;

		/**
		 * @brief	Creates and initializes a new RasterizerState.
		 */
		RasterizerStatePtr createRasterizerState(const RASTERIZER_STATE_DESC& desc) const;

		/**
		 * @brief	Creates and initializes a new BlendState.
		 */
		BlendStatePtr createBlendState(const BLEND_STATE_DESC& desc) const;

		/**
		 * @brief	Creates an uninitialized sampler state. Requires manual initialization
		 *			after creation.
		 *
		 * @note	Internal method.
		 */
		SamplerStatePtr _createSamplerStatePtr(const SAMPLER_STATE_DESC& desc) const;

		/**
		 * @brief	Creates an uninitialized depth-stencil state. Requires manual initialization
		 *			after creation.
		 *
		 * @note	Internal method.
		 */
		DepthStencilStatePtr _createDepthStencilStatePtr(const DEPTH_STENCIL_STATE_DESC& desc) const;

		/**
		 * @brief	Creates an uninitialized rasterizer state. Requires manual initialization
		 *			after creation.
		 *
		 * @note	Internal method.
		 */
		RasterizerStatePtr _createRasterizerStatePtr(const RASTERIZER_STATE_DESC& desc) const;

		/**
		 * @brief	Creates an uninitialized blend state. Requires manual initialization
		 *			after creation.
		 *
		 * @note	Internal method.
		 */
		BlendStatePtr _createBlendStatePtr(const BLEND_STATE_DESC& desc) const;

		/**
		 * @brief	Gets a sampler state initialized with default options.
		 */
		const SamplerStatePtr& getDefaultSamplerState() const;

		/**
		 * @brief	Gets a blend state initialized with default options.
		 */
		const BlendStatePtr& getDefaultBlendState() const;

		/**
		 * @brief	Gets a rasterizer state initialized with default options.
		 */
		const RasterizerStatePtr& getDefaultRasterizerState() const;

		/**
		 * @brief	Gets a depth stencil state initialized with default options.
		 */
		const DepthStencilStatePtr& getDefaultDepthStencilState() const;

	private:
		friend class SamplerState;
		friend class BlendState;
		friend class RasterizerState;
		friend class DepthStencilState;

		mutable SamplerStatePtr mDefaultSamplerState;
		mutable BlendStatePtr mDefaultBlendState;
		mutable RasterizerStatePtr mDefaultRasterizerState;
		mutable DepthStencilStatePtr mDefaultDepthStencilState;
	};

	/**
	 * @brief	Handles creation of various render states.
	 */
	class BS_CORE_EXPORT RenderStateCoreManager : public Module<RenderStateCoreManager>
	{
	private:
		/**
		 * @brief	Contains data about a cached blend state
		 */
		struct CachedBlendState
		{
			CachedBlendState()
				:id(0)
			{ }

			CachedBlendState(UINT32 id)
				:id(id)
			{ }

			std::weak_ptr<BlendStateCore> state;
			UINT32 id;
		};

		/**
		 * @brief	Contains data about a cached blend state
		 */
		struct CachedRasterizerState
		{
			CachedRasterizerState()
				:id(0)
			{ }

			CachedRasterizerState(UINT32 id)
				:id(id)
			{ }

			std::weak_ptr<RasterizerStateCore> state;
			UINT32 id;
		};

		/**
		 * @brief	Contains data about a cached blend state
		 */
		struct CachedDepthStencilState
		{
			CachedDepthStencilState()
				:id(0)
			{ }

			CachedDepthStencilState(UINT32 id)
				:id(id)
			{ }

			std::weak_ptr<DepthStencilStateCore> state;
			UINT32 id;
		};

	public:
		RenderStateCoreManager();

		/**
		 * @copydoc	RenderStateManager::createSamplerState
		 */
		SPtr<SamplerStateCore> createSamplerState(const SAMPLER_STATE_DESC& desc) const;

		/**
		 * @copydoc	RenderStateManager::createDepthStencilState
		 */
		SPtr<DepthStencilStateCore> createDepthStencilState(const DEPTH_STENCIL_STATE_DESC& desc) const;

		/**
		 * @copydoc	RenderStateManager::createRasterizerState
		 */
		SPtr<RasterizerStateCore> createRasterizerState(const RASTERIZER_STATE_DESC& desc) const;

		/**
		 * @copydoc	RenderStateManager::createBlendState
		 */
		SPtr<BlendStateCore> createBlendState(const BLEND_STATE_DESC& desc) const;

		/**
		 * @brief	Creates an uninitialized sampler state. Requires manual initialization
		 *			after creation.
		 *
		 * @note	Internal method.
		 */
		SPtr<SamplerStateCore> _createSamplerState(const SAMPLER_STATE_DESC& desc) const;

		/**
		 * @brief	Creates an uninitialized depth-stencil state. Requires manual initialization
		 *			after creation.
		 *
		 * @note	Internal method.
		 */
		SPtr<DepthStencilStateCore> _createDepthStencilState(const DEPTH_STENCIL_STATE_DESC& desc) const;

		/**
		 * @brief	Creates an uninitialized rasterizer state. Requires manual initialization
		 *			after creation.
		 *
		 * @note	Internal method.
		 */
		SPtr<RasterizerStateCore> _createRasterizerState(const RASTERIZER_STATE_DESC& desc) const;

		/**
		 * @brief	Creates an uninitialized blend state. Requires manual initialization
		 *			after creation.
		 *
		 * @note	Internal method.
		 */
		SPtr<BlendStateCore> _createBlendState(const BLEND_STATE_DESC& desc) const;

		/**
		 * @brief	Gets a sampler state initialized with default options.
		 */
		const SPtr<SamplerStateCore>& getDefaultSamplerState() const;

		/**
		 * @brief	Gets a blend state initialized with default options.
		 */
		const SPtr<BlendStateCore>& getDefaultBlendState() const;

		/**
		 * @brief	Gets a rasterizer state initialized with default options.
		 */
		const SPtr<RasterizerStateCore>& getDefaultRasterizerState() const;

		/**
		 * @brief	Gets a depth stencil state initialized with default options.
		 */
		const SPtr<DepthStencilStateCore>& getDefaultDepthStencilState() const;

	protected:
		friend class SamplerState;
		friend class BlendState;
		friend class RasterizerState;
		friend class DepthStencilState;
		friend class SamplerStateCore;
		friend class BlendStateCore;
		friend class RasterizerStateCore;
		friend class DepthStencilStateCore;

		/**
		 * @copydoc	Module::onShutDown
		 */
		void onShutDown() override;

		/**
		 * @copydoc	createSamplerState
		 */
		virtual SPtr<SamplerStateCore> createSamplerStateInternal(const SAMPLER_STATE_DESC& desc) const;

		/**
		 * @copydoc	createBlendState
		 */
		virtual SPtr<BlendStateCore> createBlendStateInternal(const BLEND_STATE_DESC& desc, UINT32 id) const;

		/**
		 * @copydoc	createRasterizerState
		 */
		virtual SPtr<RasterizerStateCore> createRasterizerStateInternal(const RASTERIZER_STATE_DESC& desc, UINT32 id) const;

		/**
		 * @copydoc	createDepthStencilState
		 */
		virtual SPtr<DepthStencilStateCore> createDepthStencilStateInternal(const DEPTH_STENCIL_STATE_DESC& desc, UINT32 id) const;

	private:
		/**
		 * @brief	Triggered when a new sampler state is created. 
		 */
		void notifySamplerStateCreated(const SAMPLER_STATE_DESC& desc, const SPtr<SamplerStateCore>& state) const;

		/**
		 * @brief	Triggered when a new sampler state is created. 
		 */
		void notifyBlendStateCreated(const BLEND_STATE_DESC& desc, const CachedBlendState& state) const;

		/**
		 * @brief	Triggered when a new sampler state is created. 
		 */
		void notifyRasterizerStateCreated(const RASTERIZER_STATE_DESC& desc, const CachedRasterizerState& state) const;

		/**
		 * @brief	Triggered when a new sampler state is created. 
		 */
		void notifyDepthStencilStateCreated(const DEPTH_STENCIL_STATE_DESC& desc, const CachedDepthStencilState& state) const;

		/**
		 * @brief	Triggered when the last reference to a specific sampler state is destroyed, which
		 *			means we must clear our cached version as well.
		 */
		void notifySamplerStateDestroyed(const SAMPLER_STATE_DESC& desc) const;

		/**
		 * @brief	Attempts to find a cached sampler state corresponding to the provided descriptor. 
		 *			Returns null if one doesn't exist.
		 */
		SPtr<SamplerStateCore> findCachedState(const SAMPLER_STATE_DESC& desc) const;

				/**
		 * @brief	Attempts to find a cached blend state corresponding to the provided descriptor. 
		 *			Returns null if one doesn't exist.
		 */
		SPtr<BlendStateCore> findCachedState(const BLEND_STATE_DESC& desc, UINT32& id) const;

				/**
		 * @brief	Attempts to find a cached rasterizer state corresponding to the provided descriptor. 
		 *			Returns null if one doesn't exist.
		 */
		SPtr<RasterizerStateCore> findCachedState(const RASTERIZER_STATE_DESC& desc, UINT32& id) const;

		/**
		 * @brief	Attempts to find a cached depth-stencil state corresponding to the provided descriptor. 
		 *			Returns null if one doesn't exist.
		 */
		SPtr<DepthStencilStateCore> findCachedState(const DEPTH_STENCIL_STATE_DESC& desc, UINT32& id) const;

		mutable SPtr<SamplerStateCore> mDefaultSamplerState;
		mutable SPtr<BlendStateCore> mDefaultBlendState;
		mutable SPtr<RasterizerStateCore> mDefaultRasterizerState;
		mutable SPtr<DepthStencilStateCore> mDefaultDepthStencilState;

		mutable UnorderedMap<SAMPLER_STATE_DESC, std::weak_ptr<SamplerStateCore>> mCachedSamplerStates;
		mutable UnorderedMap<BLEND_STATE_DESC, CachedBlendState> mCachedBlendStates;
		mutable UnorderedMap<RASTERIZER_STATE_DESC, CachedRasterizerState> mCachedRasterizerStates;
		mutable UnorderedMap<DEPTH_STENCIL_STATE_DESC, CachedDepthStencilState> mCachedDepthStencilStates;

		mutable UINT32 mNextBlendStateId;
		mutable UINT32 mNextRasterizerStateId;
		mutable UINT32 mNextDepthStencilStateId;

		BS_MUTEX(mMutex);
	};
}