#pragma once

#include "EditorAttachment.hpp"
#include "GraphicsTags.hpp"
#include "RcPtr.hpp"

#include <absl/container/flat_hash_map.h>
#include <glad/glad.h>
#include <filesystem>
#include <memory>
#include <string_view>
#include <vector>

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

struct ShaderVariable {
	enum Kind {
		KD_Math,
		KD_Sampler,
		KD_UniformBlock,
	};

	Kind kind;

protected:
	ShaderVariable(Kind kind)
		: kind{ kind } {}
};

struct ShaderMathVariable : public ShaderVariable {
	std::string name;
	GLuint location;
	GLenum scalarType;
	int arrayLength;
	int width;
	int height;

	ShaderMathVariable()
		: ShaderVariable(KD_Math) {}
};

struct ShaderSamplerVariable : public ShaderVariable {
	std::string name;
	GLuint location;
	GLenum type;
	int arrayLength;

	ShaderSamplerVariable()
		: ShaderVariable(KD_Sampler) {}
};

struct ShaderUniformBlockVariable : public ShaderVariable {
	std::string name;
	/// Possible values: KD_Math
	std::vector<std::unique_ptr<ShaderVariable>> items;
	GLuint index;

	ShaderUniformBlockVariable()
		: ShaderVariable(KD_UniformBlock) {}
};

struct ShaderInfo {
	enum ThingKind {
		TKD_Input,
		TKD_Output,
		TKD_Uniform,
		TKD_UniformBlock,
	};

	struct ThingId {
		ThingKind kind;
		int index;
	};

	struct InputOutputThing {
		ShaderMathVariable variable;
		Tags::VertexElementSemantic semantic = Tags::VES_Generic;
	};

	absl::flat_hash_map<std::string, ThingId> things;
	std::vector<InputOutputThing> inputs;
	std::vector<InputOutputThing> outputs;
	/// Possible values: KD_Math, KD_Sampler
	std::vector<std::unique_ptr<ShaderVariable>> uniforms;
	std::vector<ShaderUniformBlockVariable> uniformBlocks;

	bool SaveToFile(const std::filesystem::path& filePath) const;
	bool LoadFromFile(const std::filesystem::path& filePath);
};

class Shader : public RefCounted {
private:
	std::string mName;
	std::unique_ptr<ShaderInfo> mInfo;
	std::unique_ptr<EditorAttachment> mEditorAttachment;
	GLuint mHandle = 0;

public:
	Shader(std::string name = "");
	~Shader();
	Shader(const Shader&) = delete;
	Shader& operator=(const Shader&) = delete;
	Shader(Shader&&) = default;
	Shader& operator=(Shader&&) = default;

	enum ErrorCode {
		Success,
		/// Generated when Init*() functions are called on an already initialized Shader object.
		ShaderAlreadyCreated,
		/// Generated when the one-source-file text contains invalid or duplicate shader variants.
		InvalidShaderVariant,
		FileIOFailed,
		CompilationFailed,
		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);

	EditorAttachment* GetEditorAttachment() const { return mEditorAttachment.get(); }
	void SetEditorAttachment(EditorAttachment* attachment) { mEditorAttachment.reset(attachment); }

	bool CreateEmptyInfoIfAbsent();
	bool GatherInfoIfAbsent();
	ShaderInfo* GetInfo() const;
	/// If not empty, this name must not duplicate with any other shader object in the process.
	const std::string& GetName() const;
	GLuint GetProgram() const;

	bool IsValid() const;
};

class ShaderManager {
public:
	static inline ShaderManager* instance = nullptr;

private:
	absl::flat_hash_map<std::string_view, RcPtr<Shader>> mShaders;

public:
	void DiscoverShaders();

	const auto& GetShaders() const { return mShaders; }
	Shader* FindShader(std::string_view name);
};