// Copyright 2026 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 { WorkletEvents, registerMainThreadObjectType } from '@lynx-js/react/worklet-runtime/bindings';
import type { Worklet, WorkletRefImpl } from '@lynx-js/react/worklet-runtime/bindings';

import { useMemo } from './hooks/react.js';
import { getMainThreadObjectHandleType, registerMainThreadObjectHandle } from './main-thread-object-handle-registry.js';
import { allocateMainThreadRefId } from './main-thread-ref-id.js';
import { addMainThreadRefInitValue } from './main-thread-ref-init-value.js';

const mainThreadObjectTypeDefinitions = new WeakMap<
  object,
  MainThreadObjectTypeDefinition<unknown, object>
>();

/**
 * Describes how a serializable initialization payload becomes a stable object
 * in the main-thread runtime.
 *
 * @public
 */
export interface MainThreadObjectTypeDefinition<I, O extends object> {
  /** A globally unique and stable type key. */
  readonly type: string;
  /**
   * Create the realized main-thread object.
   *
   * This must be a capture-free Main Thread Function. Imports used by the
   * function must come from shared-runtime modules.
   */
  readonly create: (initialValue: I) => O;
}

/**
 * An immutable MainThreadObject type token with source-handle inspection.
 *
 * The factory function belongs to the input definition and is not exposed
 * by the returned token. In the background runtime they compile to opaque
 * Main Thread Function descriptors rather than bundling their implementation.
 *
 * @public
 */
export interface MainThreadObjectType<I, O extends object> {
  /** A globally unique and stable type key. */
  readonly type: string;
  /**
   * Narrow a value to a source-runtime handle of this exact type.
   *
   * This does not create or access the realized main-thread object.
   */
  readonly downcast: (value: unknown) => MainThreadObjectHandle<I, O> | undefined;
}

/**
 * An opaque source-runtime handle for an object realized on the main thread.
 *
 * A handle has logical identity and is serializable, but it is not the target
 * object or an RPC stub. Capturing it in a main-thread function resolves it to
 * the corresponding local target object. Development builds may wrap the
 * background handle in a throwing guard solely to diagnose invalid access;
 * that guard never forwards operations between runtimes.
 *
 * @public
 */
export abstract class MainThreadObjectHandle<I, O extends object> {
  /** @internal */
  declare private readonly _mainThreadObjectTargetType: O;
  /** @internal */
  protected _wvid: number;
  /** @internal */
  protected _initValue: I;
  /** @internal */
  protected _type: string;
  /** @internal */
  protected _releaseObserver?: unknown;

  /** @internal */
  protected constructor(initialValue: I, type: string) {
    if (__DEV__ && __JS__) {
      assertSerializableMainThreadObjectPayload(initialValue, type);
    }

    this._wvid = allocateMainThreadRefId();
    this._initValue = initialValue;
    this._type = type;
    registerMainThreadObjectHandle(this, type);

    if (__JS__) {
      addMainThreadRefInitValue(
        this._wvid,
        initialValue,
        type,
      );

      const id = this._wvid;
      this._releaseObserver = lynx.getNativeApp().createJSObjectDestructionObserver?.(() => {
        lynx.getCoreContext?.().dispatchEvent({
          type: WorkletEvents.releaseWorkletRef,
          data: { id },
        });
      });
    }
  }

  /**
   * The payload used to create the realized object.
   *
   * Readonly is a TypeScript contract; the value is not frozen at runtime.
   */
  public get creationPayload(): Readonly<I> {
    return this._initValue;
  }

  /** @internal */
  toJSON(): Pick<WorkletRefImpl<unknown>, '_wvid' | '_initValue' | '_type'> {
    return {
      _wvid: this._wvid,
      _initValue: this._initValue,
      _type: this._type,
    };
  }
}

class MainThreadObjectHandleImpl<I, O extends object> extends MainThreadObjectHandle<I, O> {
  constructor(initialValue: I, type: string) {
    super(initialValue, type);
  }
}

/**
 * Define a main-thread object type for use by a library-provided hook.
 *
 * @param definition - Stable type key and target-object factory function.
 * @returns An immutable object type definition.
 * @public
 */
export function defineMainThreadObjectType<I, O extends object>(
  definition: MainThreadObjectTypeDefinition<I, O>,
): MainThreadObjectType<I, O> {
  // Read accessor-backed fields once so validation and deferred registration
  // operate on the same definition values.
  const { type, create } = definition;
  if (typeof type !== 'string' || type.length === 0) {
    throw new Error('MainThreadObject type must be a non-empty string.');
  }
  if (type === 'main-thread') {
    throw new Error(
      'MainThreadObject type "main-thread" is reserved for MainThreadRef.',
    );
  }
  if (!isMainThreadLifecycleFunction(create)) {
    throw new Error(
      `MainThreadObject type "${type}" must provide a create Main Thread Function.`,
    );
  }
  if (__DEV__) {
    assertCaptureFreeLifecycleFunction(type, 'create', create);
  }

  // Keep the validated fields independent from the caller-owned definition.
  // Registration may be deferred until the first main-thread hook use, so a
  // later mutation of `definition` must not change the type or factory that is
  // registered for this token.
  const validatedDefinition = Object.freeze({
    type,
    create,
  });
  const objectType = Object.freeze({
    type,
    downcast(value: unknown): MainThreadObjectHandle<I, O> | undefined {
      return getMainThreadObjectHandleType(value) === type
        ? value as MainThreadObjectHandle<I, O>
        : undefined;
    },
  });
  mainThreadObjectTypeDefinitions.set(
    objectType,
    validatedDefinition as unknown as MainThreadObjectTypeDefinition<unknown, object>,
  );
  registerMainThreadObjectDefinition(validatedDefinition);
  return objectType;
}

/**
 * Create a stable handle that resolves to an object of the supplied type in
 * main-thread functions.
 *
 * This primitive is intended for library hooks. The returned target type is
 * only callable in a main-thread function; it must not be used from the
 * background runtime.
 *
 * @param objectType - Definition used to realize the target object.
 * @param initialValue - JSON-serializable initialization payload.
 * @returns The main-thread-only target type.
 * @public
 */
export function useMainThreadObject<I, O extends object>(
  objectType: MainThreadObjectType<I, O>,
  initialValue: I,
): O {
  return useMemo(() => {
    const definition = mainThreadObjectTypeDefinitions.get(objectType);
    if (definition === undefined) {
      throw new Error(
        `Invalid MainThreadObject type token for "${objectType.type}". Create it with defineMainThreadObjectType().`,
      );
    }
    // A library module normally registers its type while the MTS bundle is
    // evaluated. Register again at the first hook use so runtimes that retain
    // the module but reset their per-page worklet registry remain correct.
    // The registry treats an equivalent duplicate as an idempotent lookup.
    registerMainThreadObjectDefinition(definition);
    const handle = new MainThreadObjectHandleImpl<I, O>(initialValue, objectType.type);
    return guardBackgroundMainThreadObjectAccess(handle, objectType.type) as unknown as O;
  }, []);
}

function registerMainThreadObjectDefinition<I, O extends object>(
  definition: MainThreadObjectTypeDefinition<I, O>,
): void {
  if (__JS__) {
    return;
  }

  registerMainThreadObjectType(
    definition.type,
    definition.create as ((initialValue: unknown) => object) | Worklet,
  );
}

function isMainThreadLifecycleFunction(
  value: unknown,
): value is ((...args: never[]) => unknown) | Worklet {
  return typeof value === 'function'
    || (typeof value === 'object' && value !== null
      && typeof (value as Partial<Worklet>)._wkltId === 'string');
}

function assertCaptureFreeLifecycleFunction(
  type: string,
  name: 'create',
  value: ((...args: never[]) => unknown) | Worklet,
): void {
  if (typeof value === 'function') {
    return;
  }
  const descriptors = Object.getOwnPropertyDescriptors(value);
  const hasCaptures = Reflect.ownKeys(descriptors).some(key => {
    if (typeof key !== 'string') {
      return true;
    }
    if (key === '_wkltId' || key === '_workletType' || key === '_lepusWorkletHash') {
      return false;
    }
    if (key === '_c') {
      const descriptor = descriptors[key]!;
      const closure = 'value' in descriptor ? descriptor.value : undefined;
      return !('value' in descriptor)
        || (closure !== undefined
          && (typeof closure !== 'object' || closure === null
            || Reflect.ownKeys(closure).length !== 0));
    }
    return true;
  });
  if (hasCaptures) {
    throw new Error(
      `MainThreadObject ${name} function for "${type}" must not capture values. Import dependencies from a shared-runtime module instead.`,
    );
  }
}

function guardBackgroundMainThreadObjectAccess<I, O extends object>(
  handle: MainThreadObjectHandle<I, O>,
  type: string,
): MainThreadObjectHandle<I, O> {
  if (!__DEV__ || !__JS__) {
    return handle;
  }

  // This development-only guard never forwards a call to the main thread. It
  // only turns an otherwise opaque "method is not a function" failure into a
  // diagnostic. Production handles remain plain serializable objects.
  const guardedHandle = new Proxy(handle, {
    get(target, property, receiver): unknown {
      if (property in target) {
        return Reflect.get(target, property, receiver) as unknown;
      }
      if (
        typeof property === 'symbol'
        || property === 'then'
        || property === '$$typeof'
      ) {
        return undefined;
      }
      throw new Error(
        `MainThreadObject handle for "${type}" cannot access "${
          String(property)
        }" in the background runtime. Use the object only inside a main-thread function.`,
      );
    },
    set(target, property, value, receiver) {
      if (property in target) {
        return Reflect.set(target, property, value, receiver);
      }
      throw new Error(
        `MainThreadObject handle for "${type}" cannot set "${
          String(property)
        }" in the background runtime. Use the object only inside a main-thread function.`,
      );
    },
  });
  registerMainThreadObjectHandle(guardedHandle, type);
  return guardedHandle;
}

function assertSerializableMainThreadObjectPayload(value: unknown, type: string): void {
  const invalidPath = findNonSerializablePath(value, '$', new Set<object>());
  if (invalidPath !== undefined) {
    throw new Error(
      `MainThreadObject initial value for "${type}" must be JSON-serializable; invalid value at ${invalidPath}.`,
    );
  }
}

function findNonSerializablePath(
  value: unknown,
  path: string,
  ancestors: Set<object>,
): string | undefined {
  if (
    value === undefined
    || typeof value === 'function'
    || typeof value === 'symbol'
    || typeof value === 'bigint'
    || (typeof value === 'number' && !Number.isFinite(value))
  ) {
    return path;
  }
  if (value === null || typeof value !== 'object') {
    return undefined;
  }
  if (ancestors.has(value)) {
    return path;
  }

  const prototype = Object.getPrototypeOf(value) as object | null;
  if (!Array.isArray(value) && prototype !== Object.prototype && prototype !== null) {
    return path;
  }

  ancestors.add(value);
  if (Array.isArray(value)) {
    for (let index = 0; index < value.length; index++) {
      if (!Object.prototype.hasOwnProperty.call(value, index)) {
        continue;
      }
      const invalidPath = findNonSerializablePath(
        value[index],
        `${path}.${index}`,
        ancestors,
      );
      if (invalidPath !== undefined) {
        ancestors.delete(value);
        return invalidPath;
      }
    }
  } else {
    for (const [key, item] of Object.entries(value)) {
      const invalidPath = findNonSerializablePath(item, `${path}.${key}`, ancestors);
      if (invalidPath !== undefined) {
        ancestors.delete(value);
        return invalidPath;
      }
    }
  }
  ancestors.delete(value);
  return undefined;
}