//	GGVRStereoscopicVisionTest.swift
//
//	For use in psychology experiments with Elyssa Twedt and her students,
//	to screen participants for normal stereoscopic vision.
//
//	© 2026 by Jeff Weeks
//	See TermsOfUse.txt

import CompositorServices


actor GGVRStereoscopicVisionTestRenderer {

	let layerRenderer: LayerRenderer
	var openControlPanel: @Sendable () -> Void

	let arKitSession: ARKitSession
	let worldTrackingProvider: WorldTrackingProvider

	//	I'm slightly uncomfortable using implicitly unwrapped optionals
	//	-- I hadn't intended to use them at all, for run-time safety --
	//	but in this case I really would rather not have to get
	//	all the Metal objects set up directly in init(),
	//	without being able to make any function calls.
	//	So let's go ahead and use implicitly unwrapped optionals here.
	//
	var device: MTLDevice!
	var commandQueue: MTLCommandQueue!
	var pipelineState: MTLRenderPipelineState!
	var depthStencilState: MTLDepthStencilState!

	var triangleA: MTLBuffer!	//	array of SIMD3<Float32>
	var triangleB: MTLBuffer!	//	array of SIMD3<Float32>
	var triangleC: MTLBuffer!	//	array of SIMD3<Float32>


	init?(
		layerRenderer: LayerRenderer,
		openControlPanel: @Sendable @escaping () -> Void
	) async {
	
		self.layerRenderer = layerRenderer
		self.openControlPanel = openControlPanel

		self.arKitSession = ARKitSession()
		self.worldTrackingProvider =  WorldTrackingProvider()

		if !setUpMetalObjects() {
			return nil
		}
	}
	
	deinit {
		openControlPanel()
	}

// MARK: -
// MARK: World tracking

	func startWorldTracking() async {
	
		do {
			try await arKitSession.run([worldTrackingProvider])
		} catch {
			preconditionFailure("ARKitSession error: \(error)")
		}
	}

// MARK: -
// MARK: Set up Metal objects

	func setUpMetalObjects() -> Bool {
		
		device = layerRenderer.device
		
		guard let theCommandQueue = device.makeCommandQueue()
		else { return false }
		commandQueue = theCommandQueue

		if  setUpPipelineState()
		 && setUpDepthStencilState()
		 && setUpMeshes()
		{
			return true
			
		} else {
		
			assertionFailure("Unable to set up Metal resources")
			return false
		}
	}

	func setUpPipelineState() -> Bool {

		guard let theGPUFunctionLibrary = device.makeDefaultLibrary() else {
			assertionFailure("Failed to create default GPU function library")
			return false
		}
		
		guard let theGPUVertexFunction = theGPUFunctionLibrary.makeFunction(
			name: "SVTVertexFunction")
		else {
			assertionFailure("Failed to load SVTVertexFunction")
			return false
		}
		guard let theGPUFragmentFunction = theGPUFunctionLibrary.makeFunction(
			name: "SVTFragmentFunction")
		else {
			assertionFailure("Failed to load SVTFragmentFunction")
			return false
		}
		
		let theVertexDescriptor = MTLVertexDescriptor()

		theVertexDescriptor.attributes[0].format = MTLVertexFormat.float3
		theVertexDescriptor.attributes[0].bufferIndex = SVTVertexAttributes
		theVertexDescriptor.attributes[0].offset = 0

		theVertexDescriptor.layouts[SVTVertexAttributes].stepFunction = MTLVertexStepFunction.perVertex
		theVertexDescriptor.layouts[SVTVertexAttributes].stride = MemoryLayout<SIMD3<Float32>>.stride

		let thePipelineDescriptor = MTLRenderPipelineDescriptor()
		thePipelineDescriptor.label = "Stereoscopic vision test render pipeline"
		thePipelineDescriptor.vertexFunction = theGPUVertexFunction
		thePipelineDescriptor.fragmentFunction = theGPUFragmentFunction
		thePipelineDescriptor.vertexDescriptor = theVertexDescriptor
		thePipelineDescriptor.colorAttachments[0].pixelFormat = layerRenderer.configuration.colorFormat
		thePipelineDescriptor.colorAttachments[0].isBlendingEnabled = false
		thePipelineDescriptor.depthAttachmentPixelFormat = layerRenderer.configuration.depthFormat
		thePipelineDescriptor.stencilAttachmentPixelFormat = MTLPixelFormat.invalid
		thePipelineDescriptor.maxVertexAmplificationCount = layerRenderer.properties.viewCount

		do {

			pipelineState = try device.makeRenderPipelineState(
									descriptor: thePipelineDescriptor)
			return true

		} catch {

			assertionFailure("Couldn't create the pipelineState: \(error)")
			return false
		}
	}

	func setUpDepthStencilState() -> Bool {

		guard let theDepthStencilState = ggvrDepthStencilState(
										device: device,
										depthCompareEnabled: true,
										depthWriteEnabled: true)
		else { return false }
		depthStencilState = theDepthStencilState

		return true
	}
	
	func setUpMeshes() -> Bool {
	
		let theRawVerticesA = [
			SIMD4<Float32>(-0.10, -0.10, 0.0, 1.0),
			SIMD4<Float32>(-0.05,  0.12, 0.0, 1.0),
			SIMD4<Float32>( 0.12, -0.06, 0.0, 1.0)
		]
		let theTranslationA = simd_float4x4(
			SIMD4<Float32>( 1.0,  0.0,  0.0,  0.0 ),
			SIMD4<Float32>( 0.0,  1.0,  0.0,  0.0 ),
			SIMD4<Float32>( 0.0,  0.0,  1.0,  0.0 ),
			SIMD4<Float32>( 0.0,  0.1, -1.8,  1.0 )
		)
		let theVerticesA = theRawVerticesA.map() { v in
			let vv = theTranslationA * v
			return SIMD3<Float32>(vv.x, vv.y, vv.z)
		}
		guard let theTriangleA = device.makeBuffer(
			bytes: theVerticesA,
			length: theVerticesA.count * MemoryLayout<SIMD3<Float32>>.stride,
			options: [])
		else {
			assertionFailure("Failed to create theTriangleA")
			return false
		}
		triangleA = theTriangleA

		let theRawVerticesB = [
			SIMD4<Float32>(-0.02, -0.08, 0.0, 1.0),
			SIMD4<Float32>(-0.06,  0.08, 0.0, 1.0),
			SIMD4<Float32>( 0.09,  0.02, 0.0, 1.0)
		]
		let theTranslationB = simd_float4x4(
			SIMD4<Float32>( 1.0,  0.0,  0.0,  0.0 ),
			SIMD4<Float32>( 0.0,  1.0,  0.0,  0.0 ),
			SIMD4<Float32>( 0.0,  0.0,  1.0,  0.0 ),
			SIMD4<Float32>(-0.2, -0.1, -1.4,  1.0 )
		)
		let theVerticesB = theRawVerticesB.map() { v in
			let vv = theTranslationB * v
			return SIMD3<Float32>(vv.x, vv.y, vv.z)
		}
		guard let theTriangleB = device.makeBuffer(
			bytes: theVerticesB,
			length: theVerticesB.count * MemoryLayout<SIMD3<Float32>>.stride,
			options: [])
		else {
			assertionFailure("Failed to create theTriangleB")
			return false
		}
		triangleB = theTriangleB

		let theRawVerticesC = [
			SIMD4<Float32>(-0.06, -0.02, 0.0, 1.0),
			SIMD4<Float32>( 0.00,  0.05, 0.0, 1.0),
			SIMD4<Float32>( 0.01, -0.04, 0.0, 1.0)
		]
		let theTranslationC = simd_float4x4(
			SIMD4<Float32>( 1.0,  0.0,  0.0,  0.0 ),
			SIMD4<Float32>( 0.0,  1.0,  0.0,  0.0 ),
			SIMD4<Float32>( 0.0,  0.0,  1.0,  0.0 ),
			SIMD4<Float32>( 0.2,  0.0, -1.6,  1.0 )
		)
		let theVerticesC = theRawVerticesC.map() { v in
			let vv = theTranslationC * v
			return SIMD3<Float32>(vv.x, vv.y, vv.z)
		}
		guard let theTriangleC = device.makeBuffer(
			bytes: theVerticesC,
			length: theVerticesC.count * MemoryLayout<SIMD3<Float32>>.stride,
			options: [])
		else {
			assertionFailure("Failed to create theTriangleC")
			return false
		}
		triangleC = theTriangleC

		return true
	}

// MARK: -
// MARK: Render loop

	func runRenderLoop() {

		var keepRendering = true	//	remains true until the immersive space gets closed

		while keepRendering {
		
			switch layerRenderer.state {
			
			case .paused:
				layerRenderer.waitUntilRunning()
				
			case .running:
				autoreleasepool {
					renderOneFrame()
				}
				
			case .invalidated:
				fallthrough
			@unknown default:
				keepRendering = false
			}
		}
		
		arKitSession.stop()
	}

	func renderOneFrame() {

		guard let theFrame = layerRenderer.queryNextFrame(),
			  let theTiming = theFrame.predictTiming()
		else {
			assertionFailure("Couldn't get theFrame and/or theTiming")
			return
		}
		
		//----------
		//	Update
		//----------

		theFrame.startUpdate()
		theFrame.endUpdate()


		//-----------------
		//	Get drawables
		//-----------------

		LayerRenderer.Clock().wait(until: theTiming.optimalInputTime)

		let theDrawables = theFrame.queryDrawables()
		if theDrawables.isEmpty { return }
		
		
		//----------------
		//	Submit frame
		//----------------

		theFrame.startSubmission()
		for theDrawable in theDrawables {
			renderTheFrame(drawable: theDrawable)
		}
		theFrame.endSubmission()
	}

	func renderTheFrame(
		drawable: LayerRenderer.Drawable
	) {

		let thePresentationTime = LayerRenderer.Clock.Instant.epoch.duration(to:
			drawable.frameTiming.presentationTime).timeInterval

		drawable.deviceAnchor	//	could be nil
			= worldTrackingProvider.queryDeviceAnchor(atTimestamp: thePresentationTime)

		guard let theCommandBuffer = commandQueue.makeCommandBuffer() else {
			preconditionFailure("Failed to make theCommandBuffer")
		}

		let theRenderPassDescriptor = ggvrRenderPassDescriptor(
										layerRenderer: layerRenderer,
										drawable: drawable)
		
		let theViewProjections = (0 ..< drawable.views.count).map() { eyeIndex in

			//	This Stereoscopic Vision Test completely ignores
			//	the player's head motions. The headset stays fixed
			//	at the identity pose in world space.
			let theEyePoseInHeadsetSpace = drawable.views[eyeIndex].transform
			let theHeadsetPoseInWorldSpace = matrix_identity_float4x4
			let theEyePoseInWorldSpace = theHeadsetPoseInWorldSpace	//	right-to-left matrix action
									   * theEyePoseInHeadsetSpace
			let theViewMatrix = theEyePoseInWorldSpace.inverse
			let theProjectionMatrix = drawable.computeProjection(viewIndex: eyeIndex)
			let theViewProjectionMatrix = theProjectionMatrix		//	right-to-left matrix action
										* theViewMatrix
			
			return theViewProjectionMatrix
		}
		
		if let theRenderEncoder = theCommandBuffer.makeRenderCommandEncoder(
										descriptor: theRenderPassDescriptor) {
			theRenderEncoder.label = "Stereoscopic vision test render encoder"

			ggvrEncodeViewports(
				drawable: drawable,
				renderEncoder: theRenderEncoder)

			encodeCommands(
				renderEncoder: theRenderEncoder,
				viewProjections: theViewProjections)
				
			theRenderEncoder.endEncoding()
		}

		drawable.encodePresent(commandBuffer: theCommandBuffer)
		theCommandBuffer.commit()
	}

	func encodeCommands(
		renderEncoder: MTLRenderCommandEncoder,
		viewProjections: [ simd_float4x4 ]
	) {

		var theColorA = SIMD3<Float16>(0.00, 0.75, 0.25)
		var theColorB = SIMD3<Float16>(0.25, 0.00, 1.00)
		var theColorC = SIMD3<Float16>(1.00, 0.84, 0.00)

		renderEncoder.setRenderPipelineState(pipelineState)
		renderEncoder.setDepthStencilState(depthStencilState)
		renderEncoder.setCullMode(MTLCullMode.back)
		renderEncoder.setFrontFacing(MTLWinding.clockwise)	//	the default

		renderEncoder.setVertexBytes(
						viewProjections,
						length: viewProjections.count * MemoryLayout<simd_float4x4>.stride,
						index: SVTViewProjection)

		renderEncoder.setVertexBuffer(
						triangleA,
						offset: 0,
						index: SVTVertexAttributes)
		renderEncoder.setFragmentBytes(
						&theColorA,
						length: MemoryLayout.size(ofValue: theColorA),
						index: SVTColor)
		renderEncoder.drawPrimitives(
						type: .triangleStrip,
						vertexStart: 0,
						vertexCount: triangleA.length / MemoryLayout<SIMD3<Float32>>.stride)

		renderEncoder.setVertexBuffer(
						triangleB,
						offset: 0,
						index: SVTVertexAttributes)
		renderEncoder.setFragmentBytes(
						&theColorB,
						length: MemoryLayout.size(ofValue: theColorB),
						index: SVTColor)
		renderEncoder.drawPrimitives(
						type: .triangleStrip,
						vertexStart: 0,
						vertexCount: triangleB.length / MemoryLayout<SIMD3<Float32>>.stride)

		renderEncoder.setVertexBuffer(
						triangleC,
						offset: 0,
						index: SVTVertexAttributes)
		renderEncoder.setFragmentBytes(
						&theColorC,
						length: MemoryLayout.size(ofValue: theColorC),
						index: SVTColor)
		renderEncoder.drawPrimitives(
						type: .triangleStrip,
						vertexStart: 0,
						vertexCount: triangleC.length / MemoryLayout<SIMD3<Float32>>.stride)
	}
}
