import { parentPort, MessagePort } from 'node:worker_threads';
import type { Transferable } from 'node:worker_threads';
import { registry } from './registry.ts';
import { deserializeArg, serializeArg } from './shared/reconstruct.ts';
import { collectTransferables } from './transfer.ts';
import { pipeIterable, CANCEL, DEFAULT_HIGH_WATER, deserializeStreamHandle, isSerializedStreamHandle } from './pipe.ts';
import type { StreamHandle, SerializedStreamHandle } from './pipe.ts';

const imported = new Set<string>();
const taskCache = new Map<number, unknown>();

type Fn = (...args: unknown[]) => unknown;
const fnCache = new Map<string, Fn>();

function isTaskArg(arg: unknown): arg is { __task__: number; id: string; args: unknown[] } {
  return typeof arg === 'object' && arg !== null && '__task__' in arg;
}

function needsAsyncResolve(arg: unknown): boolean {
  return arg instanceof MessagePort || isSerializedStreamHandle(arg) || isTaskArg(arg);
}

// Returns the function synchronously when cached — callers must not await
// unconditionally, or they pay a microtask hop for every dispatch.
function resolveFn(id: string): Fn | Promise<Fn> {
  const cached = fnCache.get(id);
  if (cached) return cached;
  return resolveFnSlow(id);
}

async function resolveFnSlow(id: string): Promise<Fn> {
  const url = id.slice(0, id.lastIndexOf('#'));
  if (!imported.has(url)) {
    await import(url);
    imported.add(url);
  }
  const fn = registry.get(id) as Fn | undefined;
  if (!fn) throw new Error(`Moroutine not found: ${id}`);
  fnCache.set(id, fn);
  return fn;
}

function portToAsyncIterable<T>(handle: StreamHandle): AsyncIterable<T> {
  const { port, highWater, readySignal } = handle;
  const inflight = handle.flags.fields.inflight;
  const state = handle.flags.fields.state;

  const queue: T[] = [];
  let done = false;
  let error: Error | null = null;
  let waiting: (() => void) | null = null;

  port.on('message', (msg: { value?: unknown; done?: boolean; error?: Error }) => {
    if (msg.error) {
      error = msg.error;
      done = true;
      if (waiting) {
        waiting();
        waiting = null;
      }
      return;
    }
    if (msg.done) {
      done = true;
      if (waiting) {
        waiting();
        waiting = null;
      }
      return;
    }
    queue.push(deserializeArg(msg.value) as T);
    if (waiting) {
      waiting();
      waiting = null;
    }
  });

  return {
    [Symbol.asyncIterator]() {
      return {
        async next(): Promise<IteratorResult<T>> {
          while (true) {
            if (queue.length > 0) {
              const value = queue.shift()!;
              // Atomics-based backpressure: decrement inflight. In stream
              // mode the producer parks on inflight, so we notify it; in
              // channel mode the producer parks on readySignal, and we
              // only wake it on the cap→below-cap transition.
              const prev = inflight.sub(1);
              if (readySignal) {
                if (prev === highWater) {
                  readySignal.add(1);
                  readySignal.notify();
                }
              } else {
                inflight.notify();
              }
              return { done: false, value };
            }
            if (error) throw error;
            if (done) return { done: true, value: undefined };
            await new Promise<void>((resolve) => {
              waiting = resolve;
            });
          }
        },
        async return(): Promise<IteratorResult<T>> {
          state.store(CANCEL);
          inflight.notify();
          if (readySignal) {
            readySignal.add(1);
            readySignal.notify();
          }
          port.close();
          return { done: true, value: undefined };
        },
      };
    },
  };
}

async function resolveArg(arg: unknown): Promise<unknown> {
  if (isSerializedStreamHandle(arg)) {
    return portToAsyncIterable(deserializeStreamHandle(arg));
  }
  if (isTaskArg(arg)) {
    if (taskCache.has(arg.__task__)) {
      return taskCache.get(arg.__task__);
    }
    const inputArgs = arg.args;
    const resolvedArgs = new Array(inputArgs.length);
    for (let i = 0; i < inputArgs.length; i++) {
      const a = inputArgs[i];
      resolvedArgs[i] = needsAsyncResolve(a) ? await resolveArg(a) : deserializeArg(a);
    }
    const fnM = resolveFn(arg.id);
    const fn = fnM instanceof Promise ? await fnM : fnM;
    const raw = fn(...resolvedArgs);
    const value = raw instanceof Promise ? await raw : raw;
    taskCache.set(arg.__task__, value);
    return value;
  }
  return deserializeArg(arg);
}

type TaskMsg = {
  callId?: number;
  id: string;
  args: unknown[];
  /** Present on streaming dispatches; bundles port + flags + highWater. */
  stream?: SerializedStreamHandle;
};

function postResult(callId: number, value: unknown): void {
  const returnValue = serializeArg(value);
  const transferList: Transferable[] = [];
  collectTransferables(value, transferList);
  parentPort!.postMessage({ callId, value: returnValue }, transferList);
}

function postError(callId: number, err: unknown): void {
  parentPort!.postMessage({
    callId,
    error: err instanceof Error ? err : new Error(String(err)),
  });
}

// Value-task path, structured to pay zero microtask hops when the whole
// chain (fn resolution, arg resolution, function invocation) is sync.
function handleValueTask(msg: TaskMsg): void {
  const callId = msg.callId!;
  try {
    const fnM = resolveFn(msg.id);
    if (fnM instanceof Promise) {
      fnM.then(
        (fn) => invokeWithArgs(callId, fn, msg.args),
        (err) => postError(callId, err),
      );
      return;
    }
    invokeWithArgs(callId, fnM, msg.args);
  } catch (err) {
    postError(callId, err);
  }
}

function invokeWithArgs(callId: number, fn: Fn, args: unknown[]): void {
  try {
    // Sync prologue: walk args in-place until we hit one that needs async
    // resolution. If we finish the loop synchronously, go straight to invoke.
    const resolved = new Array(args.length);
    let i = 0;
    for (; i < args.length; i++) {
      const arg = args[i];
      if (needsAsyncResolve(arg)) break;
      resolved[i] = deserializeArg(arg);
    }
    if (i === args.length) {
      invokeAndRespond(callId, fn, resolved);
      return;
    }
    // Async tail: same loop body, awaiting where needed.
    (async () => {
      for (; i < args.length; i++) {
        const arg = args[i];
        resolved[i] = needsAsyncResolve(arg) ? await resolveArg(arg) : deserializeArg(arg);
      }
      invokeAndRespond(callId, fn, resolved);
    })().catch((err) => postError(callId, err));
  } catch (err) {
    postError(callId, err);
  }
}

function invokeAndRespond(callId: number, fn: Fn, args: unknown[]): void {
  try {
    const raw = fn(...args);
    if (raw instanceof Promise) {
      raw.then(
        (v) => postResult(callId, v),
        (err) => postError(callId, err),
      );
      return;
    }
    postResult(callId, raw);
  } catch (err) {
    postError(callId, err);
  }
}

async function handleStreamTask(msg: TaskMsg): Promise<void> {
  const { id, args } = msg;
  const callId = msg.callId;
  try {
    const fnM = resolveFn(id);
    const fn = fnM instanceof Promise ? await fnM : fnM;
    const resolvedArgs = new Array(args.length);
    for (let i = 0; i < args.length; i++) {
      const a = args[i];
      resolvedArgs[i] = needsAsyncResolve(a) ? await resolveArg(a) : deserializeArg(a);
    }
    const gen = fn(...resolvedArgs) as AsyncGenerator;
    await pipeIterable(gen, deserializeStreamHandle(msg.stream!));
  } catch (err) {
    if (callId != null) postError(callId, err);
  }
}

parentPort!.on('message', (msg: TaskMsg) => {
  if (msg.stream) {
    void handleStreamTask(msg);
  } else {
    handleValueTask(msg);
  }
});