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

/** Return the first message index in the newest protected user turns. */
export function getProtectedTurnStart(messages: AgentMessage[], turns: number): number | undefined {
  if (turns <= 0) return undefined;
  const userIndices = messages.flatMap((message, index) =>
    message.role === "user" ? [index] : [],
  );
  return userIndices[Math.max(0, userIndices.length - turns)];
}

/**
 * Resolve a boundary ID (m0001 or b1) to a message array index.
 */
export function resolveBoundaryIndex(state: SessionState, boundaryId: string): number | undefined {
  const parsed = parseBoundaryId(boundaryId);
  if (!parsed) return undefined;

  if (parsed.type === "message") {
    const ref = boundaryId;
    // Reverse-lookup: find the index that maps to this ref in the runtime cache.
    // byRef now maps ref->rawId (not ref->index), so we scan byIndex instead.
    for (const [idx, r] of state.messageIds.byIndex) {
      if (r === ref) return idx;
    }
    return undefined;
  }

  if (parsed.type === "block") {
    // Find the anchor index for this block
    for (const [anchorIndex, blockId] of state.prune.messages.activeByAnchorIndex) {
      if (blockId === parsed.blockId) return anchorIndex;
    }
    return undefined;
  }

  return undefined;
}

export interface SelectionResult {
  messageIndices: number[];
  startIndex: number;
  endIndex: number;
  toolIds: string[];
  directMessageIndices: number[];
  directToolIds: string[];
  consumedBlockIds: number[];
}

/**
 * Expand a compression range to ensure all tool call chains are complete.
 * If state is provided with populated toolParameters, uses cached assistantIndex/resultIndex
 * for O(C) lookups (C = cached tool calls) instead of scanning the full message array O(N).
 * Falls back to scan when state is not provided or toolParameters is empty.
 */
export function expandRangeForToolChains(
  messages: AgentMessage[],
  startIndex: number,
  endIndex: number,
  state?: SessionState,
) {
  // Fast path: use cached indices from tool parameter entries
  if (state && state.toolParameters.size > 0) {
    return expandWithCachedIndices(messages, startIndex, endIndex, state);
  }

  // Fallback: scan-based expansion (original implementation)
  return expandByScan(messages, startIndex, endIndex);
}

function expandWithCachedIndices(
  messages: AgentMessage[],
  startIndex: number,
  endIndex: number,
  state: SessionState,
) {
  let start = startIndex;
  let end = endIndex;
  let changed = true;
  const maxIndex = messages.length - 1;

  while (changed) {
    changed = false;

    for (const [, entry] of state.toolParameters) {
      const aIdx = entry.assistantIndex;
      const rIdx = entry.resultIndex;
      if (aIdx === undefined) continue;

      // Skip entries with stale indices beyond current messages array
      if (aIdx > maxIndex || (rIdx !== undefined && rIdx > maxIndex)) continue;

      // Assistant in range but result outside → expand end
      if (rIdx !== undefined && aIdx >= start && aIdx <= end && rIdx > end) {
        end = rIdx;
        changed = true;
      }

      // Result in range but assistant outside → expand start
      if (rIdx !== undefined && rIdx >= start && rIdx <= end && aIdx < start) {
        start = aIdx;
        changed = true;
      }
    }
  }

  return { startIndex: start, endIndex: end };
}

function expandByScan(messages: AgentMessage[], startIndex: number, endIndex: number) {
  let start = startIndex;
  let end = endIndex;
  let changed = true;

  while (changed) {
    changed = false;

    // Collect all toolCall IDs from assistant messages in [start, end]
    const callIdsInRange = new Set<string>();
    for (let i = start; i <= end; i++) {
      const msg = messages[i];
      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") {
          callIdsInRange.add(p.id as string);
        }
      }
    }

    // For each toolCall in range, ensure its toolResult is also in range
    for (let i = end + 1; i < messages.length; i++) {
      const msg = messages[i];
      if (msg.role !== "toolResult") continue;
      if (callIdsInRange.has(msg.toolCallId)) {
        end = i;
        changed = true;
      }
    }

    // Collect all toolResult toolCallIds in [start, end]
    const resultCallIdsInRange = new Set<string>();
    for (let i = start; i <= end; i++) {
      const msg = messages[i];
      if (msg.role !== "toolResult") continue;
      resultCallIdsInRange.add(msg.toolCallId);
    }

    // For each toolResult in range, ensure its assistant toolCall is also in range
    for (let i = 0; i < start; i++) {
      const msg = messages[i];
      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") {
          if (resultCallIdsInRange.has(p.id as string)) {
            start = i;
            changed = true;
          }
        }
      }
    }
  }

  return { startIndex: start, endIndex: end };
}

/**
 * Collect message indices in a range [startIndex, endIndex].
 * Auto-expands the range to protect tool call chains from being split.
 * If state is provided, uses cached indices for faster expansion.
 */
export function resolveSelection(
  messages: AgentMessage[],
  startIndex: number,
  endIndex: number,
  state?: SessionState,
): SelectionResult {
  if (startIndex > endIndex) {
    throw new Error(`startId appears after endId in the conversation. Start must come before end.`);
  }

  if (startIndex < 0 || endIndex >= messages.length) {
    throw new Error(`Boundary indices out of range. Valid range: 0-${messages.length - 1}`);
  }

  const direct = expandRangeForToolChains(messages, startIndex, endIndex, state);
  const directMembership = collectMembership(messages, direct.startIndex, direct.endIndex);
  let start = direct.startIndex;
  let end = direct.endIndex;
  const consumedBlockIds = new Set<number>();
  let changed = true;

  while (changed) {
    changed = false;
    const expanded = expandRangeForToolChains(messages, start, end, state);
    if (expanded.startIndex !== start || expanded.endIndex !== end) {
      start = expanded.startIndex;
      end = expanded.endIndex;
      changed = true;
    }

    if (!state) continue;
    for (const blockId of state.prune.messages.activeBlockIds) {
      const block = state.prune.messages.blocksById.get(blockId);
      if (
        !block?.active ||
        !block.effectiveMessageIndices.some((index) => index >= start && index <= end)
      ) {
        continue;
      }

      consumedBlockIds.add(blockId);
      const blockStart = Math.min(...block.effectiveMessageIndices);
      const blockEnd = Math.max(...block.effectiveMessageIndices);
      if (blockStart < start || blockEnd > end) {
        start = Math.min(start, blockStart);
        end = Math.max(end, blockEnd);
        changed = true;
      }
    }
  }

  const membership = collectMembership(messages, start, end);

  return {
    ...membership,
    startIndex: start,
    endIndex: end,
    directMessageIndices: directMembership.messageIndices,
    directToolIds: directMembership.toolIds,
    consumedBlockIds: [...consumedBlockIds].sort((a, b) => a - b),
  };
}

function collectMembership(
  messages: AgentMessage[],
  start: number,
  end: number,
): { messageIndices: number[]; toolIds: string[] } {
  const messageIndices: number[] = [];
  const toolIds = new Set<string>();
  for (let i = start; i <= end; i++) {
    messageIndices.push(i);
    const message = messages[i];
    if (message.role !== "assistant" || !Array.isArray(message.content)) continue;
    for (const part of message.content) {
      if (typeof part !== "object" || part === null) continue;
      const call = part as unknown as Record<string, unknown>;
      if (call.type === "toolCall" && typeof call.id === "string") toolIds.add(call.id);
    }
  }

  return { messageIndices, toolIds: [...toolIds].sort() };
}