--!strict
-- ROBLOX upstream: https://github.com/facebook/react/blob/16654436039dd8f16a63928e71081c7745872e8f/packages/react-reconciler/src/ReactFiberThrow.new.js
--[[*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *
 * @flow
 ]]

local Packages = script.Parent.Parent
local ReactGlobals = require(Packages.ReactGlobals)
local LuauPolyfill = require(Packages.LuauPolyfill)
local Object = LuauPolyfill.Object
-- ROBLOX: use patched console from shared
local console = require(Packages.Shared).console
type Error = LuauPolyfill.Error
type Map<K, V> = { [K]: V }
type Object = { [string]: any }
type Set<T> = { [T]: boolean }

local ReactInternalTypes = require(script.Parent.ReactInternalTypes)
type Fiber = ReactInternalTypes.Fiber
type FiberRoot = ReactInternalTypes.FiberRoot
type ReactPriorityLevel = ReactInternalTypes.ReactPriorityLevel
local ReactFiberLane = require(script.Parent.ReactFiberLane)
type Lanes = ReactFiberLane.Lanes
type Lane = ReactFiberLane.Lane
local ReactCapturedValue = require(script.Parent.ReactCapturedValue)
type CapturedValue<T> = ReactCapturedValue.CapturedValue<T>
local ReactUpdateQueue = require(script.Parent["ReactUpdateQueue.new"])
type Update<T> = ReactInternalTypes.Update<T>

local ReactTypes = require(Packages.Shared)
type React_Component<Props, State> = ReactTypes.React_Component<Props, State>
type Thenable<T> = ReactTypes.Thenable<T>
type Wakeable = ReactTypes.Wakeable

local ReactFiberSuspenseContext = require(script.Parent["ReactFiberSuspenseContext.new"])
type SuspenseContext = ReactFiberSuspenseContext.SuspenseContext

local getComponentName = require(Packages.Shared).getComponentName

local ReactWorkTags = require(script.Parent.ReactWorkTags)
local ClassComponent = ReactWorkTags.ClassComponent
local HostRoot = ReactWorkTags.HostRoot
local SuspenseComponent = ReactWorkTags.SuspenseComponent
local IncompleteClassComponent = ReactWorkTags.IncompleteClassComponent

local ReactFiberFlags = require(script.Parent.ReactFiberFlags)
local DidCapture = ReactFiberFlags.DidCapture
local Incomplete = ReactFiberFlags.Incomplete
local NoFlags = ReactFiberFlags.NoFlags
local ShouldCapture = ReactFiberFlags.ShouldCapture
local LifecycleEffectMask = ReactFiberFlags.LifecycleEffectMask
local ForceUpdateForLegacySuspense = ReactFiberFlags.ForceUpdateForLegacySuspense
local shouldCaptureSuspense =
	require(script.Parent["ReactFiberSuspenseComponent.new"]).shouldCaptureSuspense
local ReactTypeOfMode = require(script.Parent.ReactTypeOfMode)
local NoMode = ReactTypeOfMode.NoMode
local BlockingMode = ReactTypeOfMode.BlockingMode
local DebugTracingMode = ReactTypeOfMode.DebugTracingMode
local ReactFeatureFlags = require(Packages.Shared).ReactFeatureFlags
local enableDebugTracing = ReactFeatureFlags.enableDebugTracing
local enableSchedulingProfiler = ReactFeatureFlags.enableSchedulingProfiler
local createCapturedValue = require(script.Parent.ReactCapturedValue).createCapturedValue

local enqueueCapturedUpdate = ReactUpdateQueue.enqueueCapturedUpdate
local createUpdate = ReactUpdateQueue.createUpdate
local CaptureUpdate = ReactUpdateQueue.CaptureUpdate
local ForceUpdate = ReactUpdateQueue.ForceUpdate
local enqueueUpdate = ReactUpdateQueue.enqueueUpdate
local markFailedErrorBoundaryForHotReloading = require(
	script.Parent["ReactFiberHotReloading.new"]
).markFailedErrorBoundaryForHotReloading

local hasSuspenseContext = ReactFiberSuspenseContext.hasSuspenseContext
local InvisibleParentSuspenseContext =
	ReactFiberSuspenseContext.InvisibleParentSuspenseContext
local suspenseStackCursor = ReactFiberSuspenseContext.suspenseStackCursor

-- ROBLOX FIXME: these will incur a dependency cycle
-- onUncaughtError would be very easy to extract out, or to transplant into this file
local ReactFiberWorkLoop
local markLegacyErrorBoundaryAsFailedRef, isAlreadyFailedLegacyErrorBoundaryRef, pingSuspendedRootRef

-- ROBLOX deviation: lazy initialize ReactFiberWorkLoop to prevent cyclic module dependency
local markLegacyErrorBoundaryAsFailed = function(...)
	if not markLegacyErrorBoundaryAsFailedRef then
		ReactFiberWorkLoop = require(script.Parent["ReactFiberWorkLoop.new"])
		markLegacyErrorBoundaryAsFailedRef =
			ReactFiberWorkLoop.markLegacyErrorBoundaryAsFailed
	end
	return markLegacyErrorBoundaryAsFailedRef(...)
end
-- local {
--   renderDidError,
--   onUncaughtError,
--   markLegacyErrorBoundaryAsFailed,
--   isAlreadyFailedLegacyErrorBoundary,
--   pingSuspendedRoot,
-- } = require(Packages../ReactFiberWorkLoop.new'
local pingSuspendedRoot = function(...)
	if ReactFiberWorkLoop == nil then
		ReactFiberWorkLoop = require(script.Parent["ReactFiberWorkLoop.new"])
	end
	pingSuspendedRootRef = ReactFiberWorkLoop.pingSuspendedRoot
	return pingSuspendedRootRef(...)
end
local isAlreadyFailedLegacyErrorBoundary = function(...)
	if ReactFiberWorkLoop == nil then
		ReactFiberWorkLoop = require(script.Parent["ReactFiberWorkLoop.new"])
	end
	isAlreadyFailedLegacyErrorBoundaryRef =
		ReactFiberWorkLoop.isAlreadyFailedLegacyErrorBoundary
	return isAlreadyFailedLegacyErrorBoundaryRef(...)
end

local logCapturedError = require(script.Parent.ReactFiberErrorLogger).logCapturedError
local logComponentSuspended = require(script.Parent.DebugTracing).logComponentSuspended
local markComponentSuspended =
	require(script.Parent.SchedulingProfiler).markComponentSuspended

local SyncLane = ReactFiberLane.SyncLane
local NoTimestamp = ReactFiberLane.NoTimestamp
local includesSomeLane = ReactFiberLane.includesSomeLane
local mergeLanes = ReactFiberLane.mergeLanes
local pickArbitraryLane = ReactFiberLane.pickArbitraryLane

-- local PossiblyWeakMap = typeof WeakMap == 'function' ? WeakMap : Map

function createRootErrorUpdate(
	fiber: Fiber,
	errorInfo: CapturedValue<Error>,
	lane: Lane,
	-- ROBLOX deviation: parameterize method to avoid circular dependency
	onUncaughtError
): Update<any>
	local update = createUpdate(NoTimestamp, lane)
	-- Unmount the root by rendering nil.
	update.tag = CaptureUpdate
	-- Caution: React DevTools currently depends on this property
	-- being called "element".
	update.payload = { element = Object.None }
	local _error = errorInfo.value
	update.callback = function()
		if onUncaughtError ~= nil then
			onUncaughtError(_error)
		end
		logCapturedError(fiber, errorInfo)
	end
	return update
end

function createClassErrorUpdate(
	fiber: Fiber,
	errorInfo: CapturedValue<Error>,
	lane: Lane
): Update<any>
	local update = createUpdate(NoTimestamp, lane)
	update.tag = CaptureUpdate
	local getDerivedStateFromError = (fiber.type :: React_Component<any, any>).getDerivedStateFromError
	if typeof(getDerivedStateFromError) == "function" then
		local error_ = errorInfo.value
		update.payload = function()
			logCapturedError(fiber, errorInfo)
			return getDerivedStateFromError(error_)
		end
	end

	local inst = fiber.stateNode
	if inst ~= nil and typeof(inst.componentDidCatch) == "function" then
		update.callback = function()
			if ReactGlobals.__DEV__ then
				markFailedErrorBoundaryForHotReloading(fiber)
			end
			if typeof(getDerivedStateFromError) ~= "function" then
				-- To preserve the preexisting retry behavior of error boundaries,
				-- we keep track of which ones already failed during this batch.
				-- This gets reset before we yield back to the browser.
				-- TODO: Warn in strict mode if getDerivedStateFromError is
				-- not defined.
				-- ROBLOX FIXME: used to be `this` upstream, needs verification by ReactIncremental unwinding test
				markLegacyErrorBoundaryAsFailed(inst)

				-- Only log here if componentDidCatch is the only error boundary method defined
				logCapturedError(fiber, errorInfo)
			end
			local error_ = errorInfo.value
			local stack = errorInfo.stack
			-- ROBLOX FIXME: used to be `this` upstream, needs verification by ReactIncremental unwinding test
			inst:componentDidCatch(error_, {
				componentStack = stack or "",
			})
			if ReactGlobals.__DEV__ then
				if typeof(getDerivedStateFromError) ~= "function" then
					-- If componentDidCatch is the only error boundary method defined,
					-- then it needs to call setState to recover from errors.
					-- If no state update is scheduled then the boundary will swallow the error.
					if not includesSomeLane(fiber.lanes, SyncLane) then
						console.error(
							"%s: Error boundaries should implement getDerivedStateFromError(). "
								.. "In that method, return a state update to display an error message or fallback UI.",
							getComponentName(fiber.type) or "Unknown"
						)
					end
				end
			end
		end
	elseif ReactGlobals.__DEV__ then
		update.callback = function()
			markFailedErrorBoundaryForHotReloading(fiber)
		end
	end
	return update
end

local function attachPingListener(root: FiberRoot, wakeable: Wakeable, lanes: Lanes)
	-- Attach a listener to the promise to "ping" the root and retry. But only if
	-- one does not already exist for the lanes we're currently rendering (which
	-- acts like a "thread ID" here).
	local pingCache: Map<Wakeable, (Set<any> | Map<Wakeable, Set<any>>)> | nil =
		root.pingCache
	local threadIDs
	if pingCache == nil then
		-- ROBLOX deviation: use table in place of WeakMap
		-- ROBLOX performance: slight re-ordering so we initialize the table in one shot
		threadIDs = {} :: Set<any>
		root.pingCache = {
			[wakeable] = threadIDs,
		} :: Map<Wakeable, (Set<any> | Map<Wakeable, Set<any>>)>
		pingCache = root.pingCache :: Map<Wakeable, (Set<any> | Map<Wakeable, Set<any>>)>
	else
		threadIDs = (
			pingCache :: Map<Wakeable, (Set<any> | Map<Wakeable, Set<any>>)>
		)[wakeable] :: Set<any>
		if threadIDs == nil then
			threadIDs = {} :: Set<any>;
			(pingCache :: Map<Wakeable, (Set<any> | Map<Wakeable, Set<any>>)>)[wakeable] =
				threadIDs
		end
	end
	if not threadIDs[lanes] then
		-- Memoize using the thread ID to prevent redundant listeners.
		threadIDs[lanes] = true
		local ping = function()
			return pingSuspendedRoot(root, wakeable, lanes)
		end
		wakeable:andThen(ping, ping)
	end
end

function throwException(
	root: FiberRoot,
	returnFiber: Fiber,
	sourceFiber: Fiber,
	value: any,
	rootRenderLanes: Lanes,
	onUncaughtError,
	renderDidError
)
	-- The source fiber did not complete.
	sourceFiber.flags = bit32.bor(sourceFiber.flags, Incomplete)

	if
		value ~= nil
		and typeof(value) == "table"
		and typeof(value.andThen) == "function"
	then
		-- This is a wakeable.
		local wakeable: Wakeable = value

		if ReactGlobals.__DEV__ then
			if enableDebugTracing then
				if bit32.band(sourceFiber.mode, DebugTracingMode) ~= 0 then
					local name = getComponentName(sourceFiber.type) or "Unknown"
					logComponentSuspended(name, wakeable)
				end
			end
		end

		if enableSchedulingProfiler then
			markComponentSuspended(sourceFiber, wakeable)
		end

		if bit32.band(sourceFiber.mode, BlockingMode) == NoMode then
			-- Reset the memoizedState to what it was before we attempted
			-- to render it.
			local currentSource = sourceFiber.alternate
			if currentSource then
				-- ROBLOX performance TODO: return non-nil updateQueue object to the ReactUpdateQUeue pool
				sourceFiber.updateQueue = currentSource.updateQueue
				sourceFiber.memoizedState = currentSource.memoizedState
				sourceFiber.lanes = currentSource.lanes
			else
				-- ROBLOX performance TODO: return non-nil updateQueue object to the ReactUpdateQUeue pool
				sourceFiber.updateQueue = nil
				sourceFiber.memoizedState = nil
			end
		end

		local hasInvisibleParentBoundary = hasSuspenseContext(
			suspenseStackCursor.current,
			InvisibleParentSuspenseContext :: SuspenseContext
		)

		-- Schedule the nearest Suspense to re-render the timed out view.
		local workInProgress = returnFiber
		repeat
			if
				workInProgress.tag == SuspenseComponent
				and shouldCaptureSuspense(workInProgress, hasInvisibleParentBoundary)
			then
				-- Found the nearest boundary.

				-- Stash the promise on the boundary fiber. If the boundary times out, we'll
				-- attach another listener to flip the boundary back to its normal state.
				local wakeables: Set<Wakeable> = workInProgress.updateQueue
				if wakeables == nil then
					local updateQueue = {
						[wakeable] = true,
					}
					-- ROBLOX performance TODO: return non-nil updateQueue object to the ReactUpdateQUeue pool
					workInProgress.updateQueue = updateQueue
				else
					wakeables[wakeable] = true
				end

				-- If the boundary is outside of blocking mode, we should *not*
				-- suspend the commit. Pretend as if the suspended component rendered
				-- nil and keep rendering. In the commit phase, we'll schedule a
				-- subsequent synchronous update to re-render the Suspense.
				--
				-- Note: It doesn't matter whether the component that suspended was
				-- inside a blocking mode tree. If the Suspense is outside of it, we
				-- should *not* suspend the commit.
				if bit32.band(workInProgress.mode, BlockingMode) == NoMode then
					workInProgress.flags = bit32.bor(workInProgress.flags, DidCapture)
					sourceFiber.flags =
						bit32.bor(sourceFiber.flags, ForceUpdateForLegacySuspense)

					-- We're going to commit this fiber even though it didn't complete.
					-- But we shouldn't call any lifecycle methods or callbacks. Remove
					-- all lifecycle effect tags.
					sourceFiber.flags = bit32.band(
						sourceFiber.flags,
						bit32.bnot(bit32.bor(LifecycleEffectMask, Incomplete))
					)

					if sourceFiber.tag == ClassComponent then
						local currentSourceFiber = sourceFiber.alternate
						if currentSourceFiber == nil then
							-- This is a new mount. Change the tag so it's not mistaken for a
							-- completed class component. For example, we should not call
							-- componentWillUnmount if it is deleted.
							sourceFiber.tag = IncompleteClassComponent
						else
							-- When we try rendering again, we should not reuse the current fiber,
							-- since it's known to be in an inconsistent state. Use a force update to
							-- prevent a bail out.
							local update = createUpdate(NoTimestamp, SyncLane)
							update.tag = ForceUpdate
							enqueueUpdate(sourceFiber, update)
						end
					end

					-- The source fiber did not complete. Mark it with Sync priority to
					-- indicate that it still has pending work.
					sourceFiber.lanes = mergeLanes(sourceFiber.lanes, SyncLane)

					-- Exit without suspending.
					return
				end

				-- Confirmed that the boundary is in a concurrent mode tree. Continue
				-- with the normal suspend path.
				--
				-- After this we'll use a set of heuristics to determine whether this
				-- render pass will run to completion or restart or "suspend" the commit.
				-- The actual logic for this is spread out in different places.
				--
				-- This first principle is that if we're going to suspend when we complete
				-- a root, then we should also restart if we get an update or ping that
				-- might unsuspend it, and vice versa. The only reason to suspend is
				-- because you think you might want to restart before committing. However,
				-- it doesn't make sense to restart only while in the period we're suspended.
				--
				-- Restarting too aggressively is also not good because it starves out any
				-- intermediate loading state. So we use heuristics to determine when.

				-- Suspense Heuristics
				--
				-- If nothing threw a Promise or all the same fallbacks are already showing,
				-- then don't suspend/restart.
				--
				-- If this is an initial render of a new tree of Suspense boundaries and
				-- those trigger a fallback, then don't suspend/restart. We want to ensure
				-- that we can show the initial loading state as quickly as possible.
				--
				-- If we hit a "Delayed" case, such as when we'd switch from content back into
				-- a fallback, then we should always suspend/restart. Transitions apply
				-- to this case. If none is defined, JND is used instead.
				--
				-- If we're already showing a fallback and it gets "retried", allowing us to show
				-- another level, but there's still an inner boundary that would show a fallback,
				-- then we suspend/restart for 500ms since the last time we showed a fallback
				-- anywhere in the tree. This effectively throttles progressive loading into a
				-- consistent train of commits. This also gives us an opportunity to restart to
				-- get to the completed state slightly earlier.
				--
				-- If there's ambiguity due to batching it's resolved in preference of:
				-- 1) "delayed", 2) "initial render", 3) "retry".
				--
				-- We want to ensure that a "busy" state doesn't get force committed. We want to
				-- ensure that new initial loading states can commit as soon as possible.

				attachPingListener(root, wakeable, rootRenderLanes)

				workInProgress.flags = bit32.bor(workInProgress.flags, ShouldCapture)
				workInProgress.lanes = rootRenderLanes

				return
			end
			-- This boundary already captured during this render. Continue to the next
			-- boundary.
			workInProgress = workInProgress.return_ :: Fiber -- ROBLOX TODO: Luau narrowing doesn't understand this loop until nil pattern
		until workInProgress == nil

		-- No boundary was found. Fallthrough to error mode.
		-- TODO: Use invariant so the message is stripped in prod?
		value = (getComponentName(sourceFiber.type) or "A React component")
			.. " suspended while rendering, but no fallback UI was specified.\n"
			.. "\n"
			.. "Add a <Suspense fallback=...> component higher in the tree to "
			.. "provide a loading indicator or placeholder to display."
	end

	-- We didn't find a boundary that could handle this type of exception. Start
	-- over and traverse parent path again, this time treating the exception
	-- as an error.
	renderDidError()

	value = createCapturedValue(value, sourceFiber)
	local workInProgress = returnFiber
	repeat
		if workInProgress.tag == HostRoot then
			local errorInfo = value
			workInProgress.flags = bit32.bor(workInProgress.flags, ShouldCapture)
			local lane = pickArbitraryLane(rootRenderLanes)
			workInProgress.lanes = mergeLanes(workInProgress.lanes, lane)
			-- ROBLOX deviation: parameterize method onUncaughtError to avoid circular dependency
			local update =
				createRootErrorUpdate(workInProgress, errorInfo, lane, onUncaughtError)
			enqueueCapturedUpdate(workInProgress, update)
			return
		elseif workInProgress.tag == ClassComponent then
			-- Capture and retry
			local errorInfo = value
			local ctor = workInProgress.type
			local instance = workInProgress.stateNode
			if
				bit32.band(workInProgress.flags, DidCapture) == NoFlags
				and (
					typeof(ctor.getDerivedStateFromError) == "function"
					or (
						instance ~= nil
						and typeof(instance.componentDidCatch) == "function"
						and not isAlreadyFailedLegacyErrorBoundary(instance)
					)
				)
			then
				workInProgress.flags = bit32.bor(workInProgress.flags, ShouldCapture)
				local lane = pickArbitraryLane(rootRenderLanes)
				workInProgress.lanes = mergeLanes(workInProgress.lanes, lane)
				-- Schedule the error boundary to re-render using updated state
				local update = createClassErrorUpdate(workInProgress, errorInfo, lane)
				enqueueCapturedUpdate(workInProgress, update)
				return
			end
		end
		workInProgress = workInProgress.return_ :: Fiber -- ROBLOX TODO: Luau narrowing doesn't understand this loop until nil pattern
	until workInProgress == nil
end

return {
	throwException = throwException,
	createRootErrorUpdate = createRootErrorUpdate,
	createClassErrorUpdate = createClassErrorUpdate,
}
