/*
===========================================================================
Copyright (C) 2024 Gian 'myT' Schellenbaum

This file is part of Challenge Quake 3 (CNQ3).

Challenge Quake 3 is free software; you can redistribute it
and/or modify it under the terms of the GNU General Public License as
published by the Free Software Foundation; either version 2 of the License,
or (at your option) any later version.

Challenge Quake 3 is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with Challenge Quake 3. If not, see <https://www.gnu.org/licenses/>.
===========================================================================
*/
// direct lighting from dynamic lights


#include "common.hlsli"
#include "fullscreen.hlsli"
#include "raytracing.h.hlsli"
#include "scene_view.h.hlsli"
#include "alpha_test.h.hlsli"


cbuffer RootConstants
{
	uint blueNoiseTextureIndex;
};

#define CLASS_OPAQUE 0u
#define CLASS_INVISIBLE 1u
#define CLASS_TRANSLUCENT 2u

uint ClassifyNonOpaqueTriangle(inout float3 light, StructuredBuffer<TLASInstance> tlasInstanceBuffer, uint instanceId, uint meshId, uint triangleId, float2 bary2, bool frontFace)
{
	TLASInstance instance = tlasInstanceBuffer[instanceId];
#if 0
	// @TODO: is this needed or not?
	// cull mode: 0 is front-sided, 1 is back-sided
	if((frontFace && instance.cullMode == 1) ||
		(!frontFace && instance.cullMode == 0))
	{
		return CLASS_INVISIBLE;
	}
#endif
	StructuredBuffer<BLASMesh> meshBuffer = ResourceDescriptorHeap[instance.meshBufferIndex];
	BLASMesh mesh = meshBuffer[meshId];
	float3 barycentrics = float3(1.0 - bary2.x - bary2.y, bary2.x, bary2.y);
	StructuredBuffer<BLASVertex> vertexBuffer = ResourceDescriptorHeap[instance.vertexBufferIndex];
	StructuredBuffer<uint> indexBuffer = ResourceDescriptorHeap[instance.indexBufferIndex];
	uint firstIndex = mesh.firstIndex + triangleId * 3;
	uint vtxIdx0 = mesh.firstVertex + indexBuffer[firstIndex + 0];
	uint vtxIdx1 = mesh.firstVertex + indexBuffer[firstIndex + 1];
	uint vtxIdx2 = mesh.firstVertex + indexBuffer[firstIndex + 2];
	BLASVertex v0 = vertexBuffer[vtxIdx0];
	BLASVertex v1 = vertexBuffer[vtxIdx1];
	BLASVertex v2 = vertexBuffer[vtxIdx2];
	float2 texCoords = trilerp(v0.texCoords, v1.texCoords, v2.texCoords, barycentrics);
	float4 vertexColor = trilerp(UnpackColor(v0.color), UnpackColor(v1.color), UnpackColor(v2.color), barycentrics);
	Texture2D texture0 = ResourceDescriptorHeap[mesh.textureIndex];
	SamplerState sampler0 = SamplerDescriptorHeap[mesh.samplerIndex];
	float4 textureColor = texture0.SampleLevel(sampler0, texCoords, 0);
	float4 hitColor = vertexColor * textureColor;
	if(mesh.alphaTestMode == 0)
	{
		float3 blended;
		if(mesh.blendBits == (GLS_SRCBLEND_ONE | GLS_DSTBLEND_ONE))
		{
			blended = lerp(light, hitColor.rgb, Brightness(hitColor.rgb));
		}
		else
		{
			blended = Blend(hitColor, float4(light, 1), mesh.blendBits).rgb;
		}
		if(all(blended == light))
		{
			return CLASS_INVISIBLE;
		}
		light = blended;
		return CLASS_TRANSLUCENT;
	}
	if(FailsAlphaTest(hitColor.a, mesh.alphaTestMode))
	{
		return CLASS_INVISIBLE;
	}
	return CLASS_OPAQUE;
}

float2 MapSquareToDisk(float2 square01)
{
	float radius = sqrt(square01.x);
	float angle = square01.y * 2.0 * 3.14159265359;
	float2 sinCos;
	sincos(angle, sinCos.x, sinCos.y);
	float2 result = radius * sinCos;
	
	return result;
}

float3 GetRayDirectionForSphereLight(float2 square01, float3 surfacePos, float3 lightPos, float worldRadius)
{
	float3 direction = normalize(lightPos - surfacePos);
	float  radius = worldRadius / length(lightPos - surfacePos);
	float2 pointInDisk = MapSquareToDisk(square01) * radius;
	float3 tangent = normalize(cross(direction, float3(0, 1, 0)));
	float3 bitangent = normalize(cross(tangent, direction));
	float3 result = normalize(direction + pointInDisk.x * tangent + pointInDisk.y * bitangent);

	return result;
}

// true when fully in shadow
bool TraceShadowRay(
	out float t, inout float3 light,
	RaytracingAccelerationStructure rtas, StructuredBuffer<TLASInstance> instBuffer,
	float3 position, float3 direction, float dist)
{
	RayDesc ray;
	ray.Origin = position;
	ray.Direction = direction;
	ray.TMin = 0.0;
	ray.TMax = dist;

	t = 0.0;
	float translucentT = 0.0;
	RayQuery<RAY_FLAG_NONE> q;
	q.TraceRayInline(rtas, RAY_FLAG_NONE, 0xFF, ray);
	while(q.Proceed())
	{
		if(q.CandidateType() == CANDIDATE_NON_OPAQUE_TRIANGLE)
		{
			uint type = ClassifyNonOpaqueTriangle(
				light,
				instBuffer,
				q.CandidateInstanceIndex(),
				q.CandidateGeometryIndex(),
				q.CandidatePrimitiveIndex(),
				q.CandidateTriangleBarycentrics(),
				q.CandidateTriangleFrontFace());
			if(type == CLASS_OPAQUE)
			{
				q.CommitNonOpaqueTriangleHit();
			}
			else if(type == CLASS_TRANSLUCENT)
			{
				translucentT = q.CandidateTriangleRayT();
				t = translucentT;
			}
		}
	}

	if(q.CommittedStatus() == COMMITTED_TRIANGLE_HIT)
	{
		t = q.CommittedRayT();
	}

	if(q.CommittedStatus() == COMMITTED_TRIANGLE_HIT && t > translucentT)
	{
		return true;
	}

	return false;
}

// true when fully in shadow
bool TraceShadowRayOpaqueOnly(
	out float t, inout float3 light,
	RaytracingAccelerationStructure rtas, StructuredBuffer<TLASInstance> instBuffer,
	float3 position, float3 direction, float dist)
{
	RayDesc ray;
	ray.Origin = position;
	ray.Direction = direction;
	ray.TMin = 0.0;
	ray.TMax = dist;

	t = 0.0;
	bool keepLight = false;
	RayQuery<RAY_FLAG_CULL_NON_OPAQUE> q;
	q.TraceRayInline(rtas, RAY_FLAG_NONE, 0xFF, ray);
	q.Proceed();
	if(q.CommittedStatus() == COMMITTED_TRIANGLE_HIT)
	{
		t = q.CommittedRayT();
		return true;
	}

	return false;
}

float4 ps(VOut input) : SV_Target
{
	SceneView scene = GetSceneView();
	RaytracingAccelerationStructure rtas = ResourceDescriptorHeap[scene.tlasBufferIndex];
	Texture2D<float2> normalTexture = ResourceDescriptorHeap[scene.normalTextureIndex];
	Texture2D shadingPositionTexture = ResourceDescriptorHeap[scene.shadingPositionTextureIndex];
	StructuredBuffer<TLASInstance> tlasInstanceBuffer = ResourceDescriptorHeap[scene.tlasInstanceBufferIndex];
	Texture2D blueNoiseTexture = ResourceDescriptorHeap[blueNoiseTextureIndex];

	uint2 blueNoiseTextureSize;
	blueNoiseTexture.GetDimensions(blueNoiseTextureSize.x, blueNoiseTextureSize.y);

	uint3 tc = uint3(input.position.xy, 0);
	float3 normalWS = normalize(OctDecode(normalTexture.Load(tc)));
	float3 positionWS = shadingPositionTexture.Load(tc).xyz;

	float error = 0.0;
	float3 pixelAccum = float3(0, 0, 0);
	for(uint i = 0; i < scene.lightCount; i++)
	{
		float3 lightPosition = scene.lights[i].position;
		float dist = distance(positionWS, lightPosition);
		float radius = scene.lights[i].radius;
		if(dist >= radius)
		{
			continue;
		}

		float innerRadius = radius / 100.0;
		float intensity = saturate(1.0 - dist / radius);
		float3 lightDir = normalize(lightPosition - positionWS);
		float3 lightRaw = scene.lights[i].color * intensity * max(dot(normalWS, lightDir), 0.0);
		const uint SampleCount = 4;

		float3 lightAccum = float3(0, 0, 0);
		for(uint r = 0; r < SampleCount; r++)
		{
			float3 light = lightRaw;
			uint2 pos = uint2(input.position.xy) + uint2(r * 17, r * 13 + 7);
			uint2 tc = pos % blueNoiseTextureSize;
			float2 square01 = blueNoiseTexture.Load(uint3(tc, 0)).xy;
			float3 dir = GetRayDirectionForSphereLight(square01, positionWS, lightPosition, innerRadius);
			float t;
			bool inShadow = TraceShadowRay(t, light, rtas, tlasInstanceBuffer, positionWS, dir, dist);
			error = max(error, t / radius);
			if(inShadow)
			{
				continue;
			}

			lightAccum += light;
		}

		pixelAccum += lightAccum / float(SampleCount);
	}

	float4 result = float4(pixelAccum, saturate(error));

	return result;
}
