import { Box, PageRecordType, createShapeId, react } from '@tldraw/editor'
import { vi } from 'vitest'
import { TestEditor } from './TestEditor'

let editor: TestEditor

beforeEach(() => {
	editor = new TestEditor()
	editor.updateViewportScreenBounds(new Box(0, 0, 1000, 1000))
	editor.setCamera({ x: 0, y: 0, z: 1 })
})

afterEach(() => {
	editor?.dispose()
})

// Track how many times a `getShapeIdsInsideBounds` reactor re-runs. The
// reactor's only dep is the spatial index epoch, so the delta is a direct
// proxy for "did the rbush epoch tick?". Initial subscribing run is dropped
// from the count.
function trackSpatialIndexInvalidations(target: TestEditor, bounds: Box) {
	let count = 0
	const stop = react('count spatial index invalidations', () => {
		target.getShapeIdsInsideBounds(bounds)
		count++
	})
	const baseline = count
	return {
		get delta() {
			return count - baseline
		},
		stop,
	}
}

describe('SpatialIndexManager - bounds epoch', () => {
	it('ticks when a shape moves (real bounds change)', () => {
		const id = createShapeId('mover')
		editor.createShapes([{ id, type: 'geo', x: 100, y: 100, props: { w: 100, h: 100 } }])

		const tracker = trackSpatialIndexInvalidations(editor, editor.getViewportPageBounds())

		editor.updateShapes([{ id, type: 'geo', x: 200, y: 200 }])

		expect(tracker.delta).toBe(1)
		tracker.stop()
	})

	it('does not tick on a prop-only update with unchanged bounds', () => {
		// Color change doesn't move bounds — the recheck must leave the
		// mutation batch empty and the epoch must stay still.
		const id = createShapeId('prop-only')
		editor.createShapes([
			{ id, type: 'geo', x: 100, y: 100, props: { w: 100, h: 100, color: 'black' } },
		])

		const tracker = trackSpatialIndexInvalidations(editor, editor.getViewportPageBounds())

		editor.updateShapes([{ id, type: 'geo', props: { color: 'red' } }])

		expect(tracker.delta).toBe(0)
		tracker.stop()
	})

	it('skips the rbush mutation batch on a prop-only diff', () => {
		// Direct check on the headline optimization: `RBushIndex.applyBatch`
		// must not be called when only props change.
		const id = createShapeId('upsert-spy')
		editor.createShapes([{ id, type: 'geo', x: 100, y: 100 }])
		// Initialize the computed before spying so buildFromScratch's
		// bulkLoad runs outside the spy window.
		editor.getShapeIdsInsideBounds(editor.getViewportPageBounds())

		const rbush = (editor as any)._spatialIndex.rbush
		const spy = vi.spyOn(rbush, 'applyBatch')

		editor.updateShapes([{ id, type: 'geo', props: { color: 'red' } }])
		editor.getShapeIdsInsideBounds(editor.getViewportPageBounds())
		expect(spy).not.toHaveBeenCalled()

		// Sanity: a real bounds change does apply a batch.
		editor.updateShapes([{ id, type: 'geo', x: 200, y: 200 }])
		editor.getShapeIdsInsideBounds(editor.getViewportPageBounds())
		expect(spy).toHaveBeenCalled()

		spy.mockRestore()
	})

	it('ticks exactly once when a multi-shape diff mixes prop-only and bounds changes', () => {
		// Guards against a future "if (changed) return early after step 1"
		// regression: with one mover and one prop-only shape in the same
		// transaction, the epoch must bump exactly once.
		const moverId = createShapeId('multi-mover')
		const propOnlyId = createShapeId('multi-prop-only')
		editor.createShapes([
			{ id: moverId, type: 'geo', x: 100, y: 100, props: { w: 100, h: 100 } },
			{ id: propOnlyId, type: 'geo', x: 300, y: 300, props: { w: 100, h: 100, color: 'black' } },
		])

		const tracker = trackSpatialIndexInvalidations(editor, editor.getViewportPageBounds())

		editor.run(() => {
			editor.updateShapes([
				{ id: moverId, type: 'geo', x: 200, y: 200 },
				{ id: propOnlyId, type: 'geo', props: { color: 'red' } },
			])
		})

		expect(tracker.delta).toBe(1)
		tracker.stop()
	})

	it('does not tick when removing a shape that lives on a different page', () => {
		const page1 = editor.getCurrentPageId()
		editor.createShapes([{ id: createShapeId('p1-anchor'), type: 'geo', x: 100, y: 100 }])

		// Stage a shape on a different page so its lifecycle never touches
		// page1's rbush.
		const page2 = PageRecordType.createId('page2-remove-epoch')
		editor.createPage({ name: 'page2-remove-epoch', id: page2 })
		editor.setCurrentPage(page2)
		const ghostId = createShapeId('p2-ghost')
		editor.createShapes([{ id: ghostId, type: 'geo', x: 100, y: 100 }])
		editor.setCurrentPage(page1)

		const tracker = trackSpatialIndexInvalidations(editor, editor.getViewportPageBounds())

		// Removing a shape that was never indexed on the active page must
		// not bump the bounds epoch.
		editor.deleteShape(ghostId)

		expect(tracker.delta).toBe(0)
		tracker.stop()
	})

	it('ticks exactly once on page switch (full rebuild)', () => {
		editor.createShapes([{ id: createShapeId('p1-shape'), type: 'geo', x: 100, y: 100 }])
		// Initialize the computed before tracking.
		editor.getShapeIdsInsideBounds(editor.getViewportPageBounds())

		const tracker = trackSpatialIndexInvalidations(editor, editor.getViewportPageBounds())

		const page2 = PageRecordType.createId('page2-switch')
		editor.createPage({ name: 'page2-switch', id: page2 })
		editor.setCurrentPage(page2)

		expect(tracker.delta).toBe(1)
		tracker.stop()
	})
})

describe('SpatialIndexManager - step-2 transitive bounds sweep', () => {
	it('refreshes arrow bounds when a bound geo shape changes outline without changing axis-aligned bounds', () => {
		// Regression: rectangle→ellipse keeps axis-aligned bounds the
		// same but moves the outline. The bound arrow's endpoint shifts,
		// changing the arrow's bounds. Step 1 sees the geo unchanged;
		// only step 2's sweep catches the dependent arrow.
		const targetId = createShapeId('outline-target')
		editor.createShapes([
			{
				id: targetId,
				type: 'geo',
				x: 0,
				y: 0,
				props: { w: 100, h: 100, geo: 'rectangle' },
			},
		])

		const targetBoundsBefore = editor.getShapePageBounds(targetId)!
		expect(targetBoundsBefore).toBeDefined()

		// Draw an arrow from a non-axis-aligned direction so the bound
		// endpoint lands somewhere rectangle and inscribed ellipse disagree.
		editor.setCurrentTool('arrow')
		editor.pointerDown(-100, -100)
		editor.pointerMove(50, 50)
		editor.pointerUp(50, 50)
		editor.selectNone()

		const arrow = editor.getCurrentPageShapes().find((s) => s.type === 'arrow')!
		expect(arrow).toBeDefined()

		const arrowBoundsBefore = editor.getShapePageBounds(arrow.id)!
		expect(arrowBoundsBefore).toBeDefined()

		// Flip target geometry without changing its axis-aligned bounds.
		editor.updateShapes([{ id: targetId, type: 'geo', props: { geo: 'ellipse' } }])

		// Geo's own bounds must be unchanged.
		const targetBoundsAfter = editor.getShapePageBounds(targetId)!
		expect(targetBoundsAfter.equals(targetBoundsBefore)).toBe(true)

		// Arrow's actual page bounds shifted because the bound endpoint
		// moved.
		const arrowBoundsAfter = editor.getShapePageBounds(arrow.id)!
		expect(arrowBoundsAfter.equals(arrowBoundsBefore)).toBe(false)

		// Probe a region that's only inside the new bounds, not the old.
		// If the rbush still has the old bounds, the probe misses.
		const probe = probeOnlyInAfter(arrowBoundsBefore, arrowBoundsAfter)
		const hits = editor.getShapeIdsInsideBounds(probe)
		expect(hits.has(arrow.id)).toBe(true)
	})

	it('removes an indexed shape from the source page when it is reparented to another page', () => {
		// Reparent-to-other-page hits step 1's "updated, not on page,
		// `rbush.has(id)` true → remove" branch.
		const page1 = editor.getCurrentPageId()
		const id = createShapeId('migrating')
		editor.createShapes([{ id, type: 'geo', x: 100, y: 100, props: { w: 100, h: 100 } }])
		expect(editor.getShapeIdsInsideBounds(new Box(0, 0, 1000, 1000)).has(id)).toBe(true)

		const page2 = PageRecordType.createId('page2-migrate')
		editor.createPage({ name: 'page2-migrate', id: page2 })
		editor.setCurrentPage(page1)

		const tracker = trackSpatialIndexInvalidations(editor, new Box(0, 0, 1000, 1000))
		editor.reparentShapes([id], page2)

		expect(tracker.delta).toBe(1)
		expect(editor.getShapeIdsInsideBounds(new Box(0, 0, 1000, 1000)).has(id)).toBe(false)
		tracker.stop()
	})
})

describe('SpatialIndexManager - binding-only diffs', () => {
	it('refreshes arrow bounds when bindings are created in a separate step from the shapes', () => {
		// Regression: an arrow's bounds derive from its bindings, but a
		// binding-only transaction contains no shape diff. If the index was
		// queried between createShapes and createBindings (the hover side
		// effect does this on every store change), the arrow stayed indexed
		// at its unbound bounds and parts of the bound body were
		// un-hoverable.
		const boxAId = createShapeId('boxA')
		const boxBId = createShapeId('boxB')
		const arrowId = createShapeId('arrow')

		editor.createShapes([
			{ id: boxAId, type: 'geo', x: 0, y: 0, props: { w: 100, h: 100, fill: 'none' } },
			{ id: boxBId, type: 'geo', x: 500, y: 0, props: { w: 100, h: 100, fill: 'none' } },
			// The arrow's own terminals are far from the shapes it will be
			// bound to, as happens when content is created via the API.
			{
				id: arrowId,
				type: 'arrow',
				x: 2000,
				y: 2000,
				props: { kind: 'arc', bend: -50, start: { x: 0, y: 0 }, end: { x: 100, y: 0 } },
			},
		])

		const unboundBounds = editor.getShapePageBounds(arrowId)!

		// Pull the spatial index so the arrow is indexed at its unbound
		// bounds before the bindings exist.
		expect(editor.getShapeIdsInsideBounds(unboundBounds).has(arrowId)).toBe(true)

		editor.createBindings([
			{
				type: 'arrow',
				fromId: arrowId,
				toId: boxAId,
				props: {
					terminal: 'start',
					isExact: false,
					isPrecise: false,
					normalizedAnchor: { x: 0.5, y: 0.5 },
				},
			},
			{
				type: 'arrow',
				fromId: arrowId,
				toId: boxBId,
				props: {
					terminal: 'end',
					isExact: false,
					isPrecise: false,
					normalizedAnchor: { x: 0.5, y: 0.5 },
				},
			},
		])

		const boundBounds = editor.getShapePageBounds(arrowId)!
		expect(boundBounds.equals(unboundBounds)).toBe(false)

		// The index must reflect the bound bounds...
		expect(editor.getShapeIdsInsideBounds(boundBounds).has(arrowId)).toBe(true)

		// ...and the arrow must be hit-testable on its bound body.
		const apex = { x: boundBounds.midX, y: boundBounds.minY + 1 }
		expect(editor.getShapeAtPoint(apex, { hitInside: false, margin: 8 })?.id).toBe(arrowId)
	})

	it('ticks the epoch and refreshes arrow bounds when a binding update moves the arrow', () => {
		const boxAId = createShapeId('boxA')
		const boxBId = createShapeId('boxB')
		const arrowId = createShapeId('arrow')

		editor.createShapes([
			{ id: boxAId, type: 'geo', x: 0, y: 0, props: { w: 100, h: 100, fill: 'none' } },
			{ id: boxBId, type: 'geo', x: 500, y: 500, props: { w: 100, h: 100, fill: 'none' } },
			{
				id: arrowId,
				type: 'arrow',
				x: 0,
				y: 0,
				props: { start: { x: 0, y: 0 }, end: { x: 100, y: 0 } },
			},
		])
		editor.createBindings([
			{
				type: 'arrow',
				fromId: arrowId,
				toId: boxAId,
				props: {
					terminal: 'start',
					isExact: false,
					isPrecise: true,
					normalizedAnchor: { x: 0, y: 0 },
				},
			},
			{
				type: 'arrow',
				fromId: arrowId,
				toId: boxBId,
				props: {
					terminal: 'end',
					isExact: false,
					isPrecise: true,
					normalizedAnchor: { x: 0, y: 0 },
				},
			},
		])

		const boundsBefore = editor.getShapePageBounds(arrowId)!
		const tracker = trackSpatialIndexInvalidations(editor, boundsBefore.clone())

		// Move the end terminal across the target shape with a binding-only
		// update; the arrow record itself is untouched.
		const binding = editor
			.getBindingsFromShape(arrowId, 'arrow')
			.find((b) => b.props.terminal === 'end')!
		editor.updateBinding({
			...binding,
			props: { ...binding.props, normalizedAnchor: { x: 0.5, y: 1 } },
		})

		const boundsAfter = editor.getShapePageBounds(arrowId)!
		expect(boundsAfter.equals(boundsBefore)).toBe(false)
		expect(tracker.delta).toBe(1)

		const probe = probeOnlyInAfter(boundsBefore, boundsAfter)
		expect(editor.getShapeIdsInsideBounds(probe).has(arrowId)).toBe(true)
		tracker.stop()
	})

	it('ticks the epoch and refreshes arrow bounds when a binding is deleted in its own transaction', () => {
		// The mirror of binding creation: deleting a binding without touching
		// the arrow record snaps its terminal back to the shape's own props,
		// shrinking the bounds. Exercises step 1's binding-removed branch,
		// which nothing else covers.
		const boxAId = createShapeId('boxA')
		const boxBId = createShapeId('boxB')
		const arrowId = createShapeId('arrow')

		editor.createShapes([
			{ id: boxAId, type: 'geo', x: 0, y: 0, props: { w: 100, h: 100, fill: 'none' } },
			{ id: boxBId, type: 'geo', x: 500, y: 500, props: { w: 100, h: 100, fill: 'none' } },
			{
				id: arrowId,
				type: 'arrow',
				x: 0,
				y: 0,
				props: { start: { x: 0, y: 0 }, end: { x: 100, y: 0 } },
			},
		])
		editor.createBindings([
			{
				type: 'arrow',
				fromId: arrowId,
				toId: boxAId,
				props: {
					terminal: 'start',
					isExact: false,
					isPrecise: true,
					normalizedAnchor: { x: 0, y: 0 },
				},
			},
			{
				type: 'arrow',
				fromId: arrowId,
				toId: boxBId,
				props: {
					terminal: 'end',
					isExact: false,
					isPrecise: true,
					normalizedAnchor: { x: 0, y: 0 },
				},
			},
		])

		const boundsBefore = editor.getShapePageBounds(arrowId)!
		// The bound arrow reaches across to boxB.
		expect(editor.getShapeIdsInsideBounds(new Box(490, 490, 120, 120)).has(arrowId)).toBe(true)

		const tracker = trackSpatialIndexInvalidations(editor, boundsBefore.clone())

		// Delete only the end binding; the arrow record itself is untouched, so
		// this is a binding-only transaction.
		const endBinding = editor
			.getBindingsFromShape(arrowId, 'arrow')
			.find((b) => b.props.terminal === 'end')!
		editor.deleteBindings([endBinding])

		const boundsAfter = editor.getShapePageBounds(arrowId)!
		expect(boundsAfter.equals(boundsBefore)).toBe(false)
		expect(tracker.delta).toBe(1)

		// The index must drop the stale reach toward boxB and hold the new,
		// smaller bounds.
		expect(editor.getShapeIdsInsideBounds(new Box(490, 490, 120, 120)).has(arrowId)).toBe(false)
		expect(editor.getShapeIdsInsideBounds(boundsAfter).has(arrowId)).toBe(true)
		tracker.stop()
	})

	it('rechecks the old arrow when a binding is reassigned to a different arrow', () => {
		// Defensive: no built-in flow changes a binding's fromId in place, but
		// a custom binding util could. Reassigning the binding must recheck the
		// old arrow, whose bounds revert now that it has lost the binding.
		const boxId = createShapeId('box')
		const arrow1Id = createShapeId('arrow1')
		const arrow2Id = createShapeId('arrow2')

		editor.createShapes([
			{ id: boxId, type: 'geo', x: 1000, y: 0, props: { w: 100, h: 100, fill: 'none' } },
			{
				id: arrow1Id,
				type: 'arrow',
				x: 0,
				y: 0,
				props: { kind: 'arc', bend: 40, start: { x: 0, y: 0 }, end: { x: 50, y: 0 } },
			},
			{
				id: arrow2Id,
				type: 'arrow',
				x: 0,
				y: 500,
				props: { start: { x: 0, y: 0 }, end: { x: 50, y: 0 } },
			},
		])
		editor.createBindings([
			{
				type: 'arrow',
				fromId: arrow1Id,
				toId: boxId,
				props: {
					terminal: 'end',
					isExact: false,
					isPrecise: true,
					normalizedAnchor: { x: 0.5, y: 0.5 },
				},
			},
		])

		// arrow1 now reaches across to the box; arrow2 does not.
		const boxRegion = new Box(930, -20, 220, 220)
		expect(editor.getShapeIdsInsideBounds(boxRegion).has(arrow1Id)).toBe(true)
		expect(editor.getShapeIdsInsideBounds(boxRegion).has(arrow2Id)).toBe(false)

		// Reassign the binding from arrow1 to arrow2 in place (binding-only).
		const binding = editor.getBindingsFromShape(arrow1Id, 'arrow')[0]
		editor.updateBinding({ ...binding, fromId: arrow2Id })

		// arrow2 inherits the reach toward the box; arrow1 must drop it.
		expect(editor.getShapeIdsInsideBounds(boxRegion).has(arrow2Id)).toBe(true)
		expect(editor.getShapeIdsInsideBounds(boxRegion).has(arrow1Id)).toBe(false)
	})
})

describe('SpatialIndexManager - basic queries', () => {
	it('returns shapes inside the queried bounds', () => {
		const insideId = createShapeId('inside')
		const outsideId = createShapeId('outside')
		editor.createShapes([
			{ id: insideId, type: 'geo', x: 100, y: 100, props: { w: 100, h: 100 } },
			{ id: outsideId, type: 'geo', x: 5000, y: 5000, props: { w: 100, h: 100 } },
		])

		const hits = editor.getShapeIdsInsideBounds(new Box(0, 0, 1000, 1000))
		expect(hits.has(insideId)).toBe(true)
		expect(hits.has(outsideId)).toBe(false)
	})

	it('removes a shape from the index when it is deleted', () => {
		const id = createShapeId('doomed')
		editor.createShapes([{ id, type: 'geo', x: 100, y: 100 }])
		expect(editor.getShapeIdsInsideBounds(new Box(0, 0, 1000, 1000)).has(id)).toBe(true)

		editor.deleteShape(id)
		expect(editor.getShapeIdsInsideBounds(new Box(0, 0, 1000, 1000)).has(id)).toBe(false)
	})
})

function probeOnlyInAfter(before: Box, after: Box): Box {
	const eps = 0.01
	const thickness = 1
	if (after.maxX > before.maxX + eps) {
		const x = before.maxX + eps
		return new Box(x, after.minY, Math.min(thickness, after.maxX - x), Math.max(after.h, eps))
	}
	if (after.maxY > before.maxY + eps) {
		const y = before.maxY + eps
		return new Box(after.minX, y, Math.max(after.w, eps), Math.min(thickness, after.maxY - y))
	}
	if (after.minX < before.minX - eps) {
		return new Box(
			after.minX,
			after.minY,
			Math.min(thickness, before.minX - after.minX - eps),
			Math.max(after.h, eps)
		)
	}
	if (after.minY < before.minY - eps) {
		return new Box(
			after.minX,
			after.minY,
			Math.max(after.w, eps),
			Math.min(thickness, before.minY - after.minY - eps)
		)
	}
	throw new Error('expected after-bounds to extend past before-bounds on at least one edge')
}