// Copyright 2024 The Lynx Authors. All rights reserved.
// Licensed under the Apache License Version 2.0 that can be found in the
// LICENSE file in the root directory of this source tree.

import type {
  ClosureValueType,
  JsFnHandle,
  Worklet,
  WorkletRef,
  WorkletRefId,
  WorkletRefImpl,
} from './bindings/index.js';
import { profile } from './utils/profile.js';
import { hydrateWorkletValue, isHydratedWorkletValue } from './workletRef.js';

/**
 * Hydrates a Worklet context with data from a first-screen Worklet context.
 * This process is typically used to run all delayed `runOnBackground` functions
 * and initialize `WorkletRef` values modified before hydration.
 *
 * @param ctx The target Worklet context to be hydrated.
 * @param firstScreenCtx The Worklet context from the first screen rendering,
 *                       containing the data to hydrate with.
 */
export function hydrateCtx(ctx: Worklet, firstScreenCtx: Worklet): void {
  profile('hydrateCtx', () => {
    hydrateCtxImpl(ctx, firstScreenCtx, ctx._execId!);
  });
}

function hydrateCtxImpl(
  ctx: ClosureValueType,
  firstScreenCtx: ClosureValueType,
  execId: number,
): void {
  if (
    !ctx || typeof ctx !== 'object' || !firstScreenCtx
    || typeof firstScreenCtx !== 'object'
  ) return;

  const ctxObj = ctx as Record<string, ClosureValueType>;
  const firstScreenCtxObj = firstScreenCtx as Record<string, ClosureValueType>;

  if (ctxObj['_wkltId'] && ctxObj['_wkltId'] !== firstScreenCtxObj['_wkltId']) {
    return;
  }

  // Typed handles (including `_type: 'main-thread'` mutable cells) hydrate
  // atomically without walking their initialization payload. Untyped legacy
  // refs continue through the existing `_wvid` path below.
  if (isTypedWorkletRefDescriptor(ctxObj)) {
    hydrateWorkletValueHandle(ctxObj, firstScreenCtxObj);
    return;
  }
  for (const key in ctx) {
    if (key === '_wvid') {
      hydrateMainThreadRef(
        ctxObj[key] as WorkletRefId,
        firstScreenCtxObj as unknown as WorkletRefImpl<unknown>,
      );
    } else if (key === '_jsFn') {
      hydrateDelayRunOnBackgroundTasks(
        ctxObj[key] as Record<string, JsFnHandle>,
        firstScreenCtxObj[key] as Record<string, JsFnHandle>,
        execId,
      );
    } else {
      const firstScreenValue = typeof firstScreenCtxObj[key] === 'function'
        ? (firstScreenCtxObj[key] as { boundCtx?: object }).boundCtx as ClosureValueType
        : firstScreenCtxObj[key];
      hydrateCtxImpl(ctxObj[key], firstScreenValue, execId);
    }
  }
}

function isTypedWorkletRefDescriptor(
  value: object,
): value is WorkletRefImpl<unknown> {
  const descriptor = value as Partial<WorkletRefImpl<unknown>>;
  return typeof descriptor._wvid === 'number'
    && typeof descriptor._type === 'string';
}

/**
 * Hydrates a WorkletRef on the main thread.
 * This maps the positive-ID background handle to the target realized during
 * first-screen rendering, while leaving an unused first-screen handle alone.
 */
function hydrateMainThreadRef(
  refId: WorkletRefId,
  value: WorkletRefImpl<unknown>,
): void {
  if ('_initValue' in value) {
    return;
  }
  lynxWorkletImpl!._refImpl._workletRefMap[refId] = value as WorkletRef<unknown>;
}

/**
 * Hydrates a typed worklet handle on the main thread.
 * The target is the object or mutable cell realized from the matching first-screen handle.
 */
function hydrateWorkletValueHandle(
  handle: WorkletRefImpl<unknown>,
  value: ClosureValueType,
): void {
  if (!isHydratedWorkletValue(value)) {
    // The handle was not accessed during first-screen rendering.
    return;
  }
  hydrateWorkletValue(handle, value);
}

/**
 * Hydrates delayed `runOnBackground` tasks.
 * This function ensures that any `runOnBackground` calls that were delayed
 * during the first-screen rendering are correctly associated with their
 * respective JavaScript function handles in the hydrated Worklet context.
 */
function hydrateDelayRunOnBackgroundTasks(
  fnObjs: Record<string, JsFnHandle>, // example: {"_jsFn1":{"_jsFnId":1}}
  firstScreenFnObjs: Record<string, JsFnHandle>, // example: {"_jsFn1":{"_isFirstScreen":true,"_delayIndices":[0]}}
  execId: number,
) {
  for (const fnName in fnObjs) {
    const fnObj = fnObjs[fnName]!;
    const firstScreenFnObj: JsFnHandle | undefined = firstScreenFnObjs[fnName];
    if (!firstScreenFnObj?._delayIndices) {
      if (firstScreenFnObj) {
        firstScreenFnObj._isFirstScreen = false;
        firstScreenFnObj._execId = execId;
        Object.assign(firstScreenFnObj, fnObj);
      }
      continue;
    }
    for (const index of firstScreenFnObj._delayIndices) {
      const details = lynxWorkletImpl!._runOnBackgroundDelayImpl
        .delayedBackgroundFunctionArray[index];
      if (!details) {
        continue;
      }
      fnObj._execId = execId;
      details.jsFnHandle = fnObj;
    }
  }
}