import AsyncStorage from '@react-native-async-storage/async-storage';
import { AppState } from 'react-native';
import type { StateStorage } from 'zustand/middleware';

export interface DebouncedStorage extends StateStorage {
    /** Writes every pending value right away. */
    flush: () => Promise<void>;
}

/**
 * Wraps a storage so bursts of writes collapse into one trailing write.
 *
 * The chat store is persisted on every state change and a streamed answer
 * changes it once per token, which would serialize every saved session and hit
 * AsyncStorage dozens of times per second.
 */
export function createDebouncedStorage(base: StateStorage, delayMs = 400): DebouncedStorage {
    const pending = new Map<string, string>();
    let timer: ReturnType<typeof setTimeout> | null = null;

    const flush = async () => {
        if (timer) {
            clearTimeout(timer);
            timer = null;
        }
        const entries = [...pending.entries()];
        pending.clear();
        await Promise.all(entries.map(([name, value]) => base.setItem(name, value)));
    };

    return {
        getItem: (name) => {
            const value = pending.get(name);
            return value === undefined ? base.getItem(name) : Promise.resolve(value);
        },
        setItem: (name, value) => {
            pending.set(name, value);
            if (timer) clearTimeout(timer);
            timer = setTimeout(flush, delayMs);
        },
        removeItem: (name) => {
            pending.delete(name);
            return base.removeItem(name);
        },
        flush,
    };
}

export const appStorage = createDebouncedStorage(AsyncStorage);

// Do not lose the last tokens when the app is sent to the background or closed.
AppState.addEventListener('change', (state) => {
    if (state !== 'active') void appStorage.flush();
});