import { describe, expect, it, vi } from 'vitest'
import { memoizedFn, refCountedMemoizedFn } from './memoized.ts'

describe('memoizedFn', () => {
	it('returns cached result for same args', () => {
		const fn = vi.fn((a: number, b: number) => a + b)
		const memo = memoizedFn('add', fn)

		expect(memo(1, 2)).toBe(3)
		expect(memo(1, 2)).toBe(3)
		expect(fn).toHaveBeenCalledTimes(1)
	})

	it('computes fresh result for different args', () => {
		const fn = vi.fn((a: number) => a * 2)
		const memo = memoizedFn('double', fn)

		expect(memo(3)).toBe(6)
		expect(memo(5)).toBe(10)
		expect(fn).toHaveBeenCalledTimes(2)
	})

	it('deep-compares non-Thread args', () => {
		const fn = vi.fn((obj: { x: number }) => obj.x * 2)
		const memo = memoizedFn('deep', fn)

		expect(memo({ x: 3 })).toBe(6)
		expect(memo({ x: 3 })).toBe(6) // new object, same shape
		expect(fn).toHaveBeenCalledTimes(1)
	})

	it('evicts oldest entry when maxSize exceeded', () => {
		const fn = vi.fn((n: number) => n * 10)
		const memo = memoizedFn('lru', fn, { maxSize: 2 })

		memo(1) // cache: [1]
		memo(2) // cache: [1, 2]
		memo(3) // cache: [2, 3] — evicts 1

		expect(fn).toHaveBeenCalledTimes(3)

		// 2 and 3 are cached
		memo(2)
		memo(3)
		expect(fn).toHaveBeenCalledTimes(3)

		// 1 was evicted, needs recomputation
		memo(1)
		expect(fn).toHaveBeenCalledTimes(4)
	})
})

describe('refCountedMemoizedFn', () => {
	it('returns same result for same args, increments refCount', () => {
		const fn = vi.fn((n: number) => ({ value: n * 2 }))
		const memo = refCountedMemoizedFn('rc', fn)

		const ref1 = memo(5)
		const ref2 = memo(5)
		expect(ref1.value).toBe(ref2.value) // same object
		expect(fn).toHaveBeenCalledTimes(1)
	})

	it('calls onCleanup when last ref released', () => {
		const cleanup = vi.fn()
		const fn = vi.fn((n: number) => n * 2)
		const memo = refCountedMemoizedFn('rc', fn, { onCleanup: cleanup })

		const ref1 = memo(5)
		const ref2 = memo(5)

		ref1.release()
		expect(cleanup).not.toHaveBeenCalled()

		ref2.release()
		expect(cleanup).toHaveBeenCalledOnce()
		expect(cleanup).toHaveBeenCalledWith(10, 5) // result, ...args
	})

	it('recomputes after full cleanup', () => {
		const fn = vi.fn((n: number) => n * 2)
		const memo = refCountedMemoizedFn('rc', fn, { onCleanup: () => {} })

		const ref1 = memo(3)
		ref1.release()
		expect(fn).toHaveBeenCalledTimes(1)

		const ref2 = memo(3)
		expect(fn).toHaveBeenCalledTimes(2)
		ref2.release()
	})

	it('release is idempotent', () => {
		const cleanup = vi.fn()
		const fn = vi.fn((n: number) => n * 2)
		const memo = refCountedMemoizedFn('rc', fn, { onCleanup: cleanup })

		const ref = memo(1)
		ref.release()
		ref.release() // second release is a no-op
		expect(cleanup).toHaveBeenCalledOnce()
	})

	it('grace period delays cleanup', () => {
		vi.useFakeTimers()
		const cleanup = vi.fn()
		const fn = vi.fn((n: number) => n * 2)
		const memo = refCountedMemoizedFn('rc', fn, {
			onCleanup: cleanup,
			gracePeriodMs: 100,
		})

		const ref = memo(5)
		ref.release()

		// Not cleaned up yet
		expect(cleanup).not.toHaveBeenCalled()

		// Re-acquire within grace period
		const ref2 = memo(5)
		expect(fn).toHaveBeenCalledTimes(1) // reused, not recomputed

		vi.advanceTimersByTime(200)
		expect(cleanup).not.toHaveBeenCalled() // still held by ref2

		ref2.release()
		vi.advanceTimersByTime(50)
		expect(cleanup).not.toHaveBeenCalled() // still in grace

		vi.advanceTimersByTime(60)
		expect(cleanup).toHaveBeenCalledOnce()

		vi.useRealTimers()
	})

	it('separate entries for different args', () => {
		const cleanup = vi.fn()
		const fn = vi.fn((n: number) => n * 2)
		const memo = refCountedMemoizedFn('rc', fn, { onCleanup: cleanup })

		const refA = memo(1)
		const refB = memo(2)
		expect(fn).toHaveBeenCalledTimes(2)

		refA.release()
		expect(cleanup).toHaveBeenCalledWith(2, 1)

		refB.release()
		expect(cleanup).toHaveBeenCalledWith(4, 2)
	})
})