#include "UgcJobs.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <sstream>

#include "json.hpp"

#include "NifFile.h"
#include "Sd0.h"
#include "UgcFormats.h"
#include "UgcModel.h"
#include "UgcModular.h"
#include "ZCompression.h"

namespace UgcJobs {
	void AddDownload(UgcStorage::Files& files, const std::string& name, const std::string& data) {
		files[name + ".gz"] = ZCompression::Gzip(data);
		files[name + ".checksum"] = UgcFormats::ChecksumXml(data);
		// Without 3D services the client downloads it as sd0 (after asking its world for the checksum)
		auto sd0 = Sd0::Compress(data);
		if (!sd0.empty()) files[name + ".sd0"] = std::move(sd0);
	}

	std::string LxfmlFromBlob(const std::string& blob) {
		if (blob.starts_with("<?xml") || blob.starts_with("<LXFML")) return blob;
		std::stringstream stream(blob);
		try {
			Sd0 sd0(stream);
			return sd0.GetAsStringUncompressed();
		} catch (...) {
			return {};
		}
	}

	namespace {
		// The icon files of a model, and false when nothing was drawn
		bool AddIcon(UgcStorage::Files& files, const UgcModel::Model& model, const UgcRender::IconOptions& options, const std::vector<float>* ao = nullptr,
			const UgcModel::Model* plain = nullptr) {
			const auto icon = UgcRender::RenderIcon(model, options, ao, plain);
			bool drawn = false;
			for (size_t i = 3; i < icon.rgba.size(); i += 4) drawn = drawn || icon.rgba[i] != 0;
			files["icon.png"] = UgcFormats::EncodePng(icon);
			AddDownload(files, "icon.dds", UgcFormats::EncodeDds(icon));
			return drawn;
		}
	}

	size_t CountParts(std::string_view lxfml) {
		size_t count = 0;
		for (size_t at = lxfml.find("<Part"); at != std::string_view::npos; at = lxfml.find("<Part", at + 5)) count++;
		return count;
	}

	uint64_t EstimateMemory(size_t parts, const Settings& settings) {
		// Measured: the renders' buffers, and per brick its mesh in each LOD (positions, normals, colors, indices,
		// the occlusion tree and copies made along the way), about 40 KB at LOD 0
		const uint64_t resolution = static_cast<uint64_t>(std::clamp(settings.hsr.resolution, 64, 4096));
		const uint64_t icon = static_cast<uint64_t>(settings.icon.size) * settings.icon.supersample;
		const uint64_t fixed = resolution * resolution * 8 + icon * icon * 20 + 1024 * 1024 * 4 + 16 * 1024 * 1024;
		return fixed + static_cast<uint64_t>(parts) * 40 * 1024 * (1 + settings.lods.size());
	}

	namespace {
		double Since(std::chrono::steady_clock::time_point start) {
			return std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - start).count();
		}

		// An opaque mesh's pieces by look ([eLook]: UgcModel::Divide's pieces), the looks without a shader of their own
		// with the plastic ones
		using LookPieces = std::array<std::vector<UgcModel::Mesh>, UgcModel::LOOK_COUNT>;
		LookPieces DivideByLook(const UgcModel::Mesh& mesh, const std::array<bool, UgcModel::LOOK_COUNT>& separate) {
			LookPieces pieces;
			const auto split = UgcModel::SplitLooks(mesh, separate);
			if (!split) {
				pieces[0] = UgcModel::Divide(mesh);
				return pieces;
			}
			for (size_t look = 0; look < UgcModel::LOOK_COUNT; look++) pieces[look] = UgcModel::Divide((*split)[look]);
			return pieces;
		}

		// A model's transparent bricks as the .nif's shapes: one per brick (or all together, `combine`), the glitter
		// ones apart when `glitterApart`
		std::array<std::vector<UgcModel::Mesh>, 2> DivideTransparent(const UgcModel::Model& model, bool combine, bool glitterApart) {
			std::array<std::vector<UgcModel::Mesh>, 2> out;
			std::array<bool, UgcModel::LOOK_COUNT> separate{};
			separate[static_cast<size_t>(UgcModel::eLook::GLITTER)] = glitterApart;
			const auto glitter = static_cast<size_t>(UgcModel::eLook::GLITTER);
			if (combine) {
				const auto split = UgcModel::SplitLooks(model.transparent, separate);
				if (!split) {
					out[0] = UgcModel::Divide(model.transparent);
				} else {
					out[0] = UgcModel::Divide((*split)[0]);
					out[1] = UgcModel::Divide((*split)[glitter]);
				}
				return out;
			}
			for (auto& piece : UgcModel::SplitAt(model.transparent, model.transparentBricks)) {
				const auto split = UgcModel::SplitLooks(piece, separate);
				if (!split) {
					out[0].push_back(std::move(piece));
					continue;
				}
				if (!(*split)[0].Empty()) out[0].push_back(std::move((*split)[0]));
				if (!(*split)[glitter].Empty()) out[1].push_back(std::move((*split)[glitter]));
			}
			return out;
		}
	}

	uint32_t Shaders::TagOf(UgcModel::eLook look) const {
		switch (look) {
			case UgcModel::eLook::METAL: return metal;
			case UgcModel::eLook::BRUSHED: return brushed;
			case UgcModel::eLook::GLOW: return glow;
			case UgcModel::eLook::GLITTER: return glitter;
			default: return 0;
		}
	}

	std::map<int32_t, UgcModel::eLook> Shaders::TagLooks() const {
		std::map<int32_t, UgcModel::eLook> looks{ { 88, UgcModel::eLook::METAL }, { 89, UgcModel::eLook::BRUSHED }, { 46, UgcModel::eLook::GLOW },
			{ 21, UgcModel::eLook::GLITTER } };
		for (const auto look : { UgcModel::eLook::METAL, UgcModel::eLook::BRUSHED, UgcModel::eLook::GLOW, UgcModel::eLook::GLITTER }) {
			if (const auto tag = TagOf(look); tag != 0) looks[static_cast<int32_t>(tag)] = look;
		}
		return looks;
	}

	std::set<int32_t> Shaders::OverlayTags() const {
		std::set<int32_t> tags{ 79 };
		if (sparkle != 0) tags.insert(static_cast<int32_t>(sparkle));
		return tags;
	}

	std::string SparkleName(const Settings& settings) {
		const auto tag = std::to_string(settings.shaders.sparkle);
		return ("S" + std::string(tag.size() < 2 ? "0" : "") + tag + "_GlitterSparkle_Model").substr(0, 60);
	}

	std::string ShapeName(const Settings& settings, UgcModel::eLook look, bool transparent) {
		if (transparent && look != UgcModel::eLook::GLITTER) return "S01_Alpha_Model";
		if (look == UgcModel::eLook::PLASTIC) return ("S" + settings.shaderOpaque + "_Opaque_Model").substr(0, 60);
		const auto tag = std::to_string(settings.shaders.TagOf(look));
		const char* name = look == UgcModel::eLook::METAL ? "_Metal_Model" : look == UgcModel::eLook::BRUSHED ? "_Brushed_Model" :
			look == UgcModel::eLook::GLOW ? "_Glow_Model" : transparent ? "_GlitterAlpha_Model" : "_Glitter_Model";
		return ("S" + std::string(tag.size() < 2 ? "0" : "") + tag + name).substr(0, 60);
	}

	bool IconFromNif(const std::string& nif, const UgcRender::IconOptions& options, UgcStorage::Files& files, std::string& error,
		const std::map<int32_t, UgcModel::eLook>& tagLooks, const std::set<int32_t>& overlayTags, const std::string* plainNif) {
		const auto readBack = NifFile::Parse(nif, 0, error);
		if (!readBack) return false;
		// The denoiser's guide: the model before its occlusion was baked in (model.noao.nif), when there is one
		std::optional<UgcModel::Model> plain;
		if (options.denoise != UgcRender::eDenoise::OFF && plainNif) {
			std::string plainError;
			if (const auto plainRead = NifFile::Parse(*plainNif, 0, plainError)) plain = UgcModel::FromNif(*plainRead, tagLooks, overlayTags);
		}
		AddIcon(files, UgcModel::FromNif(*readBack, tagLooks, overlayTags), options, nullptr, plain ? &*plain : nullptr);
		return true;
	}

	std::optional<std::string> WithIconTime(const std::string& stats, double iconMs, double& change) {
		auto parsed = nlohmann::json::parse(stats, nullptr, false);
		if (!parsed.is_object()) return std::nullopt;
		auto& ms = parsed["ms"];
		if (!ms.is_object()) ms = nlohmann::json::object();
		const double before = ms.value("icon", 0.0);
		change = iconMs - before;
		ms["icon"] = std::lround(iconMs);
		ms["total"] = std::max(0L, std::lround(ms.value("total", 0.0) + change));
		return parsed.dump();
	}

	void ApplyOptions(Settings& settings, const UgcProcessOptions::Choice& choice) {
		if (const auto rays = UgcRays::Parse(choice.rays)) settings.ao.rays = settings.icon.ao.rays = *rays;
		if (const auto denoise = UgcRender::ParseDenoise(choice.denoise)) settings.icon.denoise = *denoise;
	}

	UgcProcessOptions::Choice MadeWith(const Settings& settings) {
		const auto denoise = UgcRender::Available(settings.icon.denoise) ? settings.icon.denoise : UgcRender::eDenoise::OFF;
		return { std::string(UgcRays::Name(UgcRays::Resolve(settings.ao.rays))), std::string(UgcRender::Name(denoise)) };
	}

	Outcome ProcessModel(const std::string& blob, UgcBricks::BrickLibrary& library, const Settings& settings, uint64_t seed, const UgcIconParams::Values& iconValues) {
		Outcome outcome;
		const auto started = std::chrono::steady_clock::now();
		const auto lxfml = LxfmlFromBlob(blob);
		if (lxfml.empty()) {
			outcome.error = "the stored LXFML can't be read";
			return outcome;
		}
		std::string error;
		const auto parts = UgcModel::ParseLxfml(lxfml, error);
		if (parts.empty()) {
			outcome.error = error;
			// A model with no bricks has nothing to make; the LXFML itself is still served
			outcome.empty = UgcModel::HasNoBricks(lxfml);
			return outcome;
		}
		if (settings.maxBricks > 0 && parts.size() > settings.maxBricks) {
			outcome.error = "the model has " + std::to_string(parts.size()) + " bricks, more than max_model_bricks (" + std::to_string(settings.maxBricks) + ")";
			return outcome;
		}

		auto lods = settings.lods;
		std::erase_if(lods, [](uint32_t lod) { return lod > 3; });
		std::sort(lods.begin(), lods.end());
		lods.erase(std::unique(lods.begin(), lods.end()), lods.end());
		if (lods.empty()) lods.push_back(0);
		const auto ranges = UgcModel::LodRanges(lods, settings.lodDistances);

		nlohmann::json stats;
		stats["version"] = 1;
		stats["bricks"] = parts.size();
		auto& lodStats = stats["lods"] = nlohmann::json::array();
		double buildMs = 0, hsrMs = 0, aoMs = 0;

		// Every LOD made like LU Toolbox makes each LOD collection: colored, hidden faces removed, lighting baked, divided
		std::vector<UgcModel::Model> models;
		std::vector<LookPieces> opaquePieces;
		// Per level: the transparent bricks' pieces, plastic and (with the glitter group on) glitter
		using TransparentPieces = std::array<std::vector<UgcModel::Mesh>, 2>;
		std::vector<TransparentPieces> transparentPieces;
		// The looks with a shader of their own (UgcJobs::Shaders), each an NiLODNode apart from the plastic
		std::array<bool, UgcModel::LOOK_COUNT> separate{};
		for (size_t look = 1; look < UgcModel::LOOK_COUNT; look++) separate[look] = settings.shaders.TagOf(static_cast<UgcModel::eLook>(look)) != 0;
		const bool glowApart = separate[static_cast<size_t>(UgcModel::eLook::GLOW)];
		// The glitter bricks' sparkles, a group over both glitter groups
		const bool sparkles = separate[static_cast<size_t>(UgcModel::eLook::GLITTER)] && settings.shaders.sparkle != 0;
		UgcModel::Model preview; // LOD 0 before the lighting bake, for the dashboard
		for (size_t i = 0; i < lods.size(); i++) {
			auto options = settings.build;
			options.seed = seed;
			options.lod = lods[i];
			auto step = std::chrono::steady_clock::now();
			auto model = UgcModel::Build(parts, library, options);
			buildMs += Since(step);
			if (i == 0) {
				if (!model.missingDesigns.empty()) {
					outcome.note = "no geometry for design(s)";
					for (const auto design : model.missingDesigns) outcome.note += " " + std::to_string(design);
					stats["missingDesigns"] = model.missingDesigns;
				}
				if (model.Empty()) {
					outcome.error = "none of the model's bricks have geometry";
					if (!outcome.note.empty()) outcome.error += " (" + outcome.note + ")";
					return outcome;
				}
			}
			nlohmann::json entry{ { "lod", lods[i] }, { "near", ranges[i].first }, { "far", ranges[i].second },
				{ "opaqueBefore", model.opaque.TriangleCount() }, { "transparent", model.transparent.TriangleCount() } };
			step = std::chrono::steady_clock::now();
			const auto optimized = UgcHsr::RemoveHiddenFaces(model, settings.hsr);
			hsrMs += Since(step);
			if (i == 0 && optimized.trianglesRemoved > 0) {
				if (!outcome.note.empty()) outcome.note += "; ";
				outcome.note += "removed " + std::to_string(optimized.trianglesRemoved) + " of " + std::to_string(optimized.trianglesBefore) + " triangles";
			}
			if (i == 0) preview = model;
			// The emissive shader lerps from its own lighting to the vertex color, so glowing bricks keep their plain
			// color: no occlusion, and no glow added (it would glow twice)
			std::vector<glm::vec4> plainColors;
			if (glowApart && !model.opaque.looks.empty()) plainColors = model.opaque.colors;
			step = std::chrono::steady_clock::now();
			UgcRender::BakeAo(model, settings.ao);
			aoMs += Since(step);
			for (size_t v = 0; v < plainColors.size() && v < model.opaque.looks.size(); v++) {
				if (model.opaque.looks[v] == UgcModel::eLook::GLOW) model.opaque.colors[v] = glm::vec4(glm::vec3(plainColors[v]), 1.0f);
			}
			entry["opaqueAfter"] = model.opaque.TriangleCount();
			entry["vertices"] = model.opaque.positions.size() + model.transparent.positions.size();
			opaquePieces.push_back(DivideByLook(model.opaque, separate));
			transparentPieces.push_back(DivideTransparent(model, settings.combineTransparent, separate[static_cast<size_t>(UgcModel::eLook::GLITTER)]));
			size_t shapes = transparentPieces.back()[0].size() + transparentPieces.back()[1].size();
			for (const auto& pieces : opaquePieces.back()) shapes += pieces.size();
			entry["shapes"] = shapes;
			// Triangles per group (NiLODNode) when metal, glow or glitter have groups of their own
			if (std::find(separate.begin(), separate.end(), true) != separate.end()) {
				auto& byGroup = entry["groups"] = nlohmann::json::object();
				for (size_t look = 0; look < UgcModel::LOOK_COUNT; look++) {
					size_t triangles = 0;
					for (const auto& piece : opaquePieces.back()[look]) triangles += piece.TriangleCount();
					if (triangles > 0) byGroup[ShapeName(settings, static_cast<UgcModel::eLook>(look), false)] = triangles;
				}
				for (size_t kind = 0; kind < 2; kind++) {
					size_t triangles = 0;
					for (const auto& piece : transparentPieces.back()[kind]) triangles += piece.TriangleCount();
					if (triangles > 0) byGroup[ShapeName(settings, kind ? UgcModel::eLook::GLITTER : UgcModel::eLook::PLASTIC, true)] = triangles;
				}
				if (sparkles) {
					size_t triangles = 0;
					for (const auto& piece : opaquePieces.back()[static_cast<size_t>(UgcModel::eLook::GLITTER)]) triangles += piece.TriangleCount();
					for (const auto& piece : transparentPieces.back()[1]) triangles += piece.TriangleCount();
					if (triangles > 0) byGroup[SparkleName(settings)] = triangles;
				}
			}
			lodStats.push_back(entry);
		}
		outcome.aoBaked = settings.ao.enabled;

		// An NiLODNode for the opaque bricks and one for the transparent ones, as LU Toolbox names them, and one for
		// each look with a shader of its own between them. Every group has every level (empty where it has nothing).
		// Transparent glitter last, after the plain transparent bricks.
		const auto groups = [&](size_t levels, const std::vector<LookPieces>& opaque, const std::vector<TransparentPieces>& transparent) {
			std::vector<UgcFormats::NifLodGroup> out;
			for (size_t kind = 0; kind < UgcModel::LOOK_COUNT + 2; kind++) {
				const bool isTransparent = kind >= UgcModel::LOOK_COUNT;
				const bool glitter = kind == UgcModel::LOOK_COUNT + 1 || kind == static_cast<size_t>(UgcModel::eLook::GLITTER);
				const auto look = glitter ? UgcModel::eLook::GLITTER : isTransparent ? UgcModel::eLook::PLASTIC : static_cast<UgcModel::eLook>(kind);
				UgcFormats::NifLodGroup group{ ShapeName(settings, look, isTransparent), isTransparent, {} };
				if (look == UgcModel::eLook::GLOW) group.emissive = std::max(settings.shaders.glowEmissive, 0.0f);
				if (glitter) group.glitter = &settings.shaders.glitterParams;
				bool any = false;
				for (size_t i = 0; i < levels; i++) {
					UgcFormats::NifLod lod{ ranges[i].first, ranges[i].second, "LOD_" + std::to_string(lods[i]), {} };
					for (const auto& piece : (isTransparent ? transparent[i][glitter ? 1 : 0] : opaque[i][kind])) lod.pieces.push_back(&piece);
					any = any || !lod.pieces.empty();
					group.lods.push_back(std::move(lod));
				}
				if (any) out.push_back(std::move(group));
			}
			// Last, over everything: the sparkles over the opaque and the transparent glitter bricks
			if (sparkles) {
				UgcFormats::NifLodGroup group{ SparkleName(settings), false, {} };
				group.glitter = &settings.shaders.glitterParams;
				group.sparkle = true;
				bool any = false;
				for (size_t i = 0; i < levels; i++) {
					UgcFormats::NifLod lod{ ranges[i].first, ranges[i].second, "LOD_" + std::to_string(lods[i]), {} };
					for (const auto& piece : opaque[i][static_cast<size_t>(UgcModel::eLook::GLITTER)]) lod.pieces.push_back(&piece);
					for (const auto& piece : transparent[i][1]) lod.pieces.push_back(&piece);
					any = any || !lod.pieces.empty();
					group.lods.push_back(std::move(lod));
				}
				if (any) out.push_back(std::move(group));
			}
			return out;
		};
		const auto nif = UgcFormats::WriteLodNif("SceneNode_Model", groups(lods.size(), opaquePieces, transparentPieces));
		// Stored compressed only (the client downloads .gz; the dashboard's copies are inflated when asked for). The
		// LXFML is served from the database.
		AddDownload(outcome.files, "model.nif", nif);
		std::string plainNif;
		{
			const std::vector<LookPieces> opaque{ DivideByLook(preview.opaque, separate) };
			const std::vector<TransparentPieces> transparent{ transparentPieces[0] };
			plainNif = UgcFormats::WriteLodNif("SceneNode_Model", groups(1, opaque, transparent));
			outcome.files["model.noao.nif.gz"] = ZCompression::Gzip(plainNif);
		}

		// The icon is drawn from the .nif just made (its most detailed LOD, read back like any client .nif), so it
		// shows what the game shows: the colors with their variation, hidden faces removed, the baked lighting.
		const auto iconStart = std::chrono::steady_clock::now();
		auto iconOptions = settings.icon;
		UgcIconParams::Apply(iconOptions, iconValues);
		std::string nifError;
		if (!IconFromNif(nif, iconOptions, outcome.files, nifError, settings.shaders.TagLooks(), settings.shaders.OverlayTags(), &plainNif)) {
			outcome.error = "the .nif made can't be read back for the icon: " + nifError;
			return outcome;
		}
		const double iconMs = Since(iconStart);

		stats["ms"] = { { "build", std::lround(buildMs) }, { "hiddenSurfaces", std::lround(hsrMs) }, { "ambientOcclusion", std::lround(aoMs) },
			{ "icon", std::lround(iconMs) }, { "total", std::lround(Since(started)) } };
		stats["settings"] = { { "palette", settings.build.palette == UgcModel::ePalette::LU_TOOLBOX ? "lu_toolbox" : "brickdb" },
			{ "colorVariation", settings.build.colorVariation }, { "transparentOpacity", settings.build.transparentOpacity },
			{ "removeHiddenFaces", settings.hsr.enabled }, { "groundPlane", settings.hsr.groundPlane }, { "hsrResolution", settings.hsr.resolution },
			{ "ao", settings.ao.enabled }, { "aoDistance", settings.ao.distance }, { "aoSamples", settings.ao.samples }, { "aoStrength", settings.ao.strength } };
		const auto madeWith = MadeWith(settings);
		stats["settings"]["rays"] = madeWith.rays;
		stats["settings"]["denoise"] = madeWith.denoise;
		outcome.options = UgcProcessOptions::ToString(madeWith);
		outcome.stats = stats.dump();
		outcome.files["stats.json"] = outcome.stats;
		outcome.ok = true;
		return outcome;
	}

	std::optional<UgcModel::Model> AssembleModular(const ModularInput& input, const std::filesystem::path& res, glm::mat4& additionalRotation, std::string& error, std::string& note) {
		const auto build = UgcModular::ParseBuild(input.buildXml);
		if (!build) {
			error = "the build type has no topology in ModularBuildComponent";
			return std::nullopt;
		}
		additionalRotation = build->additionalRotation;
		std::vector<UgcModular::Module> modules;
		for (const auto& moduleInput : input.modules) {
			const auto path = UgcBricks::ResolvePath(res, moduleInput.renderAsset);
			const auto data = path ? UgcBricks::ReadFile(*path) : std::nullopt;
			if (!data) {
				note += "module " + std::to_string(moduleInput.lot) + " has no mesh (" + moduleInput.renderAsset + "); ";
				continue;
			}
			std::string nifError;
			auto nif = NifFile::Parse(*data, 0, nifError);
			if (!nif) {
				note += "module " + std::to_string(moduleInput.lot) + ": " + nifError + "; ";
				continue;
			}
			modules.push_back({ moduleInput.partCode, std::move(*nif), UgcModular::ParseModuleConnections(moduleInput.moduleXml) });
		}
		if (modules.empty()) {
			error = "none of the modules have a mesh";
			if (!note.empty()) error += " (" + note + ")";
			return std::nullopt;
		}
		auto model = UgcModular::Assemble(*build, modules, note);
		if (model.Empty()) {
			error = "the modules have no triangles";
			return std::nullopt;
		}
		return model;
	}

	Outcome ProcessModular(const ModularInput& input, const std::filesystem::path& res, const Settings& settings) {
		Outcome outcome;
		glm::mat4 additionalRotation{ 1.0f };
		const auto model = AssembleModular(input, res, additionalRotation, outcome.error, outcome.note);
		if (!model) return outcome;
		auto options = ModularIconOptions(input, settings);
		options.modelRotation = additionalRotation;
		AddIcon(outcome.files, *model, options);
		outcome.files["combo.json"] = nlohmann::json{ { "key", input.key }, { "buildType", input.buildType } }.dump();
		outcome.ok = true;
		return outcome;
	}

	std::optional<std::string> AssemblyNif(const ModularInput& input, const std::filesystem::path& res, std::string& error) {
		glm::mat4 additionalRotation{ 1.0f };
		std::string note;
		auto model = AssembleModular(input, res, additionalRotation, error, note);
		if (!model) return std::nullopt;
		// Turned as the icon renderer turns it before the pose's own rotation, so the editor's model rotation starts from here
		model->opaque.Transform(additionalRotation);
		model->transparent.Transform(additionalRotation);
		const auto opaque = UgcModel::Split(model->opaque), transparent = UgcModel::Split(model->transparent);
		std::vector<UgcFormats::NifShape> shapes;
		for (const auto& piece : opaque) shapes.push_back({ "S01_Opaque_Model", &piece, false });
		for (const auto& piece : transparent) shapes.push_back({ "S01_Alpha_Model", &piece, true });
		return UgcFormats::WriteNif("SceneNode_Assembly", shapes);
	}

	UgcRender::IconOptions ModularIconOptions(const ModularInput& input, const Settings& settings) {
		auto options = settings.icon;
		UgcIconParams::Apply(options, input.iconValues);
		return options;
	}


}