frontend/src/shared/api/resource-query.ts

import type { QueryClient } from "@tanstack/react-query";
import {
	type QueryFunctionContext,
	type QueryKey,
	type UseQueryOptions,
	type UseQueryResult,
	useQueryClient,
} from "@tanstack/react-query";
import { useEffect } from "react";
import { useAuthenticatedQuery } from "./authenticated-query";
import { getAuthenticatedSessionSnapshot } from "./client";
import type { components } from "./generated/schema";
import { apiErrorCode } from "./problem-details";

const accessEpochs = new WeakMap<QueryClient, Map<string, number>>();

type Resource = {
	data: unknown;
	meta: components["schemas"]["InvestmentResourceMetaDto"];
};
function unwrapResourceResult<T extends Resource, E>(
	result: UseQueryResult<T, E>,
) {
	const missing = result.isSuccess && result.data.data === null;
	const data: T["data"] | undefined = result.data?.data;
	const meta: T["meta"] | undefined = result.data?.meta;
	const status: UseQueryResult<T, E>["status"] = missing
		? "pending"
		: result.status;
	// Project resource fields; refetch still returns the transport envelope.
	return {
		...result,
		transport: {
			status: result.status,
			fetchStatus: result.fetchStatus,
			error: result.error,
		},
		data: data ?? undefined,
		meta,
		isPending: result.isPending || missing,
		isSuccess: result.isSuccess && !missing,
		status,
	};
}

function invalidateResourceAccess(
	client: QueryClient,
	group: string,
	error: unknown,
) {
	if (
		![
			"NOT_FOUND",
			"FORBIDDEN",
			"TBANK_NOT_CONNECTED",
			"INVALID_TBANK_TOKEN",
		].includes(apiErrorCode(error) ?? "")
	)
		return;
	const epochs = accessEpochs.get(client) ?? new Map<string, number>();
	accessEpochs.set(client, epochs);
	epochs.set(group, (epochs.get(group) ?? 0) + 1);
	const affected = {
		predicate: (cached: { meta?: Record<string, unknown> }) =>
			cached.meta?.accessGroup === group,
	};
	void client.cancelQueries(affected);
	for (const cached of client.getQueryCache().findAll(affected))
		cached.setState({
			data: undefined,
			dataUpdatedAt: 0,
			status: "error",
			error: error as Error,
			fetchStatus: "idle",
		});
}

// Observation changes transport mode, never the generated cache identity.
function useResourceQuery<T extends Resource, E, K extends QueryKey>(
	options: UseQueryOptions<T, E, T, K>,
) {
	const client = useQueryClient();
	const group =
		typeof options.meta?.accessGroup === "string"
			? options.meta.accessGroup
			: undefined;
	let epochs = accessEpochs.get(client);
	if (!epochs) {
		epochs = new Map();
		accessEpochs.set(client, epochs);
	}
	const groupEpochs = epochs;
	const queryFn = options.queryFn;
	const guarded =
		typeof queryFn === "function"
			? async (context: QueryFunctionContext<K>) => {
					const epoch = group ? (groupEpochs.get(group) ?? 0) : 0;
					const result = await queryFn(context);
					if (group && epoch !== (groupEpochs.get(group) ?? 0))
						throw new DOMException("Access changed", "AbortError");
					return result;
				}
			: queryFn;
	const query = useAuthenticatedQuery({ ...options, queryFn: guarded });
	useEffect(() => {
		if (group) invalidateResourceAccess(client, group, query.error);
	}, [client, group, query.error]);
	const observe = async (expectedStateVersion?: string) => {
		const queryFn = guarded;
		if (
			!query.isEnabled ||
			!getAuthenticatedSessionSnapshot() ||
			typeof queryFn !== "function"
		)
			return;
		// Cancellation is synchronous; start the replacement in the same boundary.
		void client.cancelQueries({ queryKey: options.queryKey, exact: true });
		return client.fetchQuery({
			...options,
			staleTime: 0,
			retry: false,
			queryFn: async (context: QueryFunctionContext<K>) => {
				const [method, path, input] = context.queryKey;
				const init = input as
					| { params?: { query?: Record<string, unknown> } }
					| undefined;
				const queryKey = [
					method,
					path,
					{
						...init,
						params: {
							...init?.params,
							query: { ...init?.params?.query, sync: "observe" },
						},
					},
				] as unknown as K;
				const response = await queryFn({ ...context, queryKey });
				if (
					expectedStateVersion &&
					response.meta.stateVersion !== expectedStateVersion
				) {
					// One more observational read; never turn a stale notification into an import.
					const current = await queryFn({ ...context, queryKey });
					if (current.meta.stateVersion !== expectedStateVersion)
						throw new DOMException("Resource changed", "AbortError");
					return current;
				}
				return response;
			},
		});
	};
	return {
		...unwrapResourceResult(query),
		queryKey: options.queryKey,
		observe,
	};
}

export { invalidateResourceAccess, unwrapResourceResult, useResourceQuery };