import {
  getTRPCErrorFromUnknown,
  getTRPCErrorShape,
  isTrackedEnvelope,
} from '@trpc/server';
import { behaviorSubject, observable } from '@trpc/server/observable';
import { TRPC_ERROR_CODES_BY_KEY, type TRPCResult } from '@trpc/server/rpc';
import {
  callProcedure,
  isAbortError,
  isAsyncIterable,
  iteratorResource,
  makeResource,
  retryableRpcCodes,
  run,
  type AnyRouter,
  type ErrorHandlerOptions,
  type inferClientTypes,
  type inferRouterContext,
} from '@trpc/server/unstable-core-do-not-import';
import { inputWithTrackedEventId } from '../internals/inputWithTrackedEventId';
import { abortSignalToPromise, raceAbortSignals } from '../internals/signals';
import { getTransformer } from '../internals/transformer';
import type { TransformerOptions } from '../internals/transformer';
import { isTRPCClientError, TRPCClientError } from '../TRPCClientError';
import type { TRPCConnectionState } from './internals/subscriptions';
import type { TRPCLink } from './types';

export type LocalLinkOptions<TRouter extends AnyRouter> = {
  router: TRouter;
  createContext: () => Promise<inferRouterContext<TRouter>>;
  onError?: (opts: ErrorHandlerOptions<inferRouterContext<TRouter>>) => void;
} & TransformerOptions<inferClientTypes<TRouter>>;

/**
 * localLink is a terminating link that allows you to make tRPC procedure calls directly in your application without going through HTTP.
 *
 * @see https://trpc.io/docs/client/links/localLink
 */
export function unstable_localLink<TRouter extends AnyRouter>(
  opts: LocalLinkOptions<TRouter>,
): TRPCLink<TRouter> {
  const transformer = getTransformer(opts.transformer);

  const transformChunk = (chunk: unknown) => {
    if (opts.transformer) {
      // assume transformer will do the right thing
      return chunk;
    }
    // Special case for undefined, because `JSON.stringify(undefined)` throws
    if (chunk === undefined) {
      return chunk;
    }
    const serialized = JSON.stringify(transformer.input.serialize(chunk));
    const deserialized = JSON.parse(transformer.output.deserialize(serialized));
    return deserialized;
  };

  return () =>
    ({ op }) =>
      observable((observer) => {
        let ctx: inferRouterContext<TRouter> | undefined = undefined;
        const ac = new AbortController();

        const signal = raceAbortSignals(op.signal, ac.signal);
        const signalPromise = abortSignalToPromise(signal);

        signalPromise.catch(() => {
          // prevent unhandled rejection
        });

        let input = op.input;
        async function runProcedure(newInput: unknown): Promise<unknown> {
          input = newInput;

          ctx = await opts.createContext();

          return callProcedure({
            router: opts.router,
            path: op.path,
            getRawInput: async () => newInput,
            ctx,
            type: op.type,
            signal,
            batchIndex: 0,
          });
        }

        function onErrorCallback(cause: unknown) {
          if (isAbortError(cause)) {
            return;
          }
          opts.onError?.({
            error: getTRPCErrorFromUnknown(cause),
            type: op.type,
            path: op.path,
            input,
            ctx,
          });
        }

        function coerceToTRPCClientError(cause: unknown) {
          if (isTRPCClientError<TRouter>(cause)) {
            return cause;
          }
          const error = getTRPCErrorFromUnknown(cause);

          const shape = getTRPCErrorShape({
            config: opts.router._def._config,
            ctx,
            error,
            input,
            path: op.path,
            type: op.type,
          });
          return TRPCClientError.from(
            { error: transformChunk(shape) },
            { cause: cause instanceof Error ? cause : undefined },
          );
        }

        run(async () => {
          switch (op.type) {
            case 'query':
            case 'mutation': {
              const result = await runProcedure(op.input);
              if (!isAsyncIterable(result)) {
                observer.next({
                  result: { data: transformChunk(result) },
                });
                observer.complete();
                break;
              }

              observer.next({
                result: {
                  data: (async function* () {
                    await using iterator = iteratorResource(result);
                    using _finally = makeResource({}, () => {
                      observer.complete();
                    });
                    try {
                      while (true) {
                        const res = await Promise.race([
                          iterator.next(),
                          signalPromise,
                        ]);
                        if (res.done) {
                          return transformChunk(res.value);
                        }
                        yield transformChunk(res.value);
                      }
                    } catch (cause) {
                      onErrorCallback(cause);
                      throw coerceToTRPCClientError(cause);
                    }
                  })(),
                },
              });
              break;
            }
            case 'subscription': {
              const connectionState = behaviorSubject<
                TRPCConnectionState<TRPCClientError<any>>
              >({
                type: 'state',
                state: 'connecting',
                error: null,
              });

              const connectionSub = connectionState.subscribe({
                next(state) {
                  observer.next({
                    result: state,
                  });
                },
              });
              let lastEventId: string | undefined = undefined;

              using _finally = makeResource({}, async () => {
                connectionState.next({
                  type: 'state',
                  state: 'idle',
                  error: null,
                });
                connectionSub.unsubscribe();
              });

              // returning from here means the stream ended, so this completes
              // below - a throw is left to the `.catch()`, which needs the
              // observer to still be open by the time it runs
              await run(async () => {
                while (true) {
                  const result = await runProcedure(
                    inputWithTrackedEventId(op.input, lastEventId),
                  );
                  if (!isAsyncIterable(result)) {
                    throw new Error('Expected an async iterable');
                  }
                  await using iterator = iteratorResource(result);

                  observer.next({
                    result: {
                      type: 'started',
                    },
                  });
                  connectionState.next({
                    type: 'state',
                    state: 'pending',
                    error: null,
                  });

                  // Use a while loop to handle errors and reconnects
                  while (true) {
                    let res;
                    try {
                      res = await Promise.race([
                        iterator.next(),
                        signalPromise,
                      ]);
                    } catch (cause) {
                      if (isAbortError(cause)) {
                        return;
                      }
                      const error = getTRPCErrorFromUnknown(cause);

                      if (
                        !retryableRpcCodes.includes(
                          TRPC_ERROR_CODES_BY_KEY[error.code],
                        )
                      ) {
                        throw coerceToTRPCClientError(error);
                      }

                      onErrorCallback(error);
                      connectionState.next({
                        type: 'state',
                        state: 'connecting',
                        error: coerceToTRPCClientError(error),
                      });

                      break;
                    }

                    if (res.done) {
                      return;
                    }
                    let chunk: TRPCResult<unknown>;
                    if (isTrackedEnvelope(res.value)) {
                      lastEventId = res.value[0];

                      chunk = {
                        id: res.value[0],
                        data: {
                          id: res.value[0],
                          data: res.value[1],
                        },
                      };
                    } else {
                      chunk = {
                        data: res.value,
                      };
                    }

                    observer.next({
                      result: {
                        ...chunk,
                        data: transformChunk(chunk.data),
                      },
                    });
                  }
                }
              });

              observer.complete();
              break;
            }
          }
        }).catch((cause) => {
          onErrorCallback(cause);
          observer.error(coerceToTRPCClientError(cause));
        });

        return () => {
          ac.abort();
        };
      });
}
/**
 * @deprecated Renamed to `unstable_localLink`. This alias will be removed in a future major release.
 */
export const experimental_localLink: typeof unstable_localLink =
  unstable_localLink;