import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { SessionState } from "../state/types.ts";

/**
 * Filter out compressed message ranges and inject summaries.
 * Messages covered by active blocks are removed and replaced with
 * a synthetic user message containing the summary at the anchor position.
 */
export function filterCompressedRanges(
  state: SessionState,
  messages: AgentMessage[],
): AgentMessage[] {
  if (state.prune.messages.activeBlockIds.size === 0) return messages;

  const result: AgentMessage[] = [];

  for (let i = 0; i < messages.length; i++) {
    // Check if there's a summary to inject at this anchor point
    const blockId = state.prune.messages.activeByAnchorIndex.get(i);
    if (blockId !== undefined) {
      const block = state.prune.messages.blocksById.get(blockId);
      if (block?.active && block.summary) {
        result.push({
          role: "user",
          content: [{ type: "text", text: block.summary }],
          timestamp: Date.now(),
        } as AgentMessage);
      }
    }

    // Skip messages that are covered by active blocks
    const entry = state.prune.messages.byMessageIndex.get(i);
    if (entry && entry.activeBlockIds.length > 0) {
      continue;
    }

    result.push(messages[i]);
  }

  // Safety net: remove orphaned toolResult messages
  return removeOrphanedToolResults(result);
}

/**
 * Remove toolResult messages whose toolCallId has no matching toolCall
 * in an assistant message in the output array.
 *
 * Known limitation: does not detect the reverse — an assistant message with
 * toolCall content whose toolResult was removed. That case is prevented by
 * Layer 1 (expandRangeForToolChains) but not caught here because handling it
 * is harder (assistant messages may have mixed content). A future hardening
 * pass could synthesize stub toolResults for orphaned toolCalls.
 */
function removeOrphanedToolResults(messages: AgentMessage[]): AgentMessage[] {
  // Collect all toolCall IDs from assistant messages
  const toolCallIds = new Set<string>();
  for (const msg of messages) {
    if (msg.role !== "assistant" || !Array.isArray(msg.content)) continue;
    for (const part of msg.content) {
      if (typeof part !== "object" || part === null) continue;
      const p = part as unknown as Record<string, unknown>;
      if (p.type === "toolCall" && typeof p.id === "string") {
        toolCallIds.add(p.id as string);
      }
    }
  }

  // Filter out toolResult messages without a matching toolCall
  return messages.filter((msg) => {
    if (msg.role !== "toolResult") return true;
    return toolCallIds.has((msg as unknown as { toolCallId: string }).toolCallId);
  });
}

const PRUNED_OUTPUT_TEXT =
  "[Output removed to save context - information superseded or no longer needed]";
const PRUNED_ERROR_INPUT_TEXT = "[input removed due to failed tool call]";

/**
 * Replace outputs of pruned tool results with placeholder text.
 * Returns a new array (does not mutate input).
 */
export function pruneToolOutputs(
  state: SessionState,
  messages: AgentMessage[],
): AgentMessage[] {
  if (state.prune.tools.size === 0) return messages;

  return messages.map((msg) => {
    if (msg.role !== "toolResult") return msg;
    if (!state.prune.tools.has(msg.toolCallId)) return msg;
    if (msg.isError) return msg;

    return {
      ...msg,
      content: [{ type: "text" as const, text: PRUNED_OUTPUT_TEXT }],
    };
  });
}

/**
 * Replace content of pruned error tool results with placeholder text.
 * Returns a new array (does not mutate input).
 */
export function pruneToolErrors(
  state: SessionState,
  messages: AgentMessage[],
): AgentMessage[] {
  if (state.prune.tools.size === 0) return messages;

  return messages.map((msg) => {
    if (msg.role !== "toolResult") return msg;
    if (!state.prune.tools.has(msg.toolCallId)) return msg;
    if (!msg.isError) return msg;

    return {
      ...msg,
      content: [{ type: "text" as const, text: PRUNED_ERROR_INPUT_TEXT }],
    };
  });
}

/**
 * Apply all pruning passes to a message array.
 * Returns a new array.
 */
export function applyPruning(
  state: SessionState,
  messages: AgentMessage[],
): AgentMessage[] {
  let result = filterCompressedRanges(state, messages);
  result = pruneToolOutputs(state, result);
  result = pruneToolErrors(state, result);
  return result;
}