import type { AgentMessage } from "@singularity-forge/pi-agent-core";

export const DEFAULT_COMPLETION_NUDGE_AFTER = 10;
export const COMPLETION_NUDGE_TOOL_NAMES = new Set([
	"sf_complete_slice",
	"sf_slice_complete",
]);

const COMPLETION_NUDGE_CUSTOM_TYPE = "sf-completion-nudge";
const LOWERED_TEMPERATURE = 0.2;

interface CompletionNudgeState {
	active: boolean;
	unitType: string;
	unitId: string;
	toolCalls: number;
	completionCalled: boolean;
	nudgeAfter: number;
	reminderSent: boolean;
	strongSent: boolean;
	lowerTemperatureForNextRequest: boolean;
}

const state: CompletionNudgeState = {
	active: false,
	unitType: "",
	unitId: "",
	toolCalls: 0,
	completionCalled: false,
	nudgeAfter: DEFAULT_COMPLETION_NUDGE_AFTER,
	reminderSent: false,
	strongSent: false,
	lowerTemperatureForNextRequest: false,
};

export function resolveCompletionNudgeAfter(value: unknown): number {
	const n = Number(value);
	if (!Number.isFinite(n)) return DEFAULT_COMPLETION_NUDGE_AFTER;
	return Math.max(0, Math.floor(n));
}

export function resetCompletionNudgeState(
	unitType: string,
	unitId: string,
	configuredNudgeAfter?: unknown,
): void {
	const nudgeAfter = resolveCompletionNudgeAfter(configuredNudgeAfter);
	state.active = unitType === "complete-slice" && nudgeAfter > 0;
	state.unitType = unitType;
	state.unitId = unitId;
	state.toolCalls = 0;
	state.completionCalled = false;
	state.nudgeAfter = nudgeAfter;
	state.reminderSent = false;
	state.strongSent = false;
	state.lowerTemperatureForNextRequest = false;
}

export function clearCompletionNudgeState(): void {
	resetCompletionNudgeState("", "", DEFAULT_COMPLETION_NUDGE_AFTER);
}

export function recordCompletionNudgeToolCall(toolName: string): void {
	if (!state.active) return;
	if (COMPLETION_NUDGE_TOOL_NAMES.has(toolName)) {
		state.completionCalled = true;
		state.lowerTemperatureForNextRequest = false;
		return;
	}
	state.toolCalls++;
}

export function getCompletionNudgeStateForTest(): Readonly<CompletionNudgeState> {
	return { ...state };
}

export function maybeInjectCompletionNudgeMessage(
	messages: AgentMessage[],
): AgentMessage[] {
	if (!state.active || state.completionCalled) return messages;

	const message = nextCompletionNudgeMessage();
	if (!message) return messages;

	return [
		...messages,
		{
			role: "custom",
			customType: COMPLETION_NUDGE_CUSTOM_TYPE,
			content: message,
			display: false,
			details: {
				unitType: state.unitType,
				unitId: state.unitId,
				toolCalls: state.toolCalls,
			},
			timestamp: Date.now(),
		},
	];
}

export function applyCompletionNudgeTemperature(payload: unknown): unknown {
	if (!state.lowerTemperatureForNextRequest || state.completionCalled)
		return payload;
	state.lowerTemperatureForNextRequest = false;
	if (!payload || typeof payload !== "object") return payload;

	lowerTemperature(payload as Record<string, unknown>);
	return payload;
}

function nextCompletionNudgeMessage(): string | null {
	const firstThreshold = state.nudgeAfter;
	const secondThreshold = state.nudgeAfter * 2;

	if (!state.strongSent && state.toolCalls >= secondThreshold) {
		state.reminderSent = true;
		state.strongSent = true;
		state.lowerTemperatureForNextRequest = true;
		return `You've performed ${state.toolCalls} tool calls without calling sf_complete_slice. Stop further investigation unless there is a specific blocker. Call sf_complete_slice now with your summary.`;
	}

	if (!state.reminderSent && state.toolCalls >= firstThreshold) {
		state.reminderSent = true;
		return `You've performed ${state.toolCalls} tool calls of investigation. Per the slice plan you should now call sf_complete_slice with your summary. If you genuinely need more context, say so explicitly; otherwise call the tool now.`;
	}

	return null;
}

function lowerTemperature(record: Record<string, unknown>): void {
	record.temperature = lowerNumber(record.temperature);

	const generationConfig = record.generationConfig;
	if (generationConfig && typeof generationConfig === "object") {
		(generationConfig as Record<string, unknown>).temperature = lowerNumber(
			(generationConfig as Record<string, unknown>).temperature,
		);
	}

	const config = record.config;
	if (config && typeof config === "object") {
		(config as Record<string, unknown>).temperature = lowerNumber(
			(config as Record<string, unknown>).temperature,
		);
	}
}

function lowerNumber(value: unknown): number {
	return typeof value === "number"
		? Math.min(value, LOWERED_TEMPERATURE)
		: LOWERED_TEMPERATURE;
}
