import turnAbortedGuidance from "../prompts/turn-aborted-guidance.md" with { type: "text" };
import type {
	Api,
	AssistantMessage,
	DeveloperMessage,
	Message,
	Model,
	ToolCall,
	ToolResultMessage,
	UserMessage,
} from "../types";

const enum ToolCallStatus {
	/** Tool call has received a result (real or synthetic for orphan) */
	Resolved = 1,
	/** Tool call was from an aborted message; synthetic result injected, skip real results */
	Aborted = 2,
}

/**
 * Normalize tool call ID for cross-provider compatibility.
 * OpenAI Responses API generates IDs that are 450+ chars with special characters like `|`.
 * Anthropic APIs require IDs matching ^[a-zA-Z0-9_-]+$ (max 64 chars).
 *
 * For aborted/errored turns, this function:
 * - Preserves tool call structure (unlike converting to text summaries)
 * - Injects synthetic "aborted" tool results
 * - Adds a <turn-aborted> guidance marker for the model
 */
/**
 * Detect directly adjacent private thinking blocks inside one assistant message's
 * content. `thinking` and `redacted_thinking` are one adjacency class: the
 * Anthropic wire contract rejects a replayed assistant turn where two such blocks
 * sit next to each other with no intervening `tool_use`/`text` block (#4416).
 *
 * This is a pure, allocation-free predicate used by defense-in-depth diagnostics
 * (issue #4443): the write-time transcript assertion (coding-agent persistence)
 * and the stream-assembler SSE diagnostic (anthropic stream completion). It never
 * inspects block payloads — only the block-type sequence — so it cannot leak
 * thinking text, signatures, or credentials.
 *
 * Blocks separated by any non-private block (`tool_use`, `text`, …) are ordinary
 * interleaved-thinking shape and return `false`.
 */
export function hasAdjacentPrivateThinkingBlocks(content: { type: string }[]): boolean {
	let previousWasPrivate = false;
	for (const block of content) {
		const isPrivate = block.type === "thinking" || block.type === "redactedThinking";
		if (isPrivate && previousWasPrivate) return true;
		previousWasPrivate = isPrivate;
	}
	return false;
}

/**
 * Collapse a run of directly adjacent `thinking` blocks inside one assistant message down
 * to its first block.
 *
 * Anthropic accepts a replayed assistant turn carrying a single thinking block, and accepts
 * thinking blocks separated by a `tool_use` (ordinary interleaved-thinking shape), but
 * rejects two directly adjacent `thinking` blocks with
 * `messages.N.content.M: thinking or redacted_thinking blocks in the latest assistant
 * message cannot be modified`, citing the *second* block of the pair. Because the offending
 * message keeps its index as history grows, a single such turn makes every later request in
 * that session fail, and the mutation repair - scoped to the latest assistant message -
 * can never reach it (#4416).
 *
 * `redactedThinking` is not folded in this phase, but the final send-boundary
 * collapse in `convertAnthropicMessages` (#4425) treats `thinking` and
 * `redacted_thinking` as one adjacency class per the API contract.
 */
function collapseAdjacentThinking<T extends { type: string }>(content: T[]): T[] {
	let previousWasThinking = false;
	let dropped = false;
	const collapsed: T[] = [];
	for (const block of content) {
		const thinking = block.type === "thinking";
		if (thinking && previousWasThinking) {
			dropped = true;
			continue;
		}
		previousWasThinking = thinking;
		collapsed.push(block);
	}
	return dropped ? collapsed : content;
}

export function transformMessages<TApi extends Api>(
	messages: Message[],
	model: Model<TApi>,
	normalizeToolCallId?: (id: string, model: Model<TApi>, source: AssistantMessage) => string,
	options?: { repairLatestAssistantThinking?: boolean; repairAllAssistantThinking?: boolean },
): Message[] {
	// Build a map of original tool call IDs to normalized IDs
	const toolCallIdMap = new Map<string, string>();

	const latestAssistantIndex = messages.findLastIndex(msg => msg.role === "assistant");
	// First pass: transform messages (thinking blocks, tool call ID normalization)
	const transformed = messages.map((msg, index) => {
		// User and developer messages pass through unchanged
		if (msg.role === "user" || msg.role === "developer") {
			return msg;
		}

		// Handle toolResult messages - normalize toolCallId if we have a mapping
		if (msg.role === "toolResult") {
			const normalizedId = toolCallIdMap.get(msg.toolCallId);
			if (normalizedId && normalizedId !== msg.toolCallId) {
				return { ...msg, toolCallId: normalizedId };
			}
			return msg;
		}

		// Assistant messages need transformation check
		if (msg.role === "assistant") {
			const assistantMsg = msg as AssistantMessage;
			const isSameModel =
				assistantMsg.provider === model.provider &&
				assistantMsg.api === model.api &&
				assistantMsg.model === model.id;

			const mustPreserveLatestAnthropicThinking =
				index === latestAssistantIndex &&
				model.api === "anthropic-messages" &&
				assistantMsg.api === "anthropic-messages";
			// Aborted/errored messages may contain partially-streamed thinking blocks.
			// Anthropic requires thinking/redacted_thinking bytes in replayed assistant
			// messages to match the original response exactly; stripping a signature,
			// well-forming text, or keeping a partial redacted block would emit a
			// modified thinking sequence. Drop those private blocks instead. Tool calls
			// are kept so the second pass can either preserve real results or synthesize
			// an explicit aborted result without leaving dangling tool_use blocks.
			const hasPartialThinking = assistantMsg.stopReason === "aborted" || assistantMsg.stopReason === "error";
			// One-shot Anthropic replay repair. `repairLatestAssistantThinking` targets the
			// "latest assistant message ... cannot be modified" 400; `repairAllAssistantThinking`
			// targets the "Invalid `signature` in `thinking` block" 400, which can cite a block
			// anywhere in the replayed history (e.g. after compaction/pruning rewrote an earlier
			// turn), so the drop must apply to every assistant message. Within each
			// message only blocks that would replay as native thinking/redacted_thinking
			// are dropped; cross-model reasoning degrades to text and is preserved.
			const dropAssistantThinkingForRepair =
				(options?.repairAllAssistantThinking === true ||
					(options?.repairLatestAssistantThinking === true && index === latestAssistantIndex)) &&
				model.api === "anthropic-messages" &&
				assistantMsg.api === "anthropic-messages";

			const transformedContent = assistantMsg.content.flatMap(block => {
				if (block.type === "thinking") {
					if (hasPartialThinking) return [];
					const sanitized = block;
					// Repair must only drop blocks that would otherwise replay as native
					// thinking. Cross-model/provider reasoning degrades to unsigned text
					// below and was never replayed as a signed block, so it cannot be the
					// signature failure — dropping it would silently lose valid context.
					const replaysAsNativeThinking = mustPreserveLatestAnthropicThinking || isSameModel;
					if (dropAssistantThinkingForRepair && replaysAsNativeThinking) return [];
					if (mustPreserveLatestAnthropicThinking) return sanitized;
					// For same model: keep thinking blocks with signatures (needed for replay)
					// even if the thinking text is empty — but only for non-Anthropic APIs where
					// the signature represents OpenAI encrypted reasoning. For anthropic-messages,
					// a signed block with empty text means clear_thinking_20251015 stripped the
					// content server-side while the stale signature remained; replaying it
					// produces `thinking ... cannot be modified` 400s on every turn (#4247).
					if (isSameModel && sanitized.thinkingSignature) {
						if (sanitized.thinking.trim() === "" && model.api === "anthropic-messages") return [];
						return sanitized;
					}
					// Skip empty thinking blocks, convert others to plain text
					if (!sanitized.thinking || sanitized.thinking.trim() === "") return [];
					if (isSameModel) return sanitized;
					return {
						type: "text" as const,
						text: sanitized.thinking,
					};
				}

				if (block.type === "redactedThinking") {
					if (hasPartialThinking) return [];
					// Same restriction as thinking blocks: cross-model/provider redacted
					// blocks already drop below, so repair only needs to cover blocks that
					// would replay as native redacted_thinking.
					if (dropAssistantThinkingForRepair && (mustPreserveLatestAnthropicThinking || isSameModel)) {
						return [];
					}
					if (mustPreserveLatestAnthropicThinking) return block;
					if (isSameModel) return block;
					return [];
				}

				if (block.type === "text") {
					if (isSameModel) return block;
					return {
						type: "text" as const,
						text: block.text,
					};
				}

				if (block.type === "toolCall") {
					const toolCall = block as ToolCall;
					let normalizedToolCall: ToolCall = toolCall;

					if (!isSameModel && toolCall.thoughtSignature) {
						normalizedToolCall = { ...toolCall };
						delete (normalizedToolCall as { thoughtSignature?: string }).thoughtSignature;
					}

					if (!isSameModel && normalizeToolCallId) {
						const normalizedId = normalizeToolCallId(toolCall.id, model, assistantMsg);
						if (normalizedId !== toolCall.id) {
							toolCallIdMap.set(toolCall.id, normalizedId);
							normalizedToolCall = { ...normalizedToolCall, id: normalizedId };
						}
					}

					return normalizedToolCall;
				}

				return block;
			});

			// Only the Anthropic wire shape rejects adjacent private blocks; other targets
			// either degrade reasoning to text above or carry their own encoding rules.
			const replayableContent =
				model.api === "anthropic-messages" ? collapseAdjacentThinking(transformedContent) : transformedContent;

			return {
				...assistantMsg,
				content: replayableContent,
			};
		}
		return msg;
	});
	const realToolResultIds = new Set(
		transformed.filter((msg): msg is ToolResultMessage => msg.role === "toolResult").map(msg => msg.toolCallId),
	);

	// Anthropic rejects `tool_result` blocks whose `tool_use_id` does not appear in a prior
	// `tool_use` block. After handoff/compaction folds an assistant turn into a summary
	// string, the user-side `toolResult` for that turn can survive while the originating
	// `tool_use` disappears — leaving an orphan that triggers HTTP 400. Track the set of
	// `tool_use` ids that survive transformation so the second pass can drop orphans cleanly.
	const validToolUseIds = new Set<string>();
	for (const msg of transformed) {
		if (msg.role !== "assistant") continue;
		for (const block of msg.content) {
			if (block.type === "toolCall") validToolUseIds.add(block.id);
		}
	}

	// Second pass: insert synthetic empty tool results for orphaned tool calls
	// and preserve aborted/errored tool results when they were already persisted.
	const result: Message[] = [];
	let pendingToolCalls: ToolCall[] = [];
	let pendingAbortedToolCalls = new Map<string, ToolCall>();
	let pendingAbortedTimestamp: number | undefined;
	// Track tool call status: whether resolved (has result) or aborted (synthetic result injected, skip later real results)
	const toolCallStatus = new Map<string, ToolCallStatus>();

	const flushPendingToolCalls = (timestamp: number): void => {
		if (pendingToolCalls.length === 0) return;
		for (const tc of pendingToolCalls) {
			if (!toolCallStatus.has(tc.id) && !realToolResultIds.has(tc.id)) {
				result.push({
					role: "toolResult",
					toolCallId: tc.id,
					toolName: tc.name,
					content: [{ type: "text", text: "No result provided" }],
					isError: true,
					timestamp,
				} as ToolResultMessage);
				toolCallStatus.set(tc.id, ToolCallStatus.Resolved);
			}
		}
		pendingToolCalls = [];
	};

	const flushPendingAbortedToolCalls = (): void => {
		if (pendingAbortedTimestamp === undefined) return;
		for (const tc of pendingAbortedToolCalls.values()) {
			if (!toolCallStatus.has(tc.id)) {
				result.push({
					role: "toolResult",
					toolCallId: tc.id,
					toolName: tc.name,
					content: [{ type: "text", text: "aborted" }],
					isError: true,
					timestamp: pendingAbortedTimestamp,
				} as ToolResultMessage);
				toolCallStatus.set(tc.id, ToolCallStatus.Aborted);
			}
		}
		result.push({
			role: "developer",
			content: turnAbortedGuidance,
			timestamp: pendingAbortedTimestamp + 1,
		} as DeveloperMessage);
		pendingAbortedToolCalls = new Map();
		pendingAbortedTimestamp = undefined;
	};

	for (let i = 0; i < transformed.length; i++) {
		const msg = transformed[i];
		const messageTimestamp = "timestamp" in msg && typeof msg.timestamp === "number" ? msg.timestamp : Date.now();

		if (msg.role === "assistant") {
			flushPendingToolCalls(messageTimestamp);
			flushPendingAbortedToolCalls();

			const assistantMsg = msg as AssistantMessage;
			const toolCalls = assistantMsg.content.filter(b => b.type === "toolCall") as ToolCall[];

			if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") {
				// Keep the assistant message with tool calls intact. If real tool results follow, preserve them;
				// otherwise synthesize aborted results before the next turn boundary.
				result.push(msg);
				pendingAbortedToolCalls = new Map(toolCalls.map(toolCall => [toolCall.id, toolCall] as const));
				pendingAbortedTimestamp = assistantMsg.timestamp;
				continue;
			}

			if (toolCalls.length > 0) {
				pendingToolCalls = toolCalls;
			}

			result.push(msg);
		} else if (msg.role === "toolResult") {
			if (pendingAbortedToolCalls.has(msg.toolCallId)) {
				pendingAbortedToolCalls.delete(msg.toolCallId);
				toolCallStatus.set(msg.toolCallId, ToolCallStatus.Resolved);
				result.push(msg);
				continue;
			}

			if (toolCallStatus.get(msg.toolCallId) === ToolCallStatus.Aborted) continue;

			if (!validToolUseIds.has(msg.toolCallId)) {
				// Orphan `tool_result`: the originating `tool_use` is not present in the
				// transformed history (typically because handoff/compaction folded the
				// assistant message into a summary string while the user-side result
				// survived). Sending the block as-is would 400 the request, so it must
				// be dropped.
				//
				// If a pending tool-call window is still open (either normal or
				// aborted), the orphan cannot be replaced with a developer note here:
				//
				// * Anthropic requires the next message after an assistant `tool_use`
				//   to be the matching `tool_result`. Inserting a developer message
				//   would break that contiguity.
				// * `flushPendingAbortedToolCalls` synthesizes "aborted" results
				//   without checking whether a real result lands later in history
				//   (unlike `flushPendingToolCalls`, which is gated by
				//   `realToolResultIds`). Calling it here would convert a legitimate
				//   later `tool_result` into a synthetic "aborted" one via the
				//   `ToolCallStatus.Aborted` skip-guard.
				//
				// Drop the orphan silently in that case; the upcoming real
				// `tool_result` will land normally on the next iteration.
				if (pendingToolCalls.length > 0 || pendingAbortedToolCalls.size > 0) {
					continue;
				}
				// No pending tool-call window: safe to preserve the text payload so the
				// model still sees what the tool returned.
				//
				// The note is emitted with `role: "user"` rather than `role: "developer"`
				// because the developer role is elevated by some providers:
				//
				// * Ollama maps `developer` -> `system` (highest instruction priority).
				// * OpenAI chat-completions reasoning models forward `developer` as
				//   `developer` (above-user instruction priority).
				//
				// Stale, model-untrusted tool output must not gain instruction priority
				// above user/developer messages it lived alongside before compaction.
				// `user` role is mapped to plain user content by every provider, so the
				// content survives without ever being treated as an instruction the
				// model should obey.
				const textParts: string[] = [];
				for (const part of msg.content) {
					if (part.type === "text" && part.text.trim() !== "") textParts.push(part.text);
				}
				if (textParts.length > 0) {
					const errorAttr = msg.isError ? ' is-error="true"' : "";
					result.push({
						role: "user",
						content: `<stale-tool-result tool="${msg.toolName}" id="${msg.toolCallId}"${errorAttr}>\n${textParts.join("\n")}\n</stale-tool-result>`,
						timestamp: messageTimestamp,
					} as UserMessage);
				}
				continue;
			}

			toolCallStatus.set(msg.toolCallId, ToolCallStatus.Resolved);
			result.push(msg);
		} else if (msg.role === "user" || msg.role === "developer") {
			flushPendingToolCalls(messageTimestamp);
			flushPendingAbortedToolCalls();
			result.push(msg);
		} else {
			flushPendingToolCalls(messageTimestamp);
			flushPendingAbortedToolCalls();
			result.push(msg);
		}
	}

	flushPendingToolCalls(Date.now());
	flushPendingAbortedToolCalls();

	return result;
}