import type { NotionDataSourceId } from "../../bridge/ids.js"
import type { QueryDataSourceMessage } from "../../bridge/messages/queryDataSource.js"
import type {
	CustomBlockHostQueryDataSourceResult,
	QuerySubscriptionState,
} from "./types.js"

export type QueryRefreshResult =
	| {
			status: "success"
			requestId: string
			snapshotId: string
			response: CustomBlockHostQueryDataSourceResult
	  }
	| { status: "missing" }

export class QuerySubscriptions {
	private readonly latestByDataSourceId = new Map<
		NotionDataSourceId,
		QuerySubscriptionState
	>()

	track(message: QueryDataSourceMessage) {
		this.latestByDataSourceId.set(message.dataSourceId, {
			requestId: message.requestId,
			snapshotId: message.snapshotId,
			limit: message.limit,
		})
	}

	resolveRefresh(args: {
		dataSourceId: NotionDataSourceId
		response: CustomBlockHostQueryDataSourceResult
	}): QueryRefreshResult {
		const queryState = this.latestByDataSourceId.get(args.dataSourceId)
		if (!queryState) {
			return { status: "missing" }
		}
		const items = args.response.items.slice(0, queryState.limit)
		return {
			status: "success",
			requestId: queryState.requestId,
			snapshotId: queryState.snapshotId,
			response: {
				items,
				hasMore:
					args.response.hasMore ?? args.response.items.length > items.length,
				error: args.response.error,
			},
		}
	}
}