import type { SystemModelMessage } from 'ai';

import { isLlmMessage } from '../../sdk/message';
import type {
	AgentDbMessage,
	AgentMessage,
	ContentToolCall,
	Message,
} from '../../types/sdk/message';
import {
	AgentMessageList,
	buildSystemMessages,
	OBSERVATION_CONTINUATION_REMINDER,
} from '../model/message-list';

function flattenSystemContent(system: SystemModelMessage | SystemModelMessage[]): string {
	if (Array.isArray(system)) {
		return system.map((entry) => entry.content).join('');
	}
	return system.content;
}

function makeUserMsg(text: string): AgentMessage {
	return { role: 'user', content: [{ type: 'text', text }] };
}

function makeDbMsg(text: string, createdAt: Date): AgentDbMessage {
	return {
		id: crypto.randomUUID(),
		createdAt,
		role: 'user',
		content: [{ type: 'text', text }],
	};
}

// ---------------------------------------------------------------------------
// Monotonic timestamp assignment
// ---------------------------------------------------------------------------

describe('AgentMessageList — monotonic timestamps', () => {
	it('assigns a Date createdAt to every message added via addInput', () => {
		const list = new AgentMessageList();
		list.addInput([makeUserMsg('hello')]);

		const [msg] = list.turnDelta();
		expect(msg.createdAt).toBeInstanceOf(Date);
	});

	it('assigns strictly increasing createdAt to a batch of input messages', () => {
		const list = new AgentMessageList();
		// Three messages added in the same synchronous call — all would share the
		// same Date.now() tick without monotonic enforcement.
		list.addInput([makeUserMsg('a'), makeUserMsg('b'), makeUserMsg('c')]);

		const [a, b, c] = list.turnDelta();

		expect(a.createdAt).toBeInstanceOf(Date);
		expect(b.createdAt).toBeInstanceOf(Date);
		expect(c.createdAt).toBeInstanceOf(Date);

		expect(b.createdAt.getTime()).toBeGreaterThan(a.createdAt.getTime());
		expect(c.createdAt.getTime()).toBeGreaterThan(b.createdAt.getTime());
	});

	it('assigns strictly increasing createdAt to response messages', () => {
		const list = new AgentMessageList();
		list.addResponse([makeUserMsg('r1'), makeUserMsg('r2'), makeUserMsg('r3')]);

		const [r1, r2, r3] = list.responseDelta();

		expect(r2.createdAt.getTime()).toBeGreaterThan(r1.createdAt.getTime());
		expect(r3.createdAt.getTime()).toBeGreaterThan(r2.createdAt.getTime());
	});

	it('assigns createdAt that is strictly greater than history timestamps', () => {
		const list = new AgentMessageList();

		// Simulate a DB-loaded message with a timestamp in the future relative to
		// wall clock — the new input message must still be later.
		const futureTs = new Date(Date.now() + 60_000);
		list.addHistory([makeDbMsg('old', futureTs)]);

		list.addInput([makeUserMsg('new')]);

		const [inputMsg] = list.turnDelta();
		expect(inputMsg.createdAt).toBeInstanceOf(Date);
		expect(inputMsg.createdAt.getTime()).toBeGreaterThan(futureTs.getTime());
	});
});

// ---------------------------------------------------------------------------
// inputDelta — input-only messages for eager persistence
// ---------------------------------------------------------------------------

describe('AgentMessageList — inputDelta', () => {
	it('returns only input messages, excluding history and responses', () => {
		const list = new AgentMessageList();
		list.addHistory([makeDbMsg('old', new Date('2024-01-01T00:00:00.000Z'))]);
		list.addInput([makeUserMsg('the user prompt')]);
		list.addResponse([makeUserMsg('assistant reply')]);

		const delta = list.inputDelta();

		expect(delta).toHaveLength(1);
		const [text] = (delta[0] as { content: Array<{ type: string; text: string }> }).content;
		expect(text.text).toBe('the user prompt');
	});

	it('returns an empty array when no input was added', () => {
		const list = new AgentMessageList();
		list.addHistory([makeDbMsg('old', new Date('2024-01-01T00:00:00.000Z'))]);

		expect(list.inputDelta()).toEqual([]);
	});
});

// ---------------------------------------------------------------------------
// History messages keep their DB-sourced createdAt
// ---------------------------------------------------------------------------

describe('AgentMessageList — chronological order', () => {
	it('reorders addHistory when the batch is not in createdAt order', () => {
		const list = new AgentMessageList();
		const t1 = new Date('2024-01-01T00:00:01.000Z');
		const t2 = new Date('2024-01-01T00:00:02.000Z');
		list.addHistory([makeDbMsg('second', t2), makeDbMsg('first', t1)]);

		const msgs = list.serialize().messages.filter(isLlmMessage) as Message[];
		expect(msgs).toHaveLength(2);
		expect(msgs[0].content[0]).toMatchObject({ type: 'text', text: 'first' });
		expect(msgs[1].content[0]).toMatchObject({ type: 'text', text: 'second' });
	});
});

// ---------------------------------------------------------------------------
// History messages keep their DB-sourced createdAt
// ---------------------------------------------------------------------------

describe('AgentMessageList — preserving DB timestamps', () => {
	it('preserves the exact createdAt of history messages loaded from the database', () => {
		const list = new AgentMessageList();
		const dbTimestamp = new Date('2020-01-01T00:00:00.000Z');

		list.addHistory([makeDbMsg('from db', dbTimestamp)]);

		const [hist] = list.serialize().messages;
		expect(hist.createdAt).toBeInstanceOf(Date);
		expect(hist.createdAt.getTime()).toBe(dbTimestamp.getTime());
	});
});

// ---------------------------------------------------------------------------
// LLM context assembly
// ---------------------------------------------------------------------------

describe('AgentMessageList — forLlm observation memory', () => {
	it('does not inject a memory section when no observation log has been rendered', () => {
		const list = new AgentMessageList();
		list.addInput([makeUserMsg('hello')]);

		const { system, messages } = list.forLlm('Base instructions');

		expect(flattenSystemContent(system)).toContain('Base instructions');
		expect(flattenSystemContent(system)).not.toContain('## Memory');
		expect(Array.isArray(system)).toBe(false);
		expect(messages).toEqual(
			expect.arrayContaining([
				expect.objectContaining({
					role: 'user',
					content: [expect.objectContaining({ type: 'text', text: 'hello' })],
				}),
			]),
		);
	});

	it('injects the rendered observation log into the system prompt', () => {
		const list = new AgentMessageList();
		list.observationLogMemory = [
			'<observations>',
			'The following is your memory of this conversation.',
			'',
			'* CRITICAL (14:30) User wants the SDK to stay unopinionated.',
			'</observations>',
		].join('\n');

		const { system } = list.forLlm('Base instructions');
		const prompt = flattenSystemContent(system);

		expect(Array.isArray(system)).toBe(true);
		expect(prompt).toContain('Base instructions');
		expect(prompt).toContain('<observations>');
		expect(prompt).toContain('* CRITICAL (14:30) User wants the SDK to stay unopinionated.');
	});

	it('places observation memory in a separate, uncached system message so only instructions carry the cache breakpoint', () => {
		const observationLog = [
			'<observations>',
			'* CRITICAL (14:30) User wants the SDK to stay unopinionated.',
			'</observations>',
		].join('\n');
		const cacheOptions = {
			anthropic: { cacheControl: { type: 'ephemeral' as const } },
		};

		const system = buildSystemMessages('Base instructions', observationLog, cacheOptions);

		expect(Array.isArray(system)).toBe(true);
		if (!Array.isArray(system)) throw new Error('Expected split system messages');
		expect(system).toHaveLength(2);
		expect(system[0]).toEqual({
			role: 'system',
			content: 'Base instructions',
			providerOptions: cacheOptions,
		});
		// The observation message changes on nearly every call, so it must never
		// be marked — doing so would just pay the cache-write premium for no read.
		expect(system[1]).toEqual({
			role: 'system',
			content: expect.stringContaining('<observations>') as unknown as string,
		});
		expect(system[1]?.content).not.toContain('Base instructions');
	});

	it('keeps recent history messages in LLM context when observation memory is empty', () => {
		const list = new AgentMessageList();
		list.addHistory([makeDbMsg('recent history', new Date('2024-01-01T00:00:00.000Z'))]);
		list.addInput([makeUserMsg('new request')]);

		const { messages } = list.forLlm('Base instructions');

		expect(messages).toEqual(
			expect.arrayContaining([
				expect.objectContaining({
					role: 'user',
					content: [expect.objectContaining({ type: 'text', text: 'recent history' })],
				}),
				expect.objectContaining({
					role: 'user',
					content: [expect.objectContaining({ type: 'text', text: 'new request' })],
				}),
			]),
		);
	});
});

describe('buildSystemMessages — volatile tool-instruction fragments', () => {
	it('places volatile instructions in an uncached second message when observation memory is absent', () => {
		const system = buildSystemMessages(
			'Base instructions',
			undefined,
			undefined,
			'<built_in_rules>\n- Newly loaded tool rule.\n</built_in_rules>',
		);

		expect(Array.isArray(system)).toBe(true);
		if (!Array.isArray(system)) throw new Error('Expected split system messages');
		expect(system).toHaveLength(2);
		expect(system[0]).toEqual({ role: 'system', content: 'Base instructions' });
		expect(system[1]?.content).toContain('Newly loaded tool rule.');
		expect(system[1]?.providerOptions).toBeUndefined();
		// The cached instructions message must stay byte-identical regardless of
		// what volatile content exists — this is the whole point of the split.
		expect(system[0]?.content).toBe('Base instructions');
	});

	it('combines volatile instructions and observation memory in the same uncached message', () => {
		const system = buildSystemMessages(
			'Base instructions',
			'<observations>\n* Some memory.\n</observations>',
			undefined,
			'<built_in_rules>\n- Newly loaded tool rule.\n</built_in_rules>',
		);

		expect(Array.isArray(system)).toBe(true);
		if (!Array.isArray(system)) throw new Error('Expected split system messages');
		expect(system).toHaveLength(2);
		expect(system[1]?.content).toContain('Newly loaded tool rule.');
		expect(system[1]?.content).toContain('Some memory.');
	});

	it('keeps the single-message shape when neither observation memory nor volatile instructions are present', () => {
		const system = buildSystemMessages('Base instructions', undefined, undefined, undefined);

		expect(Array.isArray(system)).toBe(false);
		expect(system).toEqual({ role: 'system', content: 'Base instructions' });
	});

	it('folds the mcp-connection note into the uncached volatile message, leaving cached base instructions unchanged', () => {
		const system = buildSystemMessages(
			'Base instructions',
			undefined,
			undefined,
			undefined,
			'<mcp-connection-status>\n- dead: fetch failed\n</mcp-connection-status>',
		);

		expect(Array.isArray(system)).toBe(true);
		if (!Array.isArray(system)) throw new Error('Expected split system messages');
		expect(system).toHaveLength(2);
		// Cached prefix must stay byte-identical regardless of volatile content.
		expect(system[0]?.content).toBe('Base instructions');
		expect(system[1]?.content).toContain('mcp-connection-status');
		expect(system[1]?.content).toContain('dead');
		expect(system[1]?.providerOptions).toBeUndefined();
	});

	it('combines the mcp-connection note with observation memory and volatile instructions in one uncached message', () => {
		const system = buildSystemMessages(
			'Base instructions',
			'<observations>\n* Some memory.\n</observations>',
			undefined,
			'<built_in_rules>\n- Newly loaded tool rule.\n</built_in_rules>',
			'<mcp-connection-status>\n- dead: fetch failed\n</mcp-connection-status>',
		);

		expect(Array.isArray(system)).toBe(true);
		if (!Array.isArray(system)) throw new Error('Expected split system messages');
		expect(system[0]?.content).toBe('Base instructions');
		expect(system[1]?.content).toContain('Newly loaded tool rule.');
		expect(system[1]?.content).toContain('Some memory.');
		expect(system[1]?.content).toContain('mcp-connection-status');
	});

	it('keeps the single-message shape when the mcp-connection note is absent', () => {
		const system = buildSystemMessages(
			'Base instructions',
			undefined,
			undefined,
			undefined,
			undefined,
		);

		expect(Array.isArray(system)).toBe(false);
		expect(system).toEqual({ role: 'system', content: 'Base instructions' });
	});
});

// ---------------------------------------------------------------------------
// Input / response messages use existing createdAt as a hint
// ---------------------------------------------------------------------------

describe('AgentMessageList — hint-based monotonicity for input/response', () => {
	it('keeps an input message createdAt when it is already later than lastCreatedAt', () => {
		const list = new AgentMessageList();
		const histTs = new Date('2024-01-01T00:00:01.000Z');
		list.addHistory([makeDbMsg('hist', histTs)]);

		// freshTs is well after histTs so no bump is needed
		const freshTs = new Date('2024-01-01T00:00:10.000Z');
		const msg = { ...makeUserMsg('new'), createdAt: freshTs };
		list.addInput([msg]);

		const [inputMsg] = list.turnDelta();
		expect(inputMsg.createdAt.getTime()).toBe(freshTs.getTime());
	});

	it('bumps an input message createdAt when it collides with or precedes the last timestamp', () => {
		const list = new AgentMessageList();
		const histTs = new Date('2024-01-01T00:00:05.000Z');
		list.addHistory([makeDbMsg('hist', histTs)]);

		// staleTs is before histTs — must be bumped to at least histTs + 1 ms
		const staleTs = new Date('2024-01-01T00:00:04.000Z');
		const msg = { ...makeUserMsg('stale'), createdAt: staleTs };
		list.addInput([msg]);

		const [inputMsg] = list.turnDelta();
		expect(inputMsg.createdAt.getTime()).toBeGreaterThan(histTs.getTime());
	});

	it('bumps an input message createdAt when it equals the last timestamp', () => {
		const list = new AgentMessageList();
		const ts = new Date('2024-06-01T12:00:00.000Z');
		list.addHistory([makeDbMsg('hist', ts)]);

		// sameTs equals histTs — must be bumped by at least 1 ms
		const msg = { ...makeUserMsg('same'), createdAt: new Date(ts) };
		list.addInput([msg]);

		const [inputMsg] = list.turnDelta();
		expect(inputMsg.createdAt.getTime()).toBeGreaterThan(ts.getTime());
	});
});

// ---------------------------------------------------------------------------
// Deserialization restores lastCreatedAt
// ---------------------------------------------------------------------------

describe('AgentMessageList — deserialize', () => {
	it('messages added after deserialize have createdAt later than any restored message', () => {
		const list = new AgentMessageList();

		// History message with a future timestamp (edge case: e.g. clock skew or
		// the previous turn's monotonic assignment ran ahead of wall clock).
		const futureTs = new Date(Date.now() + 60_000);
		list.addHistory([makeDbMsg('prev', futureTs)]);

		// Round-trip through serialization (simulates suspend / resume)
		const list2 = AgentMessageList.deserialize(list.serialize());
		list2.addInput([makeUserMsg('after resume')]);

		const [newMsg] = list2.turnDelta();
		expect(newMsg.createdAt).toBeInstanceOf(Date);
		expect(newMsg.createdAt.getTime()).toBeGreaterThan(futureTs.getTime());
	});
});

// ---------------------------------------------------------------------------
// setToolCallResult / setToolCallError
// ---------------------------------------------------------------------------

function makePendingToolCallMsg(toolCallId: string): AgentMessage {
	return {
		role: 'assistant',
		content: [
			{
				type: 'tool-call',
				toolCallId,
				toolName: 'my_tool',
				input: { x: 1 },
				state: 'pending',
			},
		],
	};
}

describe('AgentMessageList — setToolCallResult', () => {
	it('sets state and output on the matching tool-call block', () => {
		const list = new AgentMessageList();
		list.addResponse([makePendingToolCallMsg('id-1')]);

		const host = list.setToolCallResult('id-1', { ok: true });
		expect(host).toBeDefined();

		const block = (host as Message).content.find((c) => c.type === 'tool-call') as ContentToolCall;
		expect(block.state).toBe('resolved');
		expect((block as ContentToolCall & { state: 'resolved' }).output).toEqual({ ok: true });
	});

	it('promotes a history-only message into responseDelta after setToolCallResult', () => {
		const list = new AgentMessageList();
		const histMsg: AgentDbMessage = {
			id: 'hist-1',
			createdAt: new Date('2024-01-01T00:00:01.000Z'),
			role: 'assistant',
			content: [
				{
					type: 'tool-call',
					toolCallId: 'tc-hist',
					toolName: 'my_tool',
					input: {},
					state: 'pending',
				},
			],
		};
		list.addHistory([histMsg]);

		// Before: not in responseDelta (history only)
		expect(list.responseDelta()).toHaveLength(0);

		list.setToolCallResult('tc-hist', { done: true });

		// After: promoted to responseDelta
		const delta = list.responseDelta();
		expect(delta).toHaveLength(1);
		const block = (delta[0] as Message).content.find(
			(c) => c.type === 'tool-call',
		) as ContentToolCall;
		expect(block.state).toBe('resolved');
	});

	it('is a no-op when toolCallId is unknown', () => {
		const list = new AgentMessageList();
		list.addResponse([makePendingToolCallMsg('id-1')]);

		const result = list.setToolCallResult('unknown-id', { x: 1 });
		expect(result).toBeUndefined();
		// List unchanged
		expect(list.responseDelta()).toHaveLength(1);
	});

	it('Set semantics make repeated calls idempotent (no duplicate messages)', () => {
		const list = new AgentMessageList();
		list.addResponse([makePendingToolCallMsg('id-1')]);

		list.setToolCallResult('id-1', { ok: true });
		list.setToolCallResult('id-1', { ok: true });

		expect(list.responseDelta()).toHaveLength(1);
	});
});

describe('AgentMessageList — setToolCallError', () => {
	it('stringifies errors and clears any prior output', () => {
		const list = new AgentMessageList();
		list.addResponse([
			{
				role: 'assistant',
				content: [
					{
						type: 'tool-call',
						toolCallId: 'id-1',
						toolName: 'my_tool',
						input: {},
						state: 'resolved',
						output: { prev: true },
					},
				],
			},
		]);

		const host = list.setToolCallError('id-1', new Error('boom'));
		expect(host).toBeDefined();

		const block = (host as Message).content.find((c) => c.type === 'tool-call') as ContentToolCall;
		expect(block.state).toBe('rejected');
		expect((block as ContentToolCall & { state: 'rejected' }).error).toBe('Error: boom');
		// output should be gone
		expect((block as unknown as { output?: unknown }).output).toBeUndefined();
	});
});

describe('AgentMessageList — markToolCallSuspended', () => {
	it('stamps suspension info on a pending tool-call block', () => {
		const list = new AgentMessageList();
		list.addResponse([makePendingToolCallMsg('id-1')]);

		list.markToolCallSuspended('id-1', { message: 'Edit My Workflow (ID: abc)?', requestId: 'r1' });

		const host = list
			.turnDelta()
			.find((m) => 'content' in m && m.content.some((c) => c.type === 'tool-call')) as Message;
		const block = host.content.find((c) => c.type === 'tool-call') as ContentToolCall & {
			state: 'pending';
		};
		expect(block.state).toBe('pending');
		expect(block.suspension).toEqual({ message: 'Edit My Workflow (ID: abc)?', requestId: 'r1' });
	});

	it('is a no-op for settled blocks and unknown ids', () => {
		const list = new AgentMessageList();
		list.addResponse([makePendingToolCallMsg('id-1')]);
		list.setToolCallResult('id-1', { ok: true });

		list.markToolCallSuspended('id-1', { message: 'stale' });
		list.markToolCallSuspended('missing', { message: 'stale' });

		const host = list
			.turnDelta()
			.find((m) => 'content' in m && m.content.some((c) => c.type === 'tool-call')) as Message;
		const block = host.content.find((c) => c.type === 'tool-call') as ContentToolCall;
		expect(block.state).toBe('resolved');
		expect('suspension' in block).toBe(false);
	});
});

// ---------------------------------------------------------------------------
// Observation masking (mid-run compaction)
// ---------------------------------------------------------------------------

describe('AgentMessageList — observation masking', () => {
	const at = (second: number) => new Date(Date.UTC(2026, 0, 1, 12, 0, second));

	function dbMsgWithId(id: string, text: string, createdAt: Date): AgentDbMessage {
		return { id, createdAt, role: 'user', content: [{ type: 'text', text }] };
	}

	it('hides messages at or before the cursor from forLlm while deltas and serialize keep everything', () => {
		const list = new AgentMessageList();
		list.addHistory([dbMsgWithId('m1', 'observed history', at(1))]);
		list.addInput([
			{
				id: 'm2',
				createdAt: at(2),
				role: 'user',
				content: [{ type: 'text', text: 'visible input' }],
			},
		]);
		list.addResponse([
			{
				id: 'm3',
				createdAt: at(3),
				role: 'assistant',
				content: [{ type: 'text', text: 'fresh reply' }],
			},
		]);

		list.maskObservedMessages({ lastObservedAt: at(1), lastObservedMessageId: 'm1' });

		const serialized = JSON.stringify(list.forLlm('base').messages);
		expect(serialized).not.toContain('observed history');
		expect(serialized).toContain('visible input');
		expect(serialized).toContain('fresh reply');

		// Persistence and result consumers are mask-agnostic.
		expect(list.turnDelta().map((m) => m.id)).toEqual(['m2', 'm3']);
		expect(list.responseDelta().map((m) => m.id)).toEqual(['m3']);
		expect(list.serialize().messages.map((m) => m.id)).toEqual(['m1', 'm2', 'm3']);
	});

	it('splits messages sharing the cursor createdAt by id, matching store keyset semantics', () => {
		const list = new AgentMessageList();
		const ts = at(5);
		list.addHistory([
			dbMsgWithId('a', 'text a', ts),
			dbMsgWithId('b', 'text b', ts),
			dbMsgWithId('c', 'text c', ts),
		]);

		list.maskObservedMessages({ lastObservedAt: ts, lastObservedMessageId: 'b' });

		const serialized = JSON.stringify(list.forLlm('base').messages);
		expect(serialized).not.toContain('text a');
		expect(serialized).not.toContain('text b');
		expect(serialized).toContain('text c');
	});

	it('masks from a string-dated cursor and ignores an unparseable one', () => {
		// Store adapters are an open interface — a JSON-backed one hands back
		// lastObservedAt as an ISO string instead of a Date.
		const list = new AgentMessageList();
		list.addHistory([
			dbMsgWithId('m1', 'observed history', at(1)),
			dbMsgWithId('m2', 'still visible', at(2)),
		]);

		list.maskObservedMessages({
			lastObservedAt: at(1).toISOString() as unknown as Date,
			lastObservedMessageId: 'm1',
		});
		expect(list.llmVisibleMessages().map((m) => m.id)).toEqual(['m2']);

		// An unusable date would compare as NaN and mask every message — the
		// boundary must stay where it was instead.
		list.maskObservedMessages({
			lastObservedAt: 'not-a-date' as unknown as Date,
			lastObservedMessageId: 'm2',
		});
		expect(list.llmVisibleMessages().map((m) => m.id)).toEqual(['m2']);
	});

	it('sends exactly one continuation reminder when the whole window is masked', () => {
		const list = new AgentMessageList();
		list.addHistory([
			dbMsgWithId('m1', 'observed one', at(1)),
			dbMsgWithId('m2', 'observed two', at(2)),
		]);

		list.maskObservedMessages({ lastObservedAt: at(2), lastObservedMessageId: 'm2' });

		expect(list.forLlm('base').messages).toEqual([
			{ role: 'user', content: OBSERVATION_CONTINUATION_REMINDER },
		]);
	});

	it('prepends the reminder before a visible assistant tool-call host and keeps its result pair intact', () => {
		const list = new AgentMessageList();
		list.addHistory([
			dbMsgWithId('m1', 'earlier work', at(1)),
			{
				id: 'm2',
				createdAt: at(2),
				role: 'assistant',
				content: [
					{
						type: 'tool-call',
						toolCallId: 'tc1',
						toolName: 'my_tool',
						input: { x: 1 },
						state: 'resolved',
						output: { ok: true },
					},
				],
			},
		]);

		list.maskObservedMessages({ lastObservedAt: at(1), lastObservedMessageId: 'm1' });

		const { messages } = list.forLlm('base');
		expect(messages[0]).toEqual({ role: 'user', content: OBSERVATION_CONTINUATION_REMINDER });
		// The host explodes into assistant tool-call + tool result — masking must
		// never orphan one half of the pair.
		expect(messages.map((m) => m.role)).toEqual(['user', 'assistant', 'tool']);
	});
});