import type { WritableFamilyToken, WritableToken } from "atom.io"
import type { Canonical } from "atom.io/foundations/canonical"
import { setIntoStore } from "atom.io/internal"
import { useCallback, useContext } from "react"

import { parseStateOverloads } from "./parse-state-overloads.ts"
import { StoreContext } from "./store-context.tsx"

export function useI<T>(
	token: WritableToken<T, any, any>,
): <New extends T>(next: New | ((old: T) => New)) => void

export function useI<T, K extends Canonical>(
	token: WritableFamilyToken<T, K, any>,
	key: NoInfer<K>,
): <New extends T>(next: New | ((old: T) => New)) => void

export function useI<T, K extends Canonical>(
	...params:
		| [WritableFamilyToken<T, K, any>, NoInfer<K>]
		| [WritableToken<T, any, any>]
): <New extends T>(next: New | ((old: T) => New)) => void {
	const store = useContext(StoreContext)
	const token = parseStateOverloads(store, ...params)
	const setter = useCallback(
		(next: T | ((old: T) => T)) => {
			setIntoStore(store, token, next)
		},
		[store, token.key],
	)
	return setter
}