#pragma once

#include "GraphicsTags.hpp"
#include "Ires.hpp"
#include "RcPtr.hpp"
#include "Utils.hpp"

#include <glad/glad.h>
#include <robin_hood.h>
#include <json_dto/pub.hpp>
#include <memory>
#include <string_view>
#include <variant>
#include <vector>

// TODO move to variable after pattern matching is in the language

// Forward declarations
class Shader;
class IresShader;

struct ShaderMathVariable {
	std::string name;
	GLuint location;
	Tags::VertexElementSemantic semantic = Tags::VES_Generic;
	GLenum scalarType;
	int width;
	int height;
	int arrayLength = 1;

	void ShowInfo() const;

	template <typename TJsonIo>
	void json_io(TJsonIo& io) {
		io& json_dto::mandatory("Name", name);
		io& json_dto::mandatory("Semantic", static_cast<int>(semantic));
		io& json_dto::mandatory("ScalarType", scalarType);
		io& json_dto::mandatory("Width", width);
		io& json_dto::mandatory("Height", height);
		io& json_dto::optional("ArrayLength", arrayLength, 1);
	}
};

struct ShaderSamplerVariable {
	std::string name;
	GLuint location;
	Tags::VertexElementSemantic semantic = Tags::VES_Generic;
	GLenum samplerType;
	int arrayLength = 1;

	void ShowInfo() const;

	template <typename TJsonIo>
	void json_io(TJsonIo& io) {
		io& json_dto::mandatory("Name", name);
		io& json_dto::mandatory("Semantic", static_cast<int>(semantic));
		io& json_dto::mandatory("SamplerType", samplerType);
		io& json_dto::optional("ArrayLength", arrayLength, 1);
	}
};

struct ShaderThingId {
	enum Kind {
		KD_Input,
		KD_Output,
		KD_Uniform,
	};

	Kind kind;
	int index;
};

struct ShaderInfo {
	robin_hood::unordered_map<std::string, ShaderThingId, StringHash, StringEqual> things;
	std::vector<ShaderMathVariable> inputs;
	std::vector<ShaderMathVariable> outputs;
	std::vector<std::variant<ShaderMathVariable, ShaderSamplerVariable>> uniforms;
	static constexpr int kKindMath = 0;
	static constexpr int kKindSampler = 1;

	// Find the first variable with the matching semantic
	GLuint FindInputLocation(Tags::VertexElementSemantic semantic) const;
	GLuint FindOutputLocation(Tags::VertexElementSemantic semantic) const;

	template <typename TJsonIo>
	void json_io(TJsonIo& io) {
		io& json_dto::mandatory("Inputs", inputs);
		io& json_dto::mandatory("Outputs", outputs);
		// TODO make json_dto support std::variant
//		io& json_dto::mandatory("Uniforms", uniforms);
	}
};

class Shader : public RefCounted {
	friend class IresShader;

private:
	IresShader* mIres = nullptr;
	ShaderInfo mInfo;
	GLuint mProgram = 0;

public:
	GLuint autofill_Transform = Tags::kInvalidLocation;
	GLuint autofill_Time = Tags::kInvalidLocation;
	GLuint autofill_DeltaTime = Tags::kInvalidLocation;
	GLuint autofill_TextureAtlas = Tags::kInvalidLocation;

public:
	Shader();
	~Shader();
	Shader(const Shader&) = delete;
	Shader& operator=(const Shader&) = delete;
	Shader(Shader&&) = default;
	Shader& operator=(Shader&&) = default;

	enum ErrorCode {
		EC_Success,
		/// Generated when Init*() functions are called on an already initialized Shader object.
		EC_AlreadyInitialized,
		/// Generated when the one-source-file text contains invalid or duplicate shader variants.
		EC_InvalidShaderVariant,
		EC_FileIoFailed,
		EC_CompilationFailed,
		EC_LinkingFailed,
	};

	struct ShaderSources {
		std::string_view vertex;
		std::string_view geometry;
		std::string_view tessControl;
		std::string_view tessEval;
		std::string_view fragment;
	};

	/// Create shader by compiling each shader source file separately and then combining them together
	/// into a Shader object.
	ErrorCode InitFromSources(const ShaderSources& sources);

	/// The given text will be split into different shader sections according to #type directives,
	/// and combined to form a Shader object.
	/// For OpenGL, this process involves compililng each section separately and then linking them
	/// together.
	///
	/// There are a directive for each shader type
	/// - `#type vertex`: Vertex shader
	/// - `#type geometry`: Geometry shader
	/// - `#type tessellation_control`: Tessellation control shader
	/// - `#type tessellation_evaluation`: Tessellation evaluation shader
	/// - `#type fragment`: Fragment shader
	ErrorCode InitFromSource(std::string_view source);

	/// Rebuild info object using OpenGL shader introspection API. Requires OpenGL 4.3 or above. Overrides existing info object.
	bool GatherInfoShaderIntrospection();
	const ShaderInfo& GetInfo() const { return mInfo; }
	ShaderInfo& GetInfo() { return mInfo; }
	/// If not empty, this name must not duplicate with any other shader object in the process.
	GLuint GetProgram() const { return mProgram; }

	IresShader* GetIres() const { return mIres; }

	bool IsValid() const;
};

// Initialized in main()
inline RcPtr<Shader> gDefaultShader;

class IresShader : public IresObject {
private:
	RcPtr<Shader> mInstance;
	std::string mSourceFile;

public:
	IresShader();

	Shader* GetInstance() const;
	void InvalidateInstance();

	void ShowEditor(IEditor& editor) override;

	void Write(IresWritingContext& ctx, rapidjson::Value& value, rapidjson::Document& root) const override;
	void Read(IresLoadingContext& ctx, const rapidjson::Value& value) override;
};