import { computed } from '@tldraw/state'
import {
	BaseRecord,
	RecordId,
	SerializedSchema,
	Store,
	StoreSchema,
	UnknownRecord,
	createMigrationIds,
	createMigrationSequence,
	createRecordMigrationSequence,
	createRecordType,
} from '@tldraw/store'
import { commentSchemaRecords, createTLSchema } from '@tldraw/tlschema'
import { vi, type Mock } from 'vitest'
import { RecordOpType, ValueOpType } from '../lib/diff'
import { TLSocketServerSentEvent, getTlsyncProtocolVersion } from '../lib/protocol'
import { TLSyncClient, TLSyncErrorCloseEventReason } from '../lib/TLSyncClient'
import { RoomSnapshot, TLRoomSocket } from '../lib/TLSyncRoom'
import { TestServer } from './TestServer'
import { TestSocketPair } from './TestSocketPair'

const actualProtocol = (await vi.importActual('../lib/protocol')) as any

vi.mock('../lib/protocol', async () => {
	const actual = (await vi.importActual('../lib/protocol')) as any
	return {
		...actual,
		getTlsyncProtocolVersion: vi.fn(actual.getTlsyncProtocolVersion),
	}
})

const mockGetTlsyncProtocolVersion = getTlsyncProtocolVersion as Mock

function mockSocket<R extends UnknownRecord>(): TLRoomSocket<R> {
	return {
		isOpen: true,
		sendMessage: vi.fn(),
		close: vi.fn(),
	}
}

// @ts-expect-error
global.requestAnimationFrame = (cb: () => any) => {
	cb()
}

const disposables: Array<() => void> = []

afterEach(() => {
	for (const dispose of disposables) {
		dispose()
	}
	disposables.length = 0
})

const UserVersions = createMigrationIds('com.tldraw.user', {
	ReplaceAgeWithBirthdate: 1,
} as const)

interface UserV1 extends BaseRecord<'user', RecordId<UserV1>> {
	name: string
	age: number
}
interface PresenceV1 extends BaseRecord<'presence', RecordId<PresenceV1>> {
	name: string
	age: number
}

const PresenceV1 = createRecordType<PresenceV1>('presence', {
	scope: 'presence',
	validator: { validate: (value) => value as PresenceV1 },
})

const UserV1 = createRecordType<UserV1>('user', {
	scope: 'document',
	validator: { validate: (value) => value as UserV1 },
})

interface UserV2 extends BaseRecord<'user', RecordId<UserV2>> {
	name: string
	birthdate: string | null
}

const userV2Migrations = createRecordMigrationSequence({
	sequenceId: 'com.tldraw.user',
	recordType: 'user',
	sequence: [
		{
			id: UserVersions.ReplaceAgeWithBirthdate,
			up({ age: _age, ...user }: any) {
				return {
					...user,
					birthdate: null,
				}
			},
			down({ birthdate: _birthdate, ...user }: any) {
				return {
					...user,
					age: 0,
				}
			},
		},
	],
})
const UserV2 = createRecordType<UserV2>('user', {
	scope: 'document',
	validator: { validate: (value) => value as UserV2 },
})

type RV1 = UserV1 | PresenceV1
type RV2 = UserV2 | PresenceV1

const userV1Migrations = createMigrationSequence({
	sequenceId: 'com.tldraw.user',
	sequence: [],
	retroactive: true,
})

const schemaV1 = StoreSchema.create<RV1>(
	{ user: UserV1, presence: PresenceV1 },
	{ migrations: [userV1Migrations] }
)

const schemaV2 = StoreSchema.create<RV2>(
	{ user: UserV2, presence: PresenceV1 },
	{
		migrations: [userV2Migrations],
	}
)

const schemaV3 = StoreSchema.create<RV2>(
	{ user: UserV2, presence: PresenceV1 },
	{
		migrations: [
			userV2Migrations,
			createMigrationSequence({
				sequenceId: 'com.tldraw.store',
				retroactive: true,
				sequence: [
					{
						id: 'com.tldraw.store/1',
						scope: 'store',
						up(store: any) {
							// remove any users called joe
							const result = Object.fromEntries(
								Object.entries(store).filter(
									([_, r]) => (r as any).typeName !== 'user' || (r as any).name !== 'joe'
								)
							)
							// add a user called steve
							const id = UserV2.createId('steve')
							result[id] = UserV2.create({
								id,
								name: 'steve',
								birthdate: '2022-02-02',
							})
							return result as any
						},
					},
				],
			}),
		],
	}
)

class TestInstance {
	server: TestServer<RV2>
	oldSocketPair: TestSocketPair<RV2>
	newSocketPair: TestSocketPair<RV2>
	oldClient: TLSyncClient<RV1>
	newClient: TLSyncClient<RV2>

	hasLoaded = false

	constructor(
		snapshot?: RoomSnapshot,
		oldSchema = schemaV1,
		newSchema = schemaV2,
		roomOpts?: ConstructorParameters<typeof TestServer<RV2>>[2]
	) {
		this.server = new TestServer(newSchema, snapshot, roomOpts)
		this.oldSocketPair = new TestSocketPair('test_upgrade_old', this.server)
		this.newSocketPair = new TestSocketPair('test_upgrade_new', this.server)

		this.oldClient = new TLSyncClient<RV1>({
			store: new Store({ schema: oldSchema, props: {} }),
			socket: this.oldSocketPair.clientSocket as any,
			onLoad: () => {
				this.hasLoaded = true
			},
			onSyncError: vi.fn((reason) => {
				throw new Error('onSyncError: ' + reason)
			}),
			presence: computed('', () => null),
		})

		this.newClient = new TLSyncClient<RV2>({
			store: new Store({ schema: newSchema, props: {} }),
			socket: this.newSocketPair.clientSocket,
			onLoad: () => {
				this.hasLoaded = true
			},
			onSyncError: vi.fn((reason) => {
				throw new Error('onSyncError: ' + reason)
			}),
			presence: computed('', () => null),
		})

		disposables.push(() => {
			this.oldClient.close()
			this.newClient.close()
		})
	}

	flush() {
		this.server.flushDebouncingMessages()

		while (this.oldSocketPair.getNeedsFlushing() || this.newSocketPair.getNeedsFlushing()) {
			this.oldSocketPair.flushClientSentEvents()
			this.oldSocketPair.flushServerSentEvents()
			this.newSocketPair.flushClientSentEvents()
			this.newSocketPair.flushServerSentEvents()
		}
	}
}

describe('record types unknown to older clients', () => {
	// Mirrors a deploy that adds the comment record types: stale tabs still run a schema
	// without them, while the server registers them and serves rooms that may contain them. The
	// comment types ship guard migrations (retroactive, no `down`), so the server must reject
	// stale sessions with CLIENT_TOO_OLD (prompting a refresh) rather than sending them record
	// types their store cannot represent.
	const schemaWithComments = createTLSchema({ records: commentSchemaRecords })
	const schemaWithoutComments = createTLSchema()

	function connectWith(server: TestServer<any>, schema: SerializedSchema) {
		const id = 'stale-tab'
		const socket = mockSocket()
		server.room.handleNewSession({ sessionId: id, socket, meta: undefined, isReadonly: false })
		server.room.handleMessage(id, {
			type: 'connect',
			connectRequestId: 'test',
			lastServerClock: 0,
			protocolVersion: getTlsyncProtocolVersion(),
			schema,
		})
		return socket
	}

	it('[HS3] rejects sessions whose schema predates the comment types', () => {
		const server = new TestServer(schemaWithComments as any, undefined, {
			objectTypes: ['comment', 'comment-thread'],
		})
		const socket = connectWith(server, schemaWithoutComments.serialize())
		expect(socket.close).toHaveBeenCalledWith(4099, TLSyncErrorCloseEventReason.CLIENT_TOO_OLD)
	})

	it('[HS3] accepts sessions whose schema registers the comment types', () => {
		const server = new TestServer(schemaWithComments as any, undefined, {
			objectTypes: ['comment', 'comment-thread'],
		})
		const socket = connectWith(server, schemaWithComments.serialize())
		expect(socket.close).not.toHaveBeenCalled()
		expect((socket.sendMessage as Mock).mock.calls[0][0]).toMatchObject({ type: 'connect' })
	})
})

it('[MG2] the server can handle receiving v1 stuff from the client', () => {
	const t = new TestInstance()
	t.oldSocketPair.connect()
	t.newSocketPair.connect()

	const user = UserV1.create({ name: 'bob', age: 10 })
	t.flush()
	t.oldClient.store.put([user])
	t.flush()

	expect(t.server.storage.documents.get(user.id)?.state).toMatchObject({
		name: 'bob',
		birthdate: null,
	})
	expect(t.server.storage.documents.get(user.id)?.state).not.toMatchObject({
		name: 'bob',
		age: 10,
	})

	expect(t.newClient.store.get(user.id as any)).toMatchObject({
		name: 'bob',
		birthdate: null,
	})
	expect(t.newClient.store.get(user.id as any)).not.toMatchObject({ name: 'bob', age: 10 })
})

it('[MG1] the server can send v2 stuff to the v1 client', () => {
	const t = new TestInstance()
	t.oldSocketPair.connect()
	t.newSocketPair.connect()

	const user = UserV2.create({ name: 'bob', birthdate: '2022-01-09' })
	t.flush()
	t.newClient.store.put([user])
	t.flush()

	expect(t.server.storage.documents.get(user.id)?.state).toMatchObject({
		name: 'bob',
		birthdate: '2022-01-09',
	})

	expect(t.oldClient.store.get(user.id as any)).toMatchObject({
		name: 'bob',
		age: 0,
	})
	expect(t.oldClient.store.get(user.id as any)).not.toMatchObject({
		name: 'bob',
		birthdate: '2022-01-09',
	})
})

it('[RC3] the server will run schema migrations on a snapshot', () => {
	const bob = UserV1.create({ name: 'bob', age: 10 })
	// joe will be deleted
	const joe = UserV1.create({ name: 'joe', age: 10 })
	const t = new TestInstance(
		{
			documents: [
				{ state: bob, lastChangedClock: 5 },
				{ state: joe, lastChangedClock: 5 },
			],
			clock: 10,
			schema: schemaV1.serialize(),
			tombstones: {},
		},
		schemaV1,
		schemaV3
	)

	expect(t.server.storage.documents.get(bob.id)?.state).toMatchObject({
		name: 'bob',
		birthdate: null,
	})
	expect(t.server.storage.documents.get(joe.id)).toBeUndefined()

	// there should be someone named steve
	const snapshot = t.server.storage.getSnapshot()
	expect(snapshot.documents.find((u: any) => u.state.name === 'steve')).toBeDefined()
})

it('[MG1] clients will receive updates from a snapshot migration upon connection', () => {
	const t = new TestInstance()
	t.oldSocketPair.connect()
	t.newSocketPair.connect()

	const bob = UserV2.create({ name: 'bob', birthdate: '2022-01-09' })
	const joe = UserV2.create({ name: 'joe', birthdate: '2022-01-09' })
	t.flush()
	t.newClient.store.put([bob, joe])
	t.flush()

	const snapshot = t.server.storage.getSnapshot()

	t.oldSocketPair.disconnect()
	t.newSocketPair.disconnect()

	const newServer = new TestServer(schemaV3, snapshot)

	const newClientSocketPair = new TestSocketPair('test_upgrade__brand_new', newServer)

	// need to set these two things to get the message through
	newClientSocketPair.callbacks['onReceiveMessage'] = vi.fn()
	newClientSocketPair.clientSocket.connectionStatus = 'online'

	const id = 'test_upgrade_brand_new'
	const newClientSocket = mockSocket()
	newServer.room.handleNewSession({
		sessionId: id,
		socket: newClientSocket,
		meta: undefined,
		isReadonly: false,
	})
	newServer.room.handleMessage(id, {
		type: 'connect',
		connectRequestId: 'test',
		lastServerClock: snapshot.documentClock ?? snapshot.clock ?? 0,
		protocolVersion: getTlsyncProtocolVersion(),
		schema: schemaV3.serialize(),
	})

	expect((newClientSocket.sendMessage as Mock).mock.calls[0][0]).toMatchObject({
		// we should have added steve and deleted joe
		diff: {
			[joe.id]: [RecordOpType.Remove],
			['user:steve']: [RecordOpType.Put, { name: 'steve', birthdate: '2022-02-02' }],
		},
	})
})

it('[HS3] out-of-date clients will receive incompatibility errors', () => {
	const v3server = new TestServer(schemaV3)

	const id = 'test_upgrade_v2'
	const socket = mockSocket()

	v3server.room.handleNewSession({ sessionId: id, socket, meta: undefined, isReadonly: false })
	v3server.room.handleMessage(id, {
		type: 'connect',
		connectRequestId: 'test',
		lastServerClock: 0,
		protocolVersion: getTlsyncProtocolVersion(),
		schema: schemaV2.serialize(),
	})

	expect(socket.close).toHaveBeenCalledWith(4099, TLSyncErrorCloseEventReason.CLIENT_TOO_OLD)
})

it('[HS2] clients using an out-of-date protocol will receive compatibility errors', () => {
	const actualVersion = getTlsyncProtocolVersion()
	mockGetTlsyncProtocolVersion.mockReturnValue(actualVersion + 1)
	try {
		const v2server = new TestServer(schemaV2)

		const id = 'test_upgrade_v3'
		const socket = mockSocket()

		v2server.room.handleNewSession({ sessionId: id, socket, meta: undefined, isReadonly: false })
		v2server.room.handleMessage(id, {
			type: 'connect',
			connectRequestId: 'test',
			lastServerClock: 0,
			protocolVersion: actualVersion,
			schema: schemaV2.serialize(),
		})

		expect(socket.close).toHaveBeenCalledWith(4099, TLSyncErrorCloseEventReason.CLIENT_TOO_OLD)
	} finally {
		mockGetTlsyncProtocolVersion.mockReset()
		mockGetTlsyncProtocolVersion.mockImplementation(actualProtocol.getTlsyncProtocolVersion)
	}
})

// protocol v5 is treated as v6: accepted, but with the legacy rejection scheme
// and without string-append support
it('[HS2] v5 special case should allow connections', () => {
	const v2server = new TestServer(schemaV2)

	const id = 'test_upgrade_v5'
	const socket = mockSocket()

	v2server.room.handleNewSession({ sessionId: id, socket, meta: undefined, isReadonly: false })
	v2server.room.handleMessage(id, {
		type: 'connect',
		connectRequestId: 'test',
		lastServerClock: 0,
		protocolVersion: 5,
		schema: schemaV2.serialize(),
	})

	expect(socket.close).not.toHaveBeenCalled()
	expect((socket.sendMessage as Mock).mock.calls[0][0]).toMatchObject({
		type: 'connect',
		connectRequestId: 'test',
		protocolVersion: getTlsyncProtocolVersion(),
		isReadonly: false,
	})

	const session = v2server.room.sessions.get(id) as any
	expect(session?.state).toBe('connected')
	expect(session?.requiresLegacyRejection).toBe(true)
	expect(session?.supportsStringAppend).toBe(false)
})

it('[HS2] clients using a too-new protocol will receive compatibility errors', () => {
	const v2server = new TestServer(schemaV2)

	const id = 'test_upgrade_v3'
	const socket = mockSocket()

	v2server.room.handleNewSession({ sessionId: id, socket, meta: undefined, isReadonly: false })
	v2server.room.handleMessage(id, {
		type: 'connect',
		connectRequestId: 'test',
		lastServerClock: 0,
		protocolVersion: getTlsyncProtocolVersion() + 1,
		schema: schemaV2.serialize(),
	})

	expect(socket.close).toHaveBeenCalledWith(4099, TLSyncErrorCloseEventReason.SERVER_TOO_OLD)
})

it('[HS3] when the client is too new it cannot connect', () => {
	const steve = UserV1.create({ id: UserV1.createId('steve'), name: 'steve', age: 23 })
	const jeff = UserV1.create({ id: UserV1.createId('jeff'), name: 'jeff', age: 23 })
	const annie = UserV1.create({ id: UserV1.createId('annie'), name: 'annie', age: 23 })
	const v1Server = new TestServer(schemaV1, {
		clock: 10,
		documents: [
			{
				state: steve,
				lastChangedClock: 10,
			},
			{
				state: jeff,
				lastChangedClock: 10,
			},
			{
				state: annie,
				lastChangedClock: 10,
			},
		],
		schema: schemaV1.serialize(),
		tombstones: {},
	})

	const v2_id = 'test_upgrade_v2'
	const v2_socket = mockSocket<RV2>()

	v1Server.room.handleNewSession({
		sessionId: v2_id,
		socket: v2_socket as any,
		meta: undefined,
		isReadonly: false,
	})
	v1Server.room.handleMessage(v2_id as any, {
		type: 'connect',
		connectRequestId: 'test',
		lastServerClock: 10,
		protocolVersion: getTlsyncProtocolVersion(),
		schema: schemaV2.serialize(),
	})

	expect(v2_socket.close).toHaveBeenCalledWith(
		4099,
		// the client is on a newer schema than the server, so from its perspective the server is
		// the one that's too old.
		TLSyncErrorCloseEventReason.SERVER_TOO_OLD
	)
})

describe('when the client is too old', () => {
	function setup() {
		const steve = UserV2.create({
			id: UserV2.createId('steve'),
			name: 'steve',
			birthdate: null,
		})
		const jeff = UserV2.create({ id: UserV2.createId('jeff'), name: 'jeff', birthdate: null })
		const annie = UserV2.create({
			id: UserV2.createId('annie'),
			name: 'annie',
			birthdate: null,
		})
		const v2Server = new TestServer(schemaV2, {
			clock: 10,
			documents: [
				{
					state: steve,
					lastChangedClock: 10,
				},
				{
					state: jeff,
					lastChangedClock: 10,
				},
				{
					state: annie,
					lastChangedClock: 10,
				},
			],
			schema: schemaV1.serialize(),
			tombstones: {},
		})

		const v2Id = 'test_upgrade_v2'
		const v2Socket = mockSocket<RV2>()

		const v2SendMessage = v2Socket.sendMessage as Mock

		const v1Id = 'test_upgrade_v1'
		const v1Socket = mockSocket<RV1>()

		const v1SendMessage = v1Socket.sendMessage as Mock

		v2Server.room.handleNewSession({
			sessionId: v1Id,
			socket: v1Socket as any,
			meta: undefined,
			isReadonly: false,
		})
		v2Server.room.handleMessage(v1Id, {
			type: 'connect',
			connectRequestId: 'test',
			lastServerClock: 10,
			protocolVersion: getTlsyncProtocolVersion(),
			schema: schemaV1.serialize(),
		})

		v2Server.room.handleNewSession({
			sessionId: v2Id,
			socket: v2Socket,
			meta: undefined,
			isReadonly: false,
		})
		v2Server.room.handleMessage(v2Id, {
			type: 'connect',
			connectRequestId: 'test',
			lastServerClock: 10,
			protocolVersion: getTlsyncProtocolVersion(),
			schema: schemaV2.serialize(),
		})

		expect(v2SendMessage).toHaveBeenCalledWith({
			type: 'connect',
			connectRequestId: 'test',
			hydrationType: 'wipe_presence',
			diff: {},
			protocolVersion: getTlsyncProtocolVersion(),
			schema: schemaV2.serialize(),
			serverClock: 10,
			isReadonly: false,
			objectAccess: 'write',
		} satisfies TLSocketServerSentEvent<RV2>)

		expect(v1SendMessage).toHaveBeenCalledWith({
			type: 'connect',
			connectRequestId: 'test',
			hydrationType: 'wipe_presence',
			diff: {},
			protocolVersion: getTlsyncProtocolVersion(),
			schema: schemaV2.serialize(),
			serverClock: 10,
			isReadonly: false,
			objectAccess: 'write',
		} satisfies TLSocketServerSentEvent<RV2>)

		v2SendMessage.mockClear()
		v1SendMessage.mockClear()

		return {
			v2Server,
			v2Id,
			v1Id,
			v2SendMessage,
			v1SendMessage,
			steve,
			jeff,
			annie,
		}
	}

	let data: ReturnType<typeof setup>

	beforeEach(() => {
		data = setup()
	})

	it('[MG2] allows deletions from v1 client', () => {
		data.v2Server.room.handleMessage(data.v1Id, {
			type: 'push',
			clientClock: 1,
			diff: {
				[data.steve.id]: [RecordOpType.Remove],
			},
		})

		expect(data.v1SendMessage).toHaveBeenCalledWith({
			type: 'data',
			data: [
				{
					type: 'push_result',
					action: 'commit',
					clientClock: 1,
					serverClock: 11,
				},
			],
		} satisfies TLSocketServerSentEvent<RV2>)

		expect(data.v2SendMessage).toHaveBeenCalledWith({
			type: 'data',
			data: [
				{
					type: 'patch',
					diff: {
						[data.steve.id]: [RecordOpType.Remove],
					},
					serverClock: 11,
				},
			],
		} satisfies TLSocketServerSentEvent<RV2>)
	})

	it('[MG2] can handle patches from older clients', () => {
		data.v2Server.room.handleMessage(data.v1Id, {
			type: 'push',
			clientClock: 1,
			diff: {
				[data.steve.id]: [RecordOpType.Patch, { name: [ValueOpType.Put, 'Jeff'] }],
			},
		})

		expect(data.v1SendMessage).toHaveBeenCalledWith({
			type: 'data',
			data: [
				{
					type: 'push_result',
					action: 'commit',
					clientClock: 1,
					serverClock: 11,
				},
			],
		} satisfies TLSocketServerSentEvent<RV2>)

		expect(data.v2SendMessage).toHaveBeenCalledWith({
			type: 'data',
			data: [
				{
					type: 'patch',
					diff: {
						[data.steve.id]: [
							RecordOpType.Patch,
							{
								name: [ValueOpType.Put, 'Jeff'],
							},
						],
					},
					serverClock: 11,
				},
			],
		} satisfies TLSocketServerSentEvent<RV2>)
	})
})

describe('migration failure (TLSyncError handling)', () => {
	let consoleSpy: ReturnType<typeof vi.spyOn>
	beforeEach(() => {
		consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {})
	})
	afterEach(() => {
		consoleSpy.mockRestore()
	})

	// Create a schema where migrations will fail during push
	const UserVersionsWithFailure = createMigrationIds('com.tldraw.user.failure', {
		AddNickname: 1,
	} as const)

	interface UserWithNickname extends BaseRecord<'user_fail', RecordId<UserWithNickname>> {
		name: string
		nickname: string
	}

	const UserWithNicknameType = createRecordType<UserWithNickname>('user_fail', {
		scope: 'document',
		validator: { validate: (value) => value as UserWithNickname },
	})

	interface UserWithoutNickname extends BaseRecord<'user_fail', RecordId<UserWithoutNickname>> {
		name: string
	}

	const UserWithoutNicknameType = createRecordType<UserWithoutNickname>('user_fail', {
		scope: 'document',
		validator: { validate: (value) => value as UserWithoutNickname },
	})

	// Server schema with migration that throws on specific conditions
	// - up() throws when name === 'FAIL_UP_MIGRATION'
	// - down() throws when nickname === 'FAIL_DOWN_MIGRATION'
	const serverSchemaWithThrowingMigration = StoreSchema.create<UserWithNickname>(
		{ user_fail: UserWithNicknameType },
		{
			migrations: [
				createRecordMigrationSequence({
					sequenceId: 'com.tldraw.user.failure',
					recordType: 'user_fail',
					sequence: [
						{
							id: UserVersionsWithFailure.AddNickname,
							up(record: any) {
								// Simulate UP migration failure for specific data
								if (record.name === 'FAIL_UP_MIGRATION') {
									throw new Error('UP migration intentionally failed for testing')
								}
								return {
									...record,
									nickname: record.name.toUpperCase(),
								}
							},
							down(record: any) {
								// Simulate DOWN migration failure for specific data
								if (record.nickname === 'FAIL_DOWN_MIGRATION') {
									throw new Error('DOWN migration intentionally failed for testing')
								}
								const { nickname: _, ...rest } = record
								return rest
							},
						},
					],
				}),
			],
		}
	)

	// Client schema (old version without nickname)
	const clientSchemaWithoutNickname = StoreSchema.create<UserWithoutNickname>(
		{ user_fail: UserWithoutNicknameType },
		{
			migrations: [
				createMigrationSequence({
					sequenceId: 'com.tldraw.user.failure',
					sequence: [],
					retroactive: true,
				}),
			],
		}
	)

	it('[MG2] rejects session when addDocument up migration throws during put', () => {
		const server = new TestServer(serverSchemaWithThrowingMigration)

		const sessionId = 'failing-migration-session'
		const socket = mockSocket<UserWithNickname>()

		server.room.handleNewSession({
			sessionId,
			socket: socket as any,
			meta: undefined,
			isReadonly: false,
		})

		// Connect with old schema - should succeed since migration has down function
		server.room.handleMessage(sessionId, {
			type: 'connect',
			connectRequestId: 'test',
			lastServerClock: 0,
			protocolVersion: getTlsyncProtocolVersion(),
			schema: clientSchemaWithoutNickname.serialize(),
		})

		expect(server.room.sessions.get(sessionId)?.state).toBe('connected')
		;(socket.sendMessage as Mock).mockClear()

		// Push a record that will cause the UP migration to throw
		const failingRecord = UserWithoutNicknameType.create({
			id: UserWithoutNicknameType.createId('fail'),
			name: 'FAIL_UP_MIGRATION', // This name triggers the UP throw
		})

		server.room.handleMessage(sessionId, {
			type: 'push',
			clientClock: 1,
			diff: {
				[failingRecord.id]: [RecordOpType.Put, failingRecord as any],
			},
		})

		// Session should be rejected due to migration failure (TLSyncError)
		expect(socket.close).toHaveBeenCalledWith(4099, TLSyncErrorCloseEventReason.CLIENT_TOO_OLD)
	})

	it('[MG2] successfully migrates records when migration does not throw', () => {
		const server = new TestServer(serverSchemaWithThrowingMigration)

		const sessionId = 'successful-migration-session'
		const socket = mockSocket<UserWithNickname>()

		server.room.handleNewSession({
			sessionId,
			socket: socket as any,
			meta: undefined,
			isReadonly: false,
		})

		// Connect with old schema
		server.room.handleMessage(sessionId, {
			type: 'connect',
			connectRequestId: 'test',
			lastServerClock: 0,
			protocolVersion: getTlsyncProtocolVersion(),
			schema: clientSchemaWithoutNickname.serialize(),
		})

		expect(server.room.sessions.get(sessionId)?.state).toBe('connected')
		;(socket.sendMessage as Mock).mockClear()

		// Push a record that will NOT cause the migration to throw
		const validRecord = UserWithoutNicknameType.create({
			id: UserWithoutNicknameType.createId('valid'),
			name: 'ValidName', // This name does not trigger the throw
		})

		server.room.handleMessage(sessionId, {
			type: 'push',
			clientClock: 1,
			diff: {
				[validRecord.id]: [RecordOpType.Put, validRecord as any],
			},
		})

		// Session should still be connected
		expect(server.room.sessions.get(sessionId)?.state).toBe('connected')

		// Record should be migrated and stored
		const storedRecord = server.storage.documents.get(validRecord.id)?.state as UserWithNickname
		expect(storedRecord).toBeDefined()
		expect(storedRecord.nickname).toBe('VALIDNAME') // Uppercase from migration
	})

	it('[MG2] rejects session when patchDocument down migration throws', () => {
		// Create a server with a document that has a nickname that will fail DOWN migration
		const existingUser = UserWithNicknameType.create({
			id: UserWithNicknameType.createId('existing_down_fail'),
			name: 'ExistingUser',
			nickname: 'FAIL_DOWN_MIGRATION', // This nickname triggers the DOWN throw
		})

		const server = new TestServer(serverSchemaWithThrowingMigration, {
			clock: 10,
			documents: [{ state: existingUser, lastChangedClock: 10 }],
			schema: serverSchemaWithThrowingMigration.serialize(),
			tombstones: {},
		})

		const sessionId = 'patch-down-migration-session'
		const socket = mockSocket<UserWithNickname>()

		server.room.handleNewSession({
			sessionId,
			socket: socket as any,
			meta: undefined,
			isReadonly: false,
		})

		// Connect with old schema
		server.room.handleMessage(sessionId, {
			type: 'connect',
			connectRequestId: 'test',
			lastServerClock: 10,
			protocolVersion: getTlsyncProtocolVersion(),
			schema: clientSchemaWithoutNickname.serialize(),
		})

		expect(server.room.sessions.get(sessionId)?.state).toBe('connected')
		;(socket.sendMessage as Mock).mockClear()

		// Patch the existing record - the DOWN migration will fail because nickname is 'FAIL_DOWN_MIGRATION'
		// This tests line 951 in TLSyncRoom.ts
		server.room.handleMessage(sessionId, {
			type: 'push',
			clientClock: 1,
			diff: {
				[existingUser.id]: [RecordOpType.Patch, { name: [ValueOpType.Put, 'NewName'] }],
			},
		})

		// Session should be rejected due to DOWN migration failure
		expect(socket.close).toHaveBeenCalledWith(4099, TLSyncErrorCloseEventReason.CLIENT_TOO_OLD)
	})

	it('[MG2] rejects session when patchDocument up migration throws (after successful down migration)', () => {
		// Create a server with a document that will pass DOWN migration but fail UP migration
		const existingUser = UserWithNicknameType.create({
			id: UserWithNicknameType.createId('existing_up_fail'),
			name: 'ExistingUser',
			nickname: 'EXISTINGUSER', // Normal nickname, DOWN migration will succeed
		})

		const server = new TestServer(serverSchemaWithThrowingMigration, {
			clock: 10,
			documents: [{ state: existingUser, lastChangedClock: 10 }],
			schema: serverSchemaWithThrowingMigration.serialize(),
			tombstones: {},
		})

		const sessionId = 'patch-up-migration-session'
		const socket = mockSocket<UserWithNickname>()

		server.room.handleNewSession({
			sessionId,
			socket: socket as any,
			meta: undefined,
			isReadonly: false,
		})

		// Connect with old schema
		server.room.handleMessage(sessionId, {
			type: 'connect',
			connectRequestId: 'test',
			lastServerClock: 10,
			protocolVersion: getTlsyncProtocolVersion(),
			schema: clientSchemaWithoutNickname.serialize(),
		})

		expect(server.room.sessions.get(sessionId)?.state).toBe('connected')
		;(socket.sendMessage as Mock).mockClear()

		// Patch the existing record - changes the name to 'FAIL_UP_MIGRATION'
		// 1. DOWN migration: {name: 'ExistingUser', nickname: 'EXISTINGUSER'} -> {name: 'ExistingUser'} (succeeds)
		// 2. Patch applied: {name: 'ExistingUser'} -> {name: 'FAIL_UP_MIGRATION'}
		// 3. UP migration: {name: 'FAIL_UP_MIGRATION'} -> THROWS
		// This tests line 972 in TLSyncRoom.ts
		server.room.handleMessage(sessionId, {
			type: 'push',
			clientClock: 1,
			diff: {
				[existingUser.id]: [RecordOpType.Patch, { name: [ValueOpType.Put, 'FAIL_UP_MIGRATION'] }],
			},
		})

		// Session should be rejected due to UP migration failure during patch
		expect(socket.close).toHaveBeenCalledWith(4099, TLSyncErrorCloseEventReason.CLIENT_TOO_OLD)
	})

	it('[MG1] rejects only the affected session when down migration fails during broadcast', () => {
		const server = new TestServer(serverSchemaWithThrowingMigration)

		// An old-schema client whose outgoing diffs require down migration
		const oldId = 'broadcast-old-schema-session'
		const oldSocket = mockSocket<UserWithNickname>()
		server.room.handleNewSession({
			sessionId: oldId,
			socket: oldSocket as any,
			meta: undefined,
			isReadonly: false,
		})
		server.room.handleMessage(oldId, {
			type: 'connect',
			connectRequestId: 'test',
			lastServerClock: 0,
			protocolVersion: getTlsyncProtocolVersion(),
			schema: clientSchemaWithoutNickname.serialize(),
		})

		// An up-to-date client that pushes a record whose down migration throws
		const newId = 'broadcast-new-schema-session'
		const newSocket = mockSocket<UserWithNickname>()
		server.room.handleNewSession({
			sessionId: newId,
			socket: newSocket as any,
			meta: undefined,
			isReadonly: false,
		})
		server.room.handleMessage(newId, {
			type: 'connect',
			connectRequestId: 'test',
			lastServerClock: 0,
			protocolVersion: getTlsyncProtocolVersion(),
			schema: serverSchemaWithThrowingMigration.serialize(),
		})

		expect(server.room.sessions.get(oldId)?.state).toBe('connected')
		expect(server.room.sessions.get(newId)?.state).toBe('connected')
		;(oldSocket.sendMessage as Mock).mockClear()
		;(newSocket.sendMessage as Mock).mockClear()

		const failingRecord = UserWithNicknameType.create({
			id: UserWithNicknameType.createId('broadcast_fail'),
			name: 'BroadcastUser',
			nickname: 'FAIL_DOWN_MIGRATION', // triggers the DOWN throw when broadcast to old client
		})

		server.room.handleMessage(newId, {
			type: 'push',
			clientClock: 1,
			diff: {
				[failingRecord.id]: [RecordOpType.Put, failingRecord as any],
			},
		})

		// The push itself succeeds: the record is stored and the pusher gets a commit
		expect(server.storage.documents.get(failingRecord.id)?.state).toMatchObject({
			nickname: 'FAIL_DOWN_MIGRATION',
		})
		expect(newSocket.close).not.toHaveBeenCalled()
		expect(newSocket.sendMessage).toHaveBeenCalledWith({
			type: 'data',
			data: [
				{
					type: 'push_result',
					action: 'commit',
					clientClock: 1,
					serverClock: 1,
				},
			],
		})

		// Only the old-schema session is rejected, because broadcasting the put to it
		// requires a down migration that throws
		expect(oldSocket.close).toHaveBeenCalledWith(4099, TLSyncErrorCloseEventReason.CLIENT_TOO_OLD)
	})
})

describe('when the client is the same version', () => {
	function setup() {
		const steve = UserV2.create({
			id: UserV2.createId('steve'),
			name: 'steve',
			birthdate: null,
		})
		const v2Server = new TestServer(schemaV2, {
			clock: 10,
			documents: [
				{
					state: steve,
					lastChangedClock: 10,
				},
			],
			schema: schemaV2.serialize(),
			tombstones: {},
		})

		const aId = 'v2ClientA'
		const aSocket = mockSocket<RV2>()

		const bId = 'v2ClientB'
		const bSocket = mockSocket<RV2>()

		v2Server.room.handleNewSession({
			sessionId: aId,
			socket: aSocket,
			meta: undefined,
			isReadonly: false,
		})
		v2Server.room.handleMessage(aId, {
			type: 'connect',
			connectRequestId: 'test',
			lastServerClock: 10,
			protocolVersion: getTlsyncProtocolVersion(),
			schema: JSON.parse(JSON.stringify(schemaV2.serialize())),
		})

		v2Server.room.handleNewSession({
			sessionId: bId,
			socket: bSocket,
			meta: undefined,
			isReadonly: false,
		})
		v2Server.room.handleMessage(bId, {
			type: 'connect',
			connectRequestId: 'test',
			lastServerClock: 10,
			protocolVersion: getTlsyncProtocolVersion(),
			schema: JSON.parse(JSON.stringify(schemaV2.serialize())),
		})

		expect(aSocket.sendMessage).toHaveBeenCalledWith({
			type: 'connect',
			connectRequestId: 'test',
			hydrationType: 'wipe_presence',
			diff: {},
			protocolVersion: getTlsyncProtocolVersion(),
			schema: schemaV2.serialize(),
			serverClock: 10,
			isReadonly: false,
			objectAccess: 'write',
		} satisfies TLSocketServerSentEvent<RV2>)

		expect(bSocket.sendMessage).toHaveBeenCalledWith({
			type: 'connect',
			connectRequestId: 'test',
			hydrationType: 'wipe_presence',
			diff: {},
			protocolVersion: getTlsyncProtocolVersion(),
			schema: schemaV2.serialize(),
			serverClock: 10,
			isReadonly: false,
			objectAccess: 'write',
		} satisfies TLSocketServerSentEvent<RV2>)
		;(aSocket.sendMessage as Mock).mockClear()
		;(bSocket.sendMessage as Mock).mockClear()

		return {
			v2Server,
			aId,
			bId,
			v2ClientASendMessage: aSocket.sendMessage as Mock,
			v2ClientBSendMessage: bSocket.sendMessage as Mock,
			steve,
		}
	}

	let data: ReturnType<typeof setup>

	beforeEach(() => {
		data = setup()
	})

	it('[MG3] accepts client with same schema version', () => {
		// both clients connected with a deep-equal (but not reference-equal) serialized
		// schema; they are accepted (see setup assertions) and share the server's schema
		// object so migration is skipped entirely
		const aSession = data.v2Server.room.sessions.get(data.aId) as any
		const bSession = data.v2Server.room.sessions.get(data.bId) as any
		expect(aSession?.state).toBe('connected')
		expect(bSession?.state).toBe('connected')
		expect(aSession?.serializedSchema).toBe(data.v2Server.room.serializedSchema)
		expect(bSession?.serializedSchema).toBe(data.v2Server.room.serializedSchema)
		expect(aSession?.requiresDownMigrations).toBe(false)
		expect(bSession?.requiresDownMigrations).toBe(false)
	})

	it('[MG3] sends minimal patches', () => {
		data.v2Server.room.handleMessage(data.aId, {
			type: 'push',
			clientClock: 1,
			diff: {
				[data.steve.id]: [RecordOpType.Patch, { name: [ValueOpType.Put, 'Jeff'] }],
			},
		})

		expect(data.v2ClientASendMessage).toHaveBeenCalledWith({
			type: 'data',
			data: [
				{
					type: 'push_result',
					action: 'commit',
					clientClock: 1,
					serverClock: 11,
				},
			],
		} satisfies TLSocketServerSentEvent<RV2>)

		expect(data.v2ClientBSendMessage).toHaveBeenCalledWith({
			type: 'data',
			data: [
				{
					type: 'patch',
					diff: {
						[data.steve.id]: [
							RecordOpType.Patch,
							{
								name: [ValueOpType.Put, 'Jeff'],
							},
						],
					},
					serverClock: 11,
				},
			],
		} satisfies TLSocketServerSentEvent<RV2>)
	})
})

describe('record authorizers see server-schema records', () => {
	function makeInstance(authorizeUser: (args: any) => any) {
		return new TestInstance(undefined, schemaV1, schemaV2, {
			authorizeRecord: { user: authorizeUser },
		})
	}

	it('a put from an old client reaches the authorizer up-migrated', () => {
		const seen: any[] = []
		const t = makeInstance((args: any) => {
			seen.push(structuredClone({ type: args.type, prev: args.prev, next: args.next }))
			return args.next
		})
		t.oldSocketPair.connect()
		t.flush()

		const user = UserV1.create({ name: 'bob', age: 10 })
		t.oldClient.store.put([user])
		t.flush()

		expect(seen).toHaveLength(1)
		expect(seen[0].type).toBe('create')
		expect(seen[0].next).toMatchObject({ name: 'bob', birthdate: null })
		expect(seen[0].next).not.toHaveProperty('age')
	})

	it('a stamp applied on create survives — no migration runs after the authorizer', () => {
		const t = makeInstance((args: any) =>
			args.type === 'create' ? { ...args.next, birthdate: '2001-02-03' } : args.next
		)
		t.oldSocketPair.connect()
		t.flush()

		const user = UserV1.create({ name: 'bob', age: 10 })
		t.oldClient.store.put([user])
		t.flush()

		expect(t.server.storage.documents.get(user.id)?.state).toMatchObject({
			name: 'bob',
			birthdate: '2001-02-03',
		})
	})

	it('a patch from an old client reaches the authorizer as the upgraded committed candidate', () => {
		const seen: any[] = []
		const t = makeInstance((args: any) => {
			seen.push(structuredClone({ type: args.type, prev: args.prev, next: args.next }))
			return args.next
		})
		t.oldSocketPair.connect()
		t.newSocketPair.connect()
		t.flush()

		const user = UserV2.create({ name: 'bob', birthdate: '2022-01-09' })
		t.newClient.store.put([user])
		t.flush()
		seen.length = 0

		t.oldClient.store.update(user.id as any, (u: any) => ({ ...u, name: 'bobby' }))
		t.flush()

		expect(seen).toHaveLength(1)
		expect(seen[0].type).toBe('update')
		expect(seen[0].prev).toMatchObject({ name: 'bob', birthdate: '2022-01-09' })
		// The down→patch→up round-trip through the lossy v1 schema resets birthdate to null.
		// The authorizer must see exactly what will be committed — not a preview mixing the
		// server record with the raw v1 diff (which would still show the old birthdate).
		expect(seen[0].next).toMatchObject({ name: 'bobby', birthdate: null })
		expect(seen[0].next).not.toHaveProperty('age')
	})

	it('vetoes based on the server-schema record', () => {
		const t = makeInstance((args: any) =>
			args.type === 'update' && args.next.name === 'forged' ? null : args.next
		)
		t.oldSocketPair.connect()
		t.newSocketPair.connect()
		t.flush()

		const user = UserV2.create({ name: 'bob', birthdate: '2022-01-09' })
		t.newClient.store.put([user])
		t.flush()

		t.oldClient.store.update(user.id as any, (u: any) => ({ ...u, name: 'forged' }))
		t.flush()

		expect(t.server.storage.documents.get(user.id)?.state).toMatchObject({ name: 'bob' })
	})
})