package kha.js.vr;

import js.Syntax;
import js.lib.Float32Array;
import kha.vr.Pose;
import kha.vr.PoseState;
import kha.vr.SensorState;
import kha.vr.TimeWarpParms;
import kha.math.FastMatrix4;
import kha.math.Vector3;
import kha.math.Quaternion;
import kha.SystemImpl;

class VrInterface extends kha.vr.VrInterface {
	var vrEnabled: Bool = false;
	var isWebXR: Bool = false;

	var vrDisplay: Dynamic;
	var frameData: Dynamic;

	var xrSession: Dynamic;
	var xrRefSpace: Dynamic;
	public var xrGLLayer: Dynamic;
	public var currentFrame: Dynamic;
	public var currentViews: Dynamic;
	public var currentViewerPose: Dynamic;
	public var currentInputSources: Dynamic;
	var xrAnimationFrameHandle: Int = -1;

	public var _glContext: Dynamic;
	public var _leftViewport: Dynamic;
	public var _rightViewport: Dynamic;
	public var _cachedViewsLength: Int = 0;
	
	var savedCanvasWidth: Int = 0;
	var savedCanvasHeight: Int = 0;
	var browserRAFId: Int = -1;

	var leftProjectionMatrix: FastMatrix4 = FastMatrix4.identity();
	var rightProjectionMatrix: FastMatrix4 = FastMatrix4.identity();
	var leftViewMatrix: FastMatrix4 = FastMatrix4.identity();
	var rightViewMatrix: FastMatrix4 = FastMatrix4.identity();

	var width: Int = 0;
	var height: Int = 0;
	var vrWidth: Int = 0;
	var vrHeight: Int = 0;

	public function new() {
		super();
		#if kha_webvr
		var webXREnabled: Bool = Syntax.code("navigator.xr");
		if (webXREnabled) {
			isWebXR = true;
			vrEnabled = true;
		}
		else {
			var displayEnabled: Bool = Syntax.code("navigator.getVRDisplays");
			if (displayEnabled) {
				isWebXR = false;
				vrEnabled = true;
				getVRDisplays();
				trace("WebVR 1.1 API detected");
			}
		}
		#else
		var displayEnabled = false;
		#end
		if (displayEnabled) {
			vrEnabled = true;
			getVRDisplays();
			trace("Display enabled.");
		}
	}

	function getVRDisplays() {
		var vrDisplayInstance = Syntax.code("navigator.getVRDisplays()");
		vrDisplayInstance.then(function(displays) {
			if (displays.length > 0) {
				frameData = Syntax.code("new VRFrameData()");
				vrDisplay = Syntax.code("displays[0]");
				vrDisplay.depthNear = 0.1;
				vrDisplay.depthFar = 1024.0;

				var leftEye = vrDisplay.getEyeParameters("left");
				var rightEye = vrDisplay.getEyeParameters("right");
				width = SystemImpl.khanvas.width;
				height = SystemImpl.khanvas.height;
				vrWidth = Std.int(Math.max(leftEye.renderWidth, rightEye.renderWidth) * 2);
				vrHeight = Std.int(Math.max(leftEye.renderHeight, rightEye.renderHeight));
			}
			else {
				trace("There are no VR displays connected.");
			}
		});
	}

	public override function onVRRequestPresent() {
		if (isWebXR) {
			requestWebXRSession();
		} else {
			// WebVR 1.1
			try {
				vrDisplay.requestPresent([{source: SystemImpl.khanvas}]).then(function() {
					onResize();
					vrDisplay.requestAnimationFrame(onAnimationFrame);
				});
			}
			catch (err:Dynamic) {
				trace("Failed to requestPresent WebVR: " + err);
			}
		}
	}
	
	function requestWebXRSession() {
		var vrScaleFactor = 1.0;
		#if lnx_vr
		vrScaleFactor = leenkx.renderpath.Inc.getSuperSampling();
		trace("[VR] Using renderpath superSample as framebufferScaleFactor: " + vrScaleFactor);
		#end
		
		try {
			Syntax.code("
				

				let gl = null;
				let canvas = null;
				
				try {
					if (typeof kha_SystemImpl !== 'undefined') {
						gl = kha_SystemImpl.gl;
						canvas = kha_SystemImpl.khanvas;
					}
				} catch (e) {
					trace('kha_SystemImpl access failed: ' + e.message;
				}

				if (!canvas) {
					canvas = document.querySelector('canvas');
				}
				
				if (canvas && !gl) {
					const contextAttributes = { xrCompatible: true, antialias: true, alpha: false };
					gl = canvas.getContext('webgl2', contextAttributes) || 
					     canvas.getContext('webgl', contextAttributes) || 
					     canvas.getContext('experimental-webgl', contextAttributes);
				}
				
				if (!canvas) {
					canvas = document.getElementById('khanvas');
					if (canvas && !gl) {
						gl = canvas.getContext('webgl2') || canvas.getContext('webgl');
					}
				}
				
				if (!gl) {
					return;
				}
				const self = this;
				const glContext = gl;
				
				self._glContext = glContext;
			");

			Syntax.code("
				self._vrRenderCallback = function() {
					self.vrRenderCallback();
				};

				const checkAndRequestSession = async () => {
					try {
						const supported = await navigator.xr.isSessionSupported('immersive-vr');
						if (!supported) {
							trace('immersive-vr session not supported');
						}
					} catch (e) {
						trace('WARN: isSessionSupported failed: ' + e.message);
						// Continue anyway as some browsers do not support the check itself
					}
					
					return await navigator.xr.requestSession('immersive-vr', {
						optionalFeatures: ['local-floor', 'hand-tracking', 'bounded-floor']
					});
				};
				
				checkAndRequestSession().then(async (session) => {
					self.xrSession = session;
					if (typeof window !== 'undefined') {
						window._khaSkipWindowRender = true;
					}
					
					const contextAttributes = glContext.getContextAttributes();
					if (!contextAttributes || !contextAttributes.xrCompatible) {
						await glContext.makeXRCompatible();
					}
					self.xrGLLayer = new XRWebGLLayer(session, glContext, {
						depth: true,       // Essential for depth testing
						stencil: false,    // Not needed, wastes memory
						alpha: false,      // Not needed in VR wastes
						antialias: true,   // Smooth rendering
						framebufferScaleFactor: {0}  // VR resolution quality from renderpath
					});
					
					if (self.xrGLLayer.framebufferWidth === 0 || self.xrGLLayer.framebufferHeight === 0) {
						trace('XRWebGLLayer framebuffer has invalid dimensions');
					}
					

					session.updateRenderState({
						baseLayer: self.xrGLLayer
					});
					const handlers = {};
					
					handlers.end = () => {
						self.onSessionEnd();
					};
					session.addEventListener('end', handlers.end);
					
					handlers.select = (event) => {
						if (self.onSelect) self.onSelect(event);
					};
					session.addEventListener('select', handlers.select);
					
					handlers.selectstart = (event) => {
						if (self.onSelectStart) self.onSelectStart(event);
					};
					session.addEventListener('selectstart', handlers.selectstart);
					
					handlers.selectend = (event) => {
						if (self.onSelectEnd) self.onSelectEnd(event);
					};
					session.addEventListener('selectend', handlers.selectend);
					
					handlers.squeeze = (event) => {
						if (self.onSqueeze) self.onSqueeze(event);
					};
					session.addEventListener('squeeze', handlers.squeeze);
					
					handlers.squeezestart = (event) => {
						if (self.onSqueezeStart) self.onSqueezeStart(event);
					};
					session.addEventListener('squeezestart', handlers.squeezestart);
					
					handlers.squeezeend = (event) => {
						if (self.onSqueezeEnd) self.onSqueezeEnd(event);
					};
					session.addEventListener('squeezeend', handlers.squeezeend);
					
					session.addEventListener('inputsourceschange', handlers.inputsourceschange);
					
					handlers.visibilitychange = (event) => {
						const state = event.session.visibilityState;

					};
					session.addEventListener('visibilitychange', handlers.visibilitychange);
					
					self._eventHandlers = handlers;
					
					const requestRefSpace = async () => {
						const spaces = ['local-floor', 'local'];
						for (const space of spaces) {
							try {
								const refSpace = await session.requestReferenceSpace(space);
								return refSpace;
							} catch (e) {
								trace(space + ' not supported');
							}
						}
						trace('No reference space supported');
					};
					
					requestRefSpace().then((refSpace) => {
						self.xrRefSpace = refSpace;
						
						if (canvas && canvas.width) {
							self.savedCanvasWidth = canvas.width;
							self.savedCanvasHeight = canvas.height;
						} else {
							const canvasFallback = document.querySelector('canvas');
							if (canvasFallback) {
								self.savedCanvasWidth = canvasFallback.width;
								self.savedCanvasHeight = canvasFallback.height;
							}
						}

						const onFrame = (time, frame) => {
							try {
								if (self.xrSession) {
									self.xrAnimationFrameHandle = self.xrSession.requestAnimationFrame(onFrame);
								}
								
								if (!self._lastFrameTime) self._lastFrameTime = time;
								const deltaTime = time - self._lastFrameTime;
								self._lastFrameTime = time;
								
								if (!window._xrFrameCount) window._xrFrameCount = 0;
								window._xrFrameCount++;

								if (glContext && self.xrSession && self.xrSession.renderState && self.xrSession.renderState.baseLayer) {
									const layer = self.xrSession.renderState.baseLayer;
									if (layer.framebuffer) {
										const pose = frame.getViewerPose(self.xrRefSpace);
										if (pose && pose.views && pose.views.length > 0) {
											glContext.bindFramebuffer(glContext.FRAMEBUFFER, layer.framebuffer);
											
											let bgR = 0, bgG = 0, bgB = 0;
											if (typeof iron !== 'undefined' && iron.Scene && iron.Scene.active && iron.Scene.active.world && iron.Scene.active.world.raw) {
												const bgColor = iron.Scene.active.world.raw.background_color;
												if (bgColor !== undefined) {
													bgR = ((bgColor >> 16) & 255) / 255;
													bgG = ((bgColor >> 8) & 255) / 255;
													bgB = (bgColor & 255) / 255;
												}
											}
											
											for (const view of pose.views) {
												const vp = layer.getViewport(view);
												glContext.viewport(vp.x, vp.y, vp.width, vp.height);
												glContext.scissor(vp.x, vp.y, vp.width, vp.height);
												glContext.enable(glContext.SCISSOR_TEST);
												
												glContext.clearColor(bgR, bgG, bgB, 1.0);
												glContext.clear(glContext.COLOR_BUFFER_BIT | glContext.DEPTH_BUFFER_BIT);

											}
											glContext.disable(glContext.SCISSOR_TEST);
										}
									}
								}
								
								if (!self.xrSession) {
									return;
								}
								
								const pose = frame.getViewerPose(self.xrRefSpace);
								if (!pose) {
									return;
								}
								
								if (pose.emulatedPosition && !self._emulatedPosLogged) {
									self._emulatedPosLogged = true;
								}
								
								const views = pose.views;
								
								if (!self.xrSession.renderState || !self.xrSession.renderState.baseLayer) {
									if (!self._noRenderStateLogged) {
										self._noRenderStateLogged = true;
									}
									return;
								}
								const glLayer = self.xrSession.renderState.baseLayer;
								
								if (!views || views.length === 0) {
									return;
								}
								
								if (self.xrSession.visibilityState === 'hidden') {
									return;
								}
								
								self.currentFrame = frame;
								self.currentViews = views;
								self.currentViewerPose = pose; /
								if (self.xrSession && self.xrSession.inputSources) {
									self.currentInputSources = self.xrSession.inputSources;
								}
								
								if (glContext.isContextLost()) {
									return;
								}
								
								if (!glContext || !glLayer || !glLayer.framebuffer) {
									return;
								}
								
								if (glContext.bindVertexArray) {
									glContext.bindVertexArray(null);
								}
								
								while (glContext.getError() !== glContext.NO_ERROR) {
									// Drain error queue
								}
								
								glContext.bindFramebuffer(glContext.FRAMEBUFFER, glLayer.framebuffer);
								

								const bindError = glContext.getError();
								if (bindError !== glContext.NO_ERROR && !self._bindErrorLogged) {
									self._bindErrorLogged = true;
								}
								
								const fbStatus = glContext.checkFramebufferStatus(glContext.FRAMEBUFFER);
								if (fbStatus !== glContext.FRAMEBUFFER_COMPLETE) {
									return;
								}
								
								glContext.enable(glContext.DEPTH_TEST);
								glContext.depthFunc(glContext.LEQUAL);
								glContext.depthMask(true);
								glContext.colorMask(true, true, true, true);
								glContext.disable(glContext.BLEND);
								glContext.enable(glContext.CULL_FACE);
								glContext.cullFace(glContext.BACK);
								glContext.frontFace(glContext.CCW);
								glContext.disable(glContext.STENCIL_TEST);
								glContext.disable(glContext.POLYGON_OFFSET_FILL);
								

								if (!self._fbLogged && p) {
									const depthTest = glContext.isEnabled(glContext.DEPTH_TEST);
									const cullFace = glContext.isEnabled(glContext.CULL_FACE);
									const blend = glContext.isEnabled(glContext.BLEND);
									self._fbLogged = true;
								}
								
								if (views.length === 0) {
									return;
								}
								
								if (views.length >= 1) {
									try {
										self._leftViewport = glLayer.getViewport(views[0]);
										self._rightViewport = views.length >= 2 ? glLayer.getViewport(views[1]) : null;
										self._cachedViewsLength = views.length;
										
										if (!self._leftViewport) {
											return;
										}
									} catch (e) {
										return;
									}
								}
								
								if (views.length >= 1) {
									self.leftProjectionMatrix = self.createMatrixFromArray(views[0].projectionMatrix);
									self.leftViewMatrix = self.createMatrixFromArray(views[0].transform.inverse.matrix);
								}
								if (views.length >= 2) {
									self.rightProjectionMatrix = self.createMatrixFromArray(views[1].projectionMatrix);
									self.rightViewMatrix = self.createMatrixFromArray(views[1].transform.inverse.matrix);
								} else if (views.length === 1) {
									self.rightProjectionMatrix = self.leftProjectionMatrix;
									self.rightViewMatrix = self.leftViewMatrix;
								}

								if (self._vrRenderCallback) {
									self._vrRenderCallback();
								}

							} catch (err) {
								console.error('XR Frame Error:', err);
							} finally {
								self.currentFrame = null;
								self.currentViews = null;
								self.currentInputSources = null;
							}
						};
						
						self.xrAnimationFrameHandle = session.requestAnimationFrame(onFrame);
					}).catch((err) => {
						trace('REF SPACE FAILED: ' + err.message );
					});
				}).catch((err) => {
					trace('SESSION FAILED: ' + err);
				});
			", vrScaleFactor);
		}
		catch (err:Dynamic) {
			trace("Failed to requestSession (WebXR).");
			trace(err);
		}
	}
	
	function onSessionEnd() {
		var canvas = SystemImpl.khanvas;
		if (canvas == null) {
			canvas = Syntax.code("document.querySelector('canvas')");
		}
		
		if (canvas != null && savedCanvasWidth > 0 && savedCanvasHeight > 0) {
			canvas.width = savedCanvasWidth;
			canvas.height = savedCanvasHeight;
		}
		
		if (xrSession != null) {
			Syntax.code("
				if (this.xrAnimationFrameHandle !== -1 && this.xrSession) {
					this.xrSession.cancelAnimationFrame(this.xrAnimationFrameHandle);
					this.xrAnimationFrameHandle = -1;
				}
				
				if (this._eventHandlers && this.xrSession) {
					const handlers = this._eventHandlers;
					const events = ['end', 'select', 'selectstart', 'selectend', 'squeeze', 'squeezestart', 'squeezeend', 'inputsourceschange', 'visibilitychange'];
					for (const event of events) {
						if (handlers[event]) {
							this.xrSession.removeEventListener(event, handlers[event]);
						}
					}
					this._eventHandlers = null;
				}
			");
		}
		
		Syntax.code("
			const gl = this._glContext;
			if (gl) {
				gl.bindFramebuffer(gl.FRAMEBUFFER, null); // Restore default framebuffer
				gl.disable(gl.SCISSOR_TEST);
				gl.viewport(0, 0, gl.canvas.width, gl.canvas.height);
			}

			if (typeof window !== 'undefined') {
				window._khaSkipWindowRender = false;
			}
			
			this.xrSession = null;
			this.xrRefSpace = null;
			this.xrGLLayer = null;
			this.currentFrame = null;
			this.currentViews = null;
			this.currentInputSources = null;
			this._lastFrameTime = null;
			
			if (typeof window !== 'undefined') {
				delete window._xrFrameCount;
				delete window._ironRenderCount;
				delete window._slowFrameCount;
			}
		");
	}

	public override function onVRExitPresent() {
		if (isWebXR) {
			try {
				if (xrSession != null) {
					Syntax.code("
						if (this.xrSession) {
							this.xrSession.end().then(() => {
								trace('Session ended');
								}).catch((err) => {
								trace('Session.end() failed:', err);
							});
						}
					");
					xrSession = null;
					xrRefSpace = null;
					xrGLLayer = null;
				}
			}
			catch (err:Dynamic) {
				trace("Failed to exitPresent in WebXR");
				trace(err);
			}
		}
		else {
			// WebVR 1.1
			try {
				vrDisplay.exitPresent([{source: SystemImpl.khanvas}]).then(function() {
					onResize();
				});
			}
			catch (err:Dynamic) {
				trace("Failed to exitPresent");
				trace(err);
			}
		}
	}

	public override function onResetPose() {
		try {
			vrDisplay.resetPose();
		}
		catch (err:Dynamic) {
			trace("Failed to resetPose");
			trace(err);
		}
	}

	function onAnimationFrame(timestamp: Float): Void {
		if (vrDisplay != null && vrDisplay.isPresenting) {
			vrDisplay.requestAnimationFrame(onAnimationFrame);

			vrDisplay.getFrameData(frameData);

			leftProjectionMatrix = createMatrixFromArray(untyped frameData.leftProjectionMatrix);
			leftViewMatrix = createMatrixFromArray(untyped frameData.leftViewMatrix);

			rightProjectionMatrix = createMatrixFromArray(untyped frameData.rightProjectionMatrix);
			rightViewMatrix = createMatrixFromArray(untyped frameData.rightViewMatrix);

			// Submit the newly rendered layer to be presented by the VRDisplay
			vrDisplay.submitFrame();
		}
	}

	function onResize() {
		if (isWebXR) {
			return;
		}
		else {
			// WebVR 1.1
			if (vrDisplay != null && vrDisplay.isPresenting) {
				var canvas = SystemImpl.khanvas;
				if (canvas != null) {
					canvas.width = vrWidth;
					canvas.height = vrHeight;
				}
			}
			else {
				var canvas = SystemImpl.khanvas;
				if (canvas != null) {
					canvas.width = width;
					canvas.height = height;
				}
			}
		}
	}

	public override function GetSensorState(): SensorState {
		return GetPredictedSensorState(0.0);
	}

	public override function GetPredictedSensorState(time: Float): SensorState {
		var result: SensorState = new SensorState();

		result.Predicted = new PoseState();
		result.Recorded = result.Predicted;

		result.Predicted.AngularAcceleration = new Vector3();
		result.Predicted.AngularVelocity = new Vector3();
		result.Predicted.LinearAcceleration = new Vector3();
		result.Predicted.LinearVelocity = new Vector3();
		result.Predicted.TimeInSeconds = time;
		result.Predicted.Pose = new Pose();
		result.Predicted.Pose.Orientation = new Quaternion();
		result.Predicted.Pose.Position = new Vector3();

		var mPose = frameData.pose; // predicted pose of the vrDisplay
		if (mPose != null) {
			result.Predicted.AngularVelocity = createVectorFromArray(untyped mPose.angularVelocity);
			result.Predicted.AngularAcceleration = createVectorFromArray(untyped mPose.angularAcceleration);
			result.Predicted.LinearVelocity = createVectorFromArray(untyped mPose.linearVelocity);
			result.Predicted.LinearAcceleration = createVectorFromArray(untyped mPose.linearAcceleration);
			result.Predicted.Pose.Orientation = createQuaternion(untyped mPose.orientation);
			result.Predicted.Pose.Position = createVectorFromArray(untyped mPose.position);
		}

		return result;
	}

	// Sends a black image to the warp swap thread
	public override function WarpSwapBlack(): Void {
		// TODO: Implement
	}

	// Sends the Oculus loading symbol to the warp swap thread
	public override function WarpSwapLoadingIcon(): Void {
		// TODO: Implement
	}

	// Sends the set of images to the warp swap thread
	public override function WarpSwap(parms: TimeWarpParms): Void {
		// TODO: Implement
	}

	public override function IsPresenting(): Bool {
		var presenting = false;
		if (vrDisplay != null)
			presenting = vrDisplay.isPresenting;
		}
		return presenting;
	}

	public override function IsVrEnabled(): Bool {
		return vrEnabled;
	}

	public override function GetTimeInSeconds(): Float {
		return Scheduler.time();
	}

	public override function GetProjectionMatrix(eye: Int): FastMatrix4 {
		if (eye == 0) {
			return leftProjectionMatrix;
		}
		else {
			return rightProjectionMatrix;
		}
	}

	public override function GetViewMatrix(eye: Int): FastMatrix4 {
		if (eye == 0) {
			return leftViewMatrix;
		}
		else {
			return rightViewMatrix;
		}
	}

	function createMatrixFromArray(array: Float32Array): FastMatrix4 {
		var matrix: FastMatrix4 = FastMatrix4.identity();
		if (array == null || array.length < 16) {
			trace("Warning: Invalid matrix array, using identity");
			return matrix;
		}
		matrix._00 = array[0];
		matrix._01 = array[1];
		matrix._02 = array[2];
		matrix._03 = array[3];
		matrix._10 = array[4];
		matrix._11 = array[5];
		matrix._12 = array[6];
		matrix._13 = array[7];
		matrix._20 = array[8];
		matrix._21 = array[9];
		matrix._22 = array[10];
		matrix._23 = array[11];
		matrix._30 = array[12];
		matrix._31 = array[13];
		matrix._32 = array[14];
		matrix._33 = array[15];
		return matrix;
	}

	function createVectorFromArray(array: Float32Array): Vector3 {
		var vector: Vector3 = new Vector3(0, 0, 0);
		if (array != null) {
			vector.x = array[0];
			vector.y = array[1];
			vector.z = array[2];
		}
		return vector;
	}

	function createQuaternion(array: Float32Array): Quaternion {
		var quaternion: Quaternion = new Quaternion(0, 0, 0, 0);
		if (array != null) {
			quaternion.x = array[0];
			quaternion.y = array[1];
			quaternion.z = array[2];
			quaternion.w = array[3];
		}
		return quaternion;
	}
	
	public function vrRenderCallback(): Void {
		var g4 = kha.SystemImpl.frame != null ? kha.SystemImpl.frame.g4 : null;

		
		if (g4 != null && iron.Scene.active != null && iron.RenderPath.active != null) {

			if (untyped window._vrUpdateStarted == null) {
				untyped window._vrUpdateStarted = true;
			}
			
			iron.system.Time.update();

			iron.Scene.active.updateFrame();

			
			js.Syntax.code("
				const App = iron.App;
				if (App) {
					const frame = window._vrCallbackCount;
					const inits = App.traitInits;
					if (inits && inits.length > 0) {
						for (let i = 0; i < inits.length; i++) {
							inits[i]();
						}
						inits.length = 0;
					}
					
					const fixedUpdates = App.traitFixedUpdates;
					if (fixedUpdates) {
						for (let i = 0; i < fixedUpdates.length; i++) {
							fixedUpdates[i]();
						}
					}
					
					const updates = App.traitUpdates;
					if (updates) {
						for (let i = 0; i < updates.length; i++) {
							updates[i]();
						}
					}
					
					const lateUpdates = App.traitLateUpdates;
					if (lateUpdates) {
						for (let i = 0; i < lateUpdates.length; i++) {
							lateUpdates[i]();
						}
					}

				}
			");
			iron.Scene.active.renderFrame(g4);
		}
		else {
			if (untyped window._vrSkipLogged == null) {
				untyped window._vrSkipLogged = true;
			}
		}
	}
}