import type { DataTag, QueryClient, QueryFilters } from '@tanstack/react-query';
import type {
  TRPCClient,
  TRPCClientErrorLike,
  TRPCRequestOptions,
} from '@trpc/client';
import { getUntypedClient, TRPCUntypedClient } from '@trpc/client';
import type {
  AnyTRPCProcedure,
  AnyTRPCRootTypes,
  AnyTRPCRouter,
  inferProcedureErrorShape,
  inferProcedureInput,
  inferRouterContext,
  inferTransformedProcedureOutput,
  TRPCProcedureType,
  TRPCRouterRecord,
} from '@trpc/server';
import { callTRPCProcedure, createTRPCRecursiveProxy } from '@trpc/server';
import type { MaybePromise } from '@trpc/server/unstable-core-do-not-import';
import type {
  DecorateHandleReaders,
  DecorateInfiniteQueryHandleVerbs,
  DecorateQueryHandleVerbs,
  HandleVerb,
} from './handles';
import { prepareHandleRead, readHandle, startHandle } from './handles';
import {
  trpcInfiniteQueryOptions,
  type TRPCInfiniteQueryOptions,
} from './infiniteQueryOptions';
import type { MutationOptionsOverride } from './mutationOptions';
import {
  trpcMutationOptions,
  type TRPCMutationOptions,
} from './mutationOptions';
import { trpcQueryOptions, type TRPCQueryOptions } from './queryOptions';
import {
  trpcSubscriptionOptions,
  type TRPCSubscriptionOptions,
} from './subscriptionOptions';
import type {
  DefaultFeatureFlags,
  FeatureFlags,
  KeyPrefixOptions,
  OptionalCursorInput,
  ResolverDef,
  TRPCInfiniteData,
  TRPCMutationKey,
  TRPCQueryKey,
  WithRequired,
} from './types';
import {
  getMutationKeyInternal,
  getQueryKeyInternal,
  unwrapLazyArg,
} from './utils';

/**
 * Key and filter helpers for a router path.
 */
export interface DecorateRouterKeyable<TFeatureFlags extends FeatureFlags> {
  /**
   * Calculate the TanStack Query Key for any path, could be used to invalidate every procedure beneath this path
   *
   * @see https://tanstack.com/query/latest/docs/framework/react/guides/query-keys
   * @see https://trpc.io/docs/client/tanstack-react-query/usage#queryKey
   */
  pathKey: () => TRPCQueryKey<TFeatureFlags['keyPrefix']>;

  /**
   * Calculate a TanStack Query Filter for any path, could be used to manipulate every procedure beneath this path
   *
   * @see https://tanstack.com/query/latest/docs/framework/react/guides/filters
   * @see https://trpc.io/docs/client/tanstack-react-query/usage#queryFilter
   */
  pathFilter: (
    filters?: QueryFilters<TRPCQueryKey<TFeatureFlags['keyPrefix']>>,
  ) => WithRequired<
    QueryFilters<TRPCQueryKey<TFeatureFlags['keyPrefix']>>,
    'queryKey'
  >;
}

interface TypeHelper<TDef extends ResolverDef> {
  /**
   * @internal prefer using inferInput and inferOutput to access types
   */
  '~types': {
    input: TDef['input'];
    output: TDef['output'];
    errorShape: TDef['errorShape'];
  };
}

/**
 * The input type of a procedure on the options proxy.
 * @see https://trpc.io/docs/client/tanstack-react-query/usage#inferring-input-and-output-types
 */
export type inferInput<
  TProcedure extends
    | DecorateInfiniteQueryProcedure<any>
    | DecorateQueryProcedure<any>
    | DecorateMutationProcedure<any>,
> = TProcedure['~types']['input'];

/**
 * The output type of a procedure on the options proxy.
 * @see https://trpc.io/docs/client/tanstack-react-query/usage#inferring-input-and-output-types
 */
export type inferOutput<
  TProcedure extends
    | DecorateInfiniteQueryProcedure<any>
    | DecorateQueryProcedure<any>
    | DecorateMutationProcedure<any>,
> = TProcedure['~types']['output'];

/**
 * The helpers for a query procedure whose input has a `cursor`.
 * @see https://trpc.io/docs/client/tanstack-react-query/usage#infiniteQueryOptions
 */
export interface DecorateInfiniteQueryProcedure<TDef extends ResolverDef>
  extends TypeHelper<TDef>, DecorateInfiniteQueryHandleVerbs<TDef> {
  /**
   * Create a set of type-safe infinite query options that can be passed to `useInfiniteQuery`, `prefetchInfiniteQuery` etc.
   *
   * @see https://tanstack.com/query/latest/docs/framework/react/reference/infiniteQueryOptions#infinitequeryoptions
   * @see https://trpc.io/docs/client/tanstack-react-query/usage#infiniteQueryOptions
   */
  infiniteQueryOptions: TRPCInfiniteQueryOptions<TDef>;

  /**
   * Calculate the TanStack Query Key for a Infinite Query Procedure
   *
   * @see https://tanstack.com/query/latest/docs/framework/react/guides/query-keys
   * @see https://trpc.io/docs/client/tanstack-react-query/usage#queryKey
   */
  infiniteQueryKey: (input?: Partial<TDef['input']>) => DataTag<
    TRPCQueryKey<TDef['featureFlags']['keyPrefix']>,
    TRPCInfiniteData<TDef['input'], TDef['output']>,
    TRPCClientErrorLike<{
      transformer: TDef['transformer'];
      errorShape: TDef['errorShape'];
    }>
  >;

  /**
   * Calculate a TanStack Query Filter for a Infinite Query Procedure
   *
   * @see https://tanstack.com/query/latest/docs/framework/react/guides/filters
   * @see https://trpc.io/docs/client/tanstack-react-query/usage#queryFilter
   */
  infiniteQueryFilter: (
    input?: Partial<TDef['input']>,
    filters?: QueryFilters<
      DataTag<
        TRPCQueryKey<TDef['featureFlags']['keyPrefix']>,
        TRPCInfiniteData<TDef['input'], TDef['output']>,
        TRPCClientErrorLike<{
          transformer: TDef['transformer'];
          errorShape: TDef['errorShape'];
        }>
      >
    >,
  ) => WithRequired<
    QueryFilters<
      DataTag<
        TRPCQueryKey<TDef['featureFlags']['keyPrefix']>,
        TRPCInfiniteData<TDef['input'], TDef['output']>,
        TRPCClientErrorLike<{
          transformer: TDef['transformer'];
          errorShape: TDef['errorShape'];
        }>
      >
    >,
    'queryKey'
  >;
}

/**
 * The helpers for a query procedure.
 * @see https://trpc.io/docs/client/tanstack-react-query/usage#queryOptions
 */
export interface DecorateQueryProcedure<TDef extends ResolverDef>
  extends
    TypeHelper<TDef>,
    DecorateRouterKeyable<TDef['featureFlags']>,
    DecorateQueryHandleVerbs<TDef> {
  /**
   * Create a set of type-safe query options that can be passed to `useQuery`, `prefetchQuery` etc.
   *
   * @see https://tanstack.com/query/latest/docs/framework/react/reference/queryOptions#queryoptions
   * @see https://trpc.io/docs/client/tanstack-react-query/usage#queryOptions
   */
  queryOptions: TRPCQueryOptions<TDef>;

  /**
   * Calculate the TanStack Query Key for a Query Procedure
   *
   * @see https://tanstack.com/query/latest/docs/framework/react/guides/query-keys
   * @see https://trpc.io/docs/client/tanstack-react-query/usage#queryKey
   */
  queryKey: (input?: Partial<TDef['input']>) => DataTag<
    TRPCQueryKey<TDef['featureFlags']['keyPrefix']>,
    TDef['output'],
    TRPCClientErrorLike<{
      transformer: TDef['transformer'];
      errorShape: TDef['errorShape'];
    }>
  >;

  /**
   * Calculate a TanStack Query Filter for a Query Procedure
   *
   * @see https://tanstack.com/query/latest/docs/framework/react/guides/filters
   * @see https://trpc.io/docs/client/tanstack-react-query/usage#queryFilter
   */
  queryFilter: (
    input?: Partial<TDef['input']>,
    filters?: QueryFilters<
      DataTag<
        TRPCQueryKey<TDef['featureFlags']['keyPrefix']>,
        TDef['output'],
        TRPCClientErrorLike<{
          transformer: TDef['transformer'];
          errorShape: TDef['errorShape'];
        }>
      >
    >,
  ) => WithRequired<
    QueryFilters<
      DataTag<
        TRPCQueryKey<TDef['featureFlags']['keyPrefix']>,
        TDef['output'],
        TRPCClientErrorLike<{
          transformer: TDef['transformer'];
          errorShape: TDef['errorShape'];
        }>
      >
    >,
    'queryKey'
  >;
}

/**
 * The helpers for a mutation procedure.
 * @see https://trpc.io/docs/client/tanstack-react-query/usage#mutationOptions
 */
export interface DecorateMutationProcedure<
  TDef extends ResolverDef,
> extends TypeHelper<TDef> {
  /**
   * Create a set of type-safe mutation options that can be passed to `useMutation`
   *
   * @see https://trpc.io/docs/client/tanstack-react-query/usage#mutationOptions
   */
  mutationOptions: TRPCMutationOptions<TDef>;

  /**
   * Calculate the TanStack Mutation Key for a Mutation Procedure
   *
   * @see https://trpc.io/docs/client/tanstack-react-query/usage#mutationKey
   */
  mutationKey: () => TRPCMutationKey<TDef['featureFlags']['keyPrefix']>;
}

/**
 * The helpers for a subscription procedure.
 * @see https://trpc.io/docs/client/tanstack-react-query/usage#subscriptionOptions
 */
export interface DecorateSubscriptionProcedure<
  TDef extends ResolverDef,
> extends TypeHelper<TDef> {
  /**
   * Create a set of type-safe subscription options that can be passed to `useSubscription`
   *
   * @see https://trpc.io/docs/client/tanstack-react-query/usage#subscriptionOptions
   */
  subscriptionOptions: TRPCSubscriptionOptions<TDef>;
}

/**
 * The helpers for a procedure of the given type.
 */
export type DecorateProcedure<
  TType extends TRPCProcedureType,
  TDef extends ResolverDef,
> = TType extends 'query'
  ? DecorateQueryProcedure<TDef> &
      (TDef['input'] extends OptionalCursorInput
        ? DecorateInfiniteQueryProcedure<TDef>
        : Record<string, never>)
  : TType extends 'mutation'
    ? DecorateMutationProcedure<TDef>
    : TType extends 'subscription'
      ? DecorateSubscriptionProcedure<TDef>
      : never;

/**
 * @internal
 */
export type DecoratedRouterRecord<
  TRoot extends AnyTRPCRootTypes,
  TRecord extends TRPCRouterRecord,
  TFeatureFlags extends FeatureFlags = DefaultFeatureFlags,
> = {
  [TKey in keyof TRecord]: TRecord[TKey] extends infer $Value
    ? $Value extends TRPCRouterRecord
      ? DecoratedRouterRecord<TRoot, $Value, TFeatureFlags> &
          DecorateRouterKeyable<TFeatureFlags> &
          DecorateHandleReaders<TFeatureFlags>
      : $Value extends AnyTRPCProcedure
        ? DecorateProcedure<
            $Value['_def']['type'],
            {
              input: inferProcedureInput<$Value>;
              output: inferTransformedProcedureOutput<TRoot, $Value>;
              transformer: TRoot['transformer'];
              errorShape: inferProcedureErrorShape<TRoot, $Value>;
              featureFlags: TFeatureFlags;
            }
          >
        : never
    : never;
};

/**
 * The object returned by `useTRPC()` and `createTRPCOptionsProxy`. It mirrors your router and has query, mutation and subscription option helpers on each procedure.
 * @see https://trpc.io/docs/client/tanstack-react-query/usage
 */
export type TRPCOptionsProxy<
  TRouter extends AnyTRPCRouter,
  TFeatureFlags extends FeatureFlags = DefaultFeatureFlags,
> = DecoratedRouterRecord<
  TRouter['_def']['_config']['$types'],
  TRouter['_def']['record'],
  TFeatureFlags
> &
  DecorateRouterKeyable<TFeatureFlags> &
  DecorateHandleReaders<TFeatureFlags>;

/**
 * Options shared by every {@link createTRPCOptionsProxy} call.
 */
export type TRPCOptionsProxyOptionsBase<
  TFeatureFlags extends FeatureFlags = DefaultFeatureFlags,
> = {
  /**
   * The `QueryClient`, or a function that returns it.
   */
  queryClient: QueryClient | (() => QueryClient);
  /**
   * @internal
   */
  overrides?: {
    mutations?: MutationOptionsOverride;
  };
} & KeyPrefixOptions<TFeatureFlags>;

/**
 * Call procedures on a router directly, for example on the server.
 */
export interface TRPCOptionsProxyOptionsInternal<
  TRouter extends AnyTRPCRouter,
> {
  /**
   * The router to call.
   */
  router: TRouter;
  /**
   * The context for each call, or a function that returns it.
   */
  ctx:
    | inferRouterContext<TRouter>
    | (() => MaybePromise<inferRouterContext<TRouter>>);
}

/**
 * Call procedures through a tRPC client.
 */
export interface TRPCOptionsProxyOptionsExternal<
  TRouter extends AnyTRPCRouter,
> {
  /**
   * The tRPC client to call procedures with.
   */
  client: TRPCUntypedClient<TRouter> | TRPCClient<TRouter>;
}

/**
 * Options for {@link createTRPCOptionsProxy}. Pass either `client` or `router` and `ctx`.
 */
export type TRPCOptionsProxyOptions<
  TRouter extends AnyTRPCRouter,
  TFeatureFlags extends FeatureFlags = DefaultFeatureFlags,
> = TRPCOptionsProxyOptionsBase<TFeatureFlags> &
  (
    | TRPCOptionsProxyOptionsInternal<TRouter>
    | TRPCOptionsProxyOptionsExternal<TRouter>
  );

/**
 * Marks the `useTRPC()` proxy, whose handle verbs run during render
 * @internal
 */
export const renderHandleVerbs = Symbol('renderHandleVerbs');

type UtilsMethods =
  | keyof DecorateQueryProcedure<any>
  | keyof DecorateHandleReaders<any>
  | keyof DecorateInfiniteQueryProcedure<any>
  | keyof DecorateMutationProcedure<any>
  | keyof DecorateSubscriptionProcedure<any>;

/**
 * Creates the options proxy without React context. Works on the server too.
 *
 * @example
 * ```ts
 * export const trpc = createTRPCOptionsProxy<AppRouter>({
 *   client: trpcClient,
 *   queryClient,
 * });
 * ```
 *
 * @see https://trpc.io/docs/client/tanstack-react-query/setup#3c-set-up-without-react-context
 * @see https://trpc.io/docs/client/tanstack-react-query/server-components#5-create-a-trpc-caller-for-server-components
 */
export function createTRPCOptionsProxy<
  TRouter extends AnyTRPCRouter,
  TFeatureFlags extends FeatureFlags = DefaultFeatureFlags,
>(
  opts: TRPCOptionsProxyOptions<TRouter, TFeatureFlags>,
): TRPCOptionsProxy<TRouter, TFeatureFlags> {
  const prefix = opts.keyPrefix;

  const callIt = (type: TRPCProcedureType): any => {
    return (path: string, input: unknown, trpcOpts: TRPCRequestOptions) => {
      if ('router' in opts) {
        return Promise.resolve(unwrapLazyArg(opts.ctx)).then((ctx) =>
          callTRPCProcedure({
            router: opts.router,
            path: path,
            getRawInput: async () => input,
            ctx: ctx,
            type: type,
            signal: undefined,
            batchIndex: 0,
          }),
        );
      }

      const untypedClient =
        opts.client instanceof TRPCUntypedClient
          ? opts.client
          : getUntypedClient(opts.client);

      return untypedClient[type](path, input, trpcOpts);
    };
  };

  return createTRPCRecursiveProxy(({ args, path: _path }) => {
    const path = [..._path];
    const utilName = path.pop() as UtilsMethods;
    const [arg1, arg2] = args as any[];

    const startQueryHandle = (kind: 'query' | 'infinite', verb: HandleVerb) => {
      const inRender =
        (opts as { [renderHandleVerbs]?: boolean })[renderHandleVerbs] === true;
      if (inRender && verb !== 'prefetch') {
        throw new Error(
          `\`${[...path, utilName].join('.')}()\` on \`useTRPC()\` waits for its query, which a render can't. ` +
            `Start queries in components with \`unstable_$prefetch\` or \`unstable_$prefetchInfinite\`, ` +
            `and wait where you read the handle, with \`useSuspenseQuery\`. ` +
            `See https://trpc.io/docs/client/tanstack-react-query/handles#in-render`,
        );
      }
      const options =
        kind === 'infinite'
          ? trpcInfiniteQueryOptions({
              input: arg1,
              opts: arg2,
              path,
              queryClient: opts.queryClient,
              queryKey: getQueryKeyInternal({
                path,
                input: arg1,
                type: 'infinite',
                prefix,
              }),
              query: callIt('query'),
            })
          : trpcQueryOptions({
              input: arg1,
              opts: arg2,
              path,
              queryClient: opts.queryClient,
              queryKey: getQueryKeyInternal({
                path,
                input: arg1,
                type: 'query',
                prefix,
              }),
              query: callIt('query'),
            });
      return startHandle({
        queryClient: unwrapLazyArg(opts.queryClient),
        options,
        kind,
        verb,
        inRender,
      });
    };

    const contextMap: Record<UtilsMethods, () => unknown> = {
      '~types': undefined as any,

      pathKey: () => {
        return getQueryKeyInternal({
          path,
          type: 'any',
          prefix,
        });
      },
      pathFilter: (): QueryFilters => {
        return {
          ...arg1,
          queryKey: getQueryKeyInternal({
            path,
            type: 'any',
            prefix,
          }),
        };
      },

      queryOptions: () => {
        return trpcQueryOptions({
          input: arg1,
          opts: arg2,
          path,
          queryClient: opts.queryClient,
          queryKey: getQueryKeyInternal({
            path,
            input: arg1,
            type: 'query',
            prefix,
          }),
          query: callIt('query'),
        });
      },
      queryKey: () => {
        return getQueryKeyInternal({
          path,
          input: arg1,
          type: 'query',
          prefix,
        });
      },
      queryFilter: (): QueryFilters => {
        return {
          ...arg2,
          queryKey: getQueryKeyInternal({
            path,
            input: arg1,
            type: 'query',
            prefix,
          }),
        };
      },

      infiniteQueryOptions: () => {
        return trpcInfiniteQueryOptions({
          input: arg1,
          opts: arg2,
          path,
          queryClient: opts.queryClient,
          queryKey: getQueryKeyInternal({
            path,
            input: arg1,
            type: 'infinite',
            prefix,
          }),
          query: callIt('query'),
        });
      },
      infiniteQueryKey: () => {
        return getQueryKeyInternal({
          path,
          input: arg1,
          type: 'infinite',
          prefix,
        });
      },
      infiniteQueryFilter: (): QueryFilters => {
        return {
          ...arg2,
          queryKey: getQueryKeyInternal({
            path,
            input: arg1,
            type: 'infinite',
            prefix,
          }),
        };
      },

      unstable_$fromHandle: () => {
        const handle = readHandle(arg1, prefix);
        const options =
          handle.type === 'infinite'
            ? trpcInfiniteQueryOptions({
                input: handle.input,
                opts: {
                  initialCursor: arg1.initialCursor,
                  ...arg2,
                },
                path: handle.path,
                queryClient: opts.queryClient,
                queryKey: handle.queryKey,
                query: callIt('query'),
              })
            : trpcQueryOptions({
                input: handle.input,
                opts: arg2,
                path: handle.path,
                queryClient: opts.queryClient,
                queryKey: handle.queryKey,
                query: callIt('query'),
              });
        return prepareHandleRead(
          unwrapLazyArg(opts.queryClient),
          handle,
          options,
        );
      },
      unstable_$prefetch: () => startQueryHandle('query', 'prefetch'),
      unstable_$ensure: () => startQueryHandle('query', 'ensure'),
      unstable_$fetch: () => startQueryHandle('query', 'fetch'),
      unstable_$prefetchInfinite: () =>
        startQueryHandle('infinite', 'prefetch'),
      unstable_$ensureInfinite: () => startQueryHandle('infinite', 'ensure'),
      unstable_$fetchInfinite: () => startQueryHandle('infinite', 'fetch'),

      mutationOptions: () => {
        return trpcMutationOptions({
          opts: arg1,
          path,
          prefix,
          queryClient: opts.queryClient,
          mutate: callIt('mutation'),
          overrides: opts.overrides?.mutations,
        });
      },
      mutationKey: () => {
        return getMutationKeyInternal({
          path,
          prefix,
        });
      },

      subscriptionOptions: () => {
        return trpcSubscriptionOptions({
          opts: arg2,
          path,
          input: arg1,
          queryKey: getQueryKeyInternal({
            path,
            input: arg1,
            type: 'any',
            prefix,
          }),
          subscribe: callIt('subscription'),
        });
      },
    };

    return contextMap[utilName]();
  });
}