import assert from "node:assert/strict";
import { test } from 'vitest';
import type { AgentMessage } from "@singularity-forge/pi-agent-core";

import {
	applyCompletionNudgeTemperature,
	getCompletionNudgeStateForTest,
	maybeInjectCompletionNudgeMessage,
	recordCompletionNudgeToolCall,
	resetCompletionNudgeState,
} from "../auto-completion-nudge.ts";

type CustomTestMessage = Extract<AgentMessage, { role: "custom" }> & {
	content: string;
};

function baseMessages(): AgentMessage[] {
	return [
		{
			role: "user",
			content: [{ type: "text", text: "complete slice" }],
			timestamp: Date.now(),
		},
	];
}

test("completion nudge injects reminder after N tool calls without sf_complete_slice", () => {
	resetCompletionNudgeState("complete-slice", "M001/S01", 3);
	recordCompletionNudgeToolCall("read");
	recordCompletionNudgeToolCall("grep");
	assert.equal(maybeInjectCompletionNudgeMessage(baseMessages()).length, 1);

	recordCompletionNudgeToolCall("bash");
	const messages = maybeInjectCompletionNudgeMessage(baseMessages());

	assert.equal(messages.length, 2);
	const nudge = messages[1];
	assert.equal(nudge.role, "custom");
	assert.match(
		String(nudge.content),
		/You've performed 3 tool calls of investigation/,
	);
	assert.match(String(nudge.content), /call sf_complete_slice/);
});

test("completion nudge does not inject after sf_complete_slice is called", () => {
	resetCompletionNudgeState("complete-slice", "M001/S01", 3);
	recordCompletionNudgeToolCall("read");
	recordCompletionNudgeToolCall("grep");
	recordCompletionNudgeToolCall("sf_complete_slice");
	recordCompletionNudgeToolCall("bash");

	const messages = maybeInjectCompletionNudgeMessage(baseMessages());
	assert.equal(messages.length, 1);
	assert.equal(getCompletionNudgeStateForTest().completionCalled, true);
});

test("completion nudge injects stronger message at 2N and lowers temperature", () => {
	resetCompletionNudgeState("complete-slice", "M001/S01", 2);
	recordCompletionNudgeToolCall("read");
	recordCompletionNudgeToolCall("grep");
	maybeInjectCompletionNudgeMessage(baseMessages());

	recordCompletionNudgeToolCall("bash");
	recordCompletionNudgeToolCall("read");
	const messages = maybeInjectCompletionNudgeMessage(baseMessages());
	assert.equal(messages.length, 2);
	const nudge = messages[1] as CustomTestMessage;
	assert.match(String(nudge.content), /without calling sf_complete_slice/);

	const payload = { temperature: 0.9, generationConfig: { temperature: 0.8 } };
	applyCompletionNudgeTemperature(payload);
	assert.equal(payload.temperature, 0.2);
	assert.equal(payload.generationConfig.temperature, 0.2);
});
