#include "MaterialViewer.h"

#include <Maya/MFnAreaLight.h>
#include <Maya/MFnCamera.h>
#include <Maya/MFnMesh.h>
#include <Maya/MRenderView.h>
#include <Maya/MPlug.h>
#include <Maya/MPlugArray.h>
#include <Maya/MThreadAsync.h>

#include "Atlas/Core/BvhAccel.h"
#include "Atlas/Core/Log.h"
#include "Atlas/Shapes/Rectangle.h"

#include "Utils.h"

const char *MaterialViewer::name = "Atlas";

namespace
{
	MThreadRetVal asyncFunction(void *t)
	{
		((MaterialViewer *)t)->threadFunction();
		return (0);
	}
}

void *MaterialViewer::creator()
{
	return (new MaterialViewer);
}

bool MaterialViewer::isSafeToUnload()
{
	return (true);
}

MStatus MaterialViewer::startAsync(const JobParams &params)
{
	LOG(DEBUG, GENERAL, "StartAsync");
	atlas::NextEventEstimation::Info rendererInfo;
	rendererInfo.maxLightBounce = 9;
	rendererInfo.tmin = (Float)0.01;
	rendererInfo.samplePerPixel = 64;
	rendererInfo.backgroundColor = atlas::Spectrum(0.10);
	rendererInfo.sampler = atlas::StratifiedSampler::create();
	rendererInfo.threadCount = std::thread::hardware_concurrency() / 2;
	rendererInfo.endOfIterationCallback = [this](const atlas::Film &film)
	{
		const uint32_t size = film.resolution.x * film.resolution.y;
		const Float invColorSum = 1. / film.getPixel(0).filterWeightSum;
		std::unique_ptr<RV_PIXEL[]> pixels(new RV_PIXEL[size]);
		for (uint32_t i = 0; i < size; i++)
		{
			const atlas::Film::Pixel &p = film.getPixel(size - i);
			pixels[i].r = p.color.r * invColorSum;
			pixels[i].g = p.color.g * invColorSum;
			pixels[i].b = p.color.b * invColorSum;
			pixels[i].a = 1;
		}

		RefreshParams params;
		params.width = film.resolution.x;
		params.height = film.resolution.y;
		params.bottom = 0;
		params.top = film.resolution.y - 1;
		params.left = 0;
		params.right = film.resolution.x - 1;
		params.channels = 4;
		params.bytesPerChannel = sizeof(float);
		params.data = pixels.get();
		refresh(params);
	};
	nee = std::make_unique<atlas::NextEventEstimation>(rendererInfo);

	isRunning = true;
	MThreadAsync::init();
	MThreadAsync::createTask(asyncFunction, this, nullptr, nullptr);
	return MS::kSuccess;
}

MStatus MaterialViewer::stopAsync()
{
	LOG(DEBUG, GENERAL, "StopAsync");
	isRunning = false;
	sleepCtrl.notify_all();
	LOG(DEBUG, GENERAL, "notify all");

	ProgressParams progressParams;
	progressParams.description = "StopAsync";
	progressParams.progress = -1;
	progress(progressParams);

	MThreadAsync::release();
	return MS::kSuccess;
}

bool MaterialViewer::isRunningAsync()
{
	return (true);
}

MStatus MaterialViewer::beginSceneUpdate()
{
	LOG(DEBUG, GENERAL, "BeginSceneUpdate");
	if (nee)
		nee->cancel();
	return MS::kSuccess;
}

MStatus MaterialViewer::endSceneUpdate()
{
	LOG(DEBUG, GENERAL, "EndSceneUpdate");
	requestRender = true;
	sleepCtrl.notify_all();
	return MS::kSuccess;
}

MStatus MaterialViewer::destroyScene()
{
	LOG(DEBUG, GENERAL, "DestroyScene");
	requestRender = false;
	nee.release();
	return MS::kSuccess;
}

MStatus MaterialViewer::setProperty(const MUuid &id, const MString &name, bool value)
{
	LOG(DEBUG, GENERAL, "Set property %s %s as %d", id.asString(), name.asChar(), value);
	return MS::kSuccess;
}

MStatus MaterialViewer::setProperty(const MUuid &id, const MString &name, int value)
{
	LOG(DEBUG, GENERAL, "Set property %s %s as %d", id.asString(), name.asChar(), value);
	return MS::kSuccess;
}

MStatus MaterialViewer::setProperty(const MUuid &id, const MString &name, float value)
{
	LOG(DEBUG, GENERAL, "Set property %s %s as %f", id.asString(), name.asChar(), value);
	return MS::kSuccess;
}

MStatus MaterialViewer::setProperty(const MUuid &id, const MString &name, const MString &value)
{
	LOG(DEBUG, GENERAL, "Set property %s %s as %s", id.asString(), name.asChar(), value.asChar());
	return MS::kSuccess;
}

MStatus MaterialViewer::setShader(const MUuid &id, const MUuid &shaderId)
{
	LOG(DEBUG, GENERAL, "Shader");
	return MS::kSuccess;
}

MStatus MaterialViewer::setResolution(unsigned int width, unsigned int height)
{
	LOG(INFO, GENERAL, "resolution %d %d", width, height);
	resolution.x = width;
	resolution.y = height;
	return MS::kSuccess;
}

MStatus MaterialViewer::translateMesh(const MUuid &id, const MObject &node)
{
	LOG(DEBUG, GENERAL, "Mesh");
	mMeshHandle = node;
	return MS::kSuccess;
}

MStatus MaterialViewer::translateLightSource(const MUuid &id, const MObject &node)
{
	LOG(DEBUG, GENERAL, "load light");
	lightExporters.emplace_back(node);
	lightExporters.back().deferredBuild(node);
	LOG(DEBUG, GENERAL, "end light");
	return MS::kSuccess;
}

MStatus MaterialViewer::translateCamera(const MUuid &id, const MObject &node)
{
	LOG(DEBUG, GENERAL, "Camera");
	mCameraHandle = node;
	return MS::kSuccess;
}

MStatus MaterialViewer::translateEnvironment(const MUuid &id, EnvironmentType type)
{
	LOG(DEBUG, GENERAL, "Env");
	return MS::kSuccess;
}

MStatus MaterialViewer::translateTransform(const MUuid &id, const MUuid &childId, const MMatrix &matrix)
{
	LOG(DEBUG, GENERAL, "Transform");
	unsigned char uid;
	childId.get(&uid);
	transforms[uid] = matrix;
	return MS::kSuccess;
}

MStatus MaterialViewer::translateShader(const MUuid &id, const MObject &node)
{
	mShaderHandle = node;
	return MS::kSuccess;
}

void MaterialViewer::threadFunction()
{
	while (isRunning)
	{
		{
			std::unique_lock<std::mutex> lock(locker);
			sleepCtrl.wait(lock, [this]() {
				return (!isRunning || requestRender);
				});
		}

		if (!isRunning)
			return;
		else
			requestRender = false;

		LOG(DEBUG, GENERAL, "New Render");

		if (!mMeshHandle.isValid())
			continue;

		{
			ProgressParams progressParams;
			progressParams.description = "Rendering";
			progressParams.progress = 0;
			progress(progressParams);
		}

		atlas::FilmInfo filmInfo;
		filmInfo.filename = "D:\\Uedaki\\Documents\\unamed.ppm";
		filmInfo.resolution = resolution;
		filmInfo.filter = new atlas::BoxFilter(atlas::Vec2f(0.5, 0.5));
		atlas::Film film(filmInfo);

		atlas::Bounds2f screen;
		screen.min.x = -1.f;
		screen.max.x = 1.f;
		screen.min.y = -1.f;
		screen.max.y = 1.f;

		MFnCamera mCamera(mCameraHandle.object());
		unsigned char cameraUid;
		mCamera.uuid().get(&cameraUid);
		MPoint mOrigin = mCamera.eyePoint(MSpace::kWorld) * transforms[cameraUid];
		MVector mDirection = mCamera.viewDirection(MSpace::kWorld) * transforms[cameraUid];
		MPoint mInterstPoint = mCamera.centerOfInterestPoint(MSpace::kWorld) * transforms[cameraUid];
		MVector mUp = mCamera.upDirection(MSpace::kWorld) * transforms[cameraUid];
		atlas::Point3f origin(mOrigin.x, mOrigin.y, mOrigin.z);
		atlas::Point3f poi(mInterstPoint.x, mInterstPoint.y, mInterstPoint.z);
		atlas::Vec3f dir(mDirection.x, mDirection.y, mDirection.z);
		atlas::Transform worldToCam = atlas::Transform::lookAt(origin, poi, atlas::Vec3f(mUp.x, mUp.y, mUp.z));
		atlas::Transform cameraToWorld = worldToCam.inverse();
		double vFov = mCamera.verticalFieldOfView();
		double hFov = mCamera.horizontalFieldOfView();
		atlas::PerspectiveCameraInfo cameraInfo;
		cameraInfo.cameraToWorld = &cameraToWorld;
		cameraInfo.screenWindow = screen;
		cameraInfo.shutterOpen = 0.f;
		cameraInfo.shutterClose = 1.f;
		cameraInfo.lensRadius = 0;
		cameraInfo.focalDistance = 10;
		cameraInfo.verticalFov = atlas::degrees(vFov);
		cameraInfo.horizontalFov = atlas::degrees(hFov);
		cameraInfo.film = &film;
		atlas::PerspectiveCamera camera(cameraInfo);

		std::shared_ptr<MaterialExporter> material = std::make_shared<MaterialExporter>(mShaderHandle.object());
		material->build();

		unsigned char meshUid;
		MeshExporter mesh(mMeshHandle.object());
		mesh.getUuid().get(&meshUid);
		MMatrix *meshTransform = nullptr;
		const auto &transformIt = transforms.find(meshUid);
		if (transformIt != transforms.cend())
			meshTransform = &transformIt->second;
		mesh.build(meshTransform);
		std::vector<std::shared_ptr<atlas::Primitive>> primitives = mesh.exportPrimitives(material);
		atlas::BvhAccel scene(primitives);

		LOG(DEBUG, GENERAL, "load light list");
		std::vector<std::shared_ptr<atlas::Light>> lights;
		for (AreaLightExporter &exporter : lightExporters)
		{
			exporter.constructLightFromDeferredBuild(transforms[exporter.getUuid()]);
			lights.emplace_back(exporter.getLight());
		}
		LOG(DEBUG, GENERAL, "end light list");
		nee->render(camera, scene, lights, film);

		const uint32_t size = resolution.x * resolution.y;
		const Float invColorSum = 1. / film.getPixel(0).filterWeightSum;
		std::unique_ptr<RV_PIXEL[]> pixels = std::make_unique<RV_PIXEL[]>(size);
		for (uint32_t i = 0; i < size; i++)
		{
			const atlas::Film::Pixel &p = film.getPixel(size - i);
			pixels[i].r = p.color.r * invColorSum;
			pixels[i].g = p.color.g * invColorSum;
			pixels[i].b = p.color.b * invColorSum;
			pixels[i].a = 1;
		}

		{
			RefreshParams params;
			params.width = resolution.x;
			params.height = resolution.y;
			params.bottom = 0;
			params.top = resolution.y - 1;
			params.left = 0;
			params.right = resolution.x - 1;
			params.channels = 4;
			params.bytesPerChannel = sizeof(float);
			params.data = pixels.get();
			refresh(params);

			ProgressParams progressParams;
			progressParams.description = "Rendering";
			progressParams.progress = 1;
			progress(progressParams);
		}

		LOG(DEBUG, GENERAL, "Render Complete");
	}
	LOG(DEBUG, GENERAL, "-> Exit");
}