// Orphan tool message repair + sanitize coverage. Ported from upstream's
// `prune.test.ts` "Tool Message Handling" + top-level `sanitizeOrphanToolBlocks`
// describes — kept in a dedicated file to keep the diff against illuma's
// existing prune.test.ts narrow.
import {
  AIMessage,
  BaseMessage,
  ToolMessage,
  HumanMessage,
  SystemMessage,
} from '@langchain/core/messages';
import type * as t from '@/types';
import {
  repairOrphanedToolMessages,
  sanitizeOrphanToolBlocks,
} from '@/messages/prune';

const createTestTokenCounter = (): t.TokenCounter => {
  return (message: BaseMessage): number => {
    if (typeof message.content === 'string') {
      return Math.max(1, Math.ceil(message.content.length / 4));
    }
    if (Array.isArray(message.content)) {
      let total = 0;
      for (const block of message.content) {
        if (typeof block === 'object' && block !== null) {
          total += Math.ceil(JSON.stringify(block).length / 4);
        }
      }
      return Math.max(1, total);
    }
    return 1;
  };
};

describe('repairOrphanedToolMessages', () => {
  it('drops orphan ToolMessages that no longer have a matching AIMessage tool_call', () => {
    const tokenCounter = createTestTokenCounter();
    const context = [
      new SystemMessage('System'),
      new ToolMessage({ content: 'Orphan result', tool_call_id: 'gone' }),
      new HumanMessage('User question'),
      new AIMessage({
        content: 'thinking',
        tool_calls: [{ name: 't1', args: {}, id: 'tool1' }],
      }),
      new ToolMessage({ content: 'Tool result 1', tool_call_id: 'tool1' }),
    ];

    const result = repairOrphanedToolMessages({
      context,
      allMessages: context,
      tokenCounter,
      indexTokenCountMap: {},
    });

    expect(result.droppedOrphanCount).toBe(1);
    expect(result.context).toHaveLength(4);
    expect(
      result.context.some(
        (m) =>
          m.getType() === 'tool' && (m as ToolMessage).tool_call_id === 'gone'
      )
    ).toBe(false);
    expect(result.droppedMessages).toHaveLength(1);
  });

  it('strips orphan tool_use content blocks from AIMessages whose ToolMessage was dropped', () => {
    const tokenCounter = createTestTokenCounter();
    const context = [
      new HumanMessage('Hi'),
      new AIMessage({
        content: [
          { type: 'text', text: 'Calling tool…' },
          { type: 'tool_use', id: 'tool_orphan', name: 'foo', input: {} },
        ],
        tool_calls: [{ name: 'foo', args: {}, id: 'tool_orphan' }],
      }),
      // Note: NO matching ToolMessage for tool_orphan
    ];

    const result = repairOrphanedToolMessages({
      context,
      allMessages: context,
      tokenCounter,
      indexTokenCountMap: {},
    });

    // The AI message should still be in context, but with the tool_use block stripped
    expect(result.context).toHaveLength(2);
    const ai = result.context[1] as AIMessage;
    expect(ai.tool_calls?.length ?? 0).toBe(0);
    if (Array.isArray(ai.content)) {
      const toolUseBlocks = ai.content.filter(
        (b) =>
          typeof b === 'object' &&
          ((b as { type?: string }).type === 'tool_use' ||
            (b as { type?: string }).type === 'tool_call')
      );
      expect(toolUseBlocks).toHaveLength(0);
    }
  });

  it('drops the entire AIMessage when stripping leaves no content blocks', () => {
    const tokenCounter = createTestTokenCounter();
    const context = [
      new HumanMessage('Question'),
      new AIMessage({
        content: [
          { type: 'tool_use', id: 'orphan_only', name: 'foo', input: {} },
        ],
        tool_calls: [{ name: 'foo', args: {}, id: 'orphan_only' }],
      }),
    ];

    const result = repairOrphanedToolMessages({
      context,
      allMessages: context,
      tokenCounter,
      indexTokenCountMap: {},
    });

    expect(result.context).toHaveLength(1);
    expect(result.droppedOrphanCount).toBe(1);
    expect(result.droppedMessages).toHaveLength(1);
  });

  it('preserves matched tool_call ↔ tool_result pairs intact', () => {
    const tokenCounter = createTestTokenCounter();
    const context = [
      new HumanMessage('Question'),
      new AIMessage({
        content: 'thinking',
        tool_calls: [{ name: 'foo', args: {}, id: 'paired' }],
      }),
      new ToolMessage({ content: 'result', tool_call_id: 'paired' }),
    ];

    const result = repairOrphanedToolMessages({
      context,
      allMessages: context,
      tokenCounter,
      indexTokenCountMap: {},
    });

    expect(result.context).toEqual(context);
    expect(result.droppedOrphanCount).toBe(0);
  });
});

describe('sanitizeOrphanToolBlocks', () => {
  it('strips orphan tool_use blocks from AI messages with no matching ToolMessage', () => {
    const messages: BaseMessage[] = [
      new HumanMessage('Hello'),
      new AIMessage({
        content: [
          { type: 'text', text: 'Let me check.' },
          { type: 'tool_use', id: 'tu_orphan', name: 'search', input: {} },
        ],
        tool_calls: [{ name: 'search', args: {}, id: 'tu_orphan' }],
      }),
      // No ToolMessage with tool_call_id 'tu_orphan'
      new HumanMessage('Follow-up'),
    ];

    const out = sanitizeOrphanToolBlocks(messages);
    // The result depends on whether the stripped AI is at the end (popped by
    // the trailing-stripped-AI rule). Here it is NOT trailing, so it stays.
    expect(out).toHaveLength(3);
    const ai = out[1] as AIMessage;
    expect(ai.tool_calls?.length ?? 0).toBe(0);
    if (Array.isArray(ai.content)) {
      expect(
        ai.content.some(
          (b) =>
            typeof b === 'object' &&
            (b as { type?: string }).type === 'tool_use'
        )
      ).toBe(false);
    }
  });

  it('drops orphan ToolMessages whose matching tool_use is missing', () => {
    const messages: BaseMessage[] = [
      new HumanMessage('Hi'),
      new AIMessage('Plain response'),
      new ToolMessage({ content: 'orphaned', tool_call_id: 'gone' }),
    ];

    const out = sanitizeOrphanToolBlocks(messages);
    expect(out).toHaveLength(2);
    expect(out.find((m) => m.getType() === 'tool')).toBeUndefined();
  });

  it('returns the original array unchanged when every pair is matched (fast path)', () => {
    const messages: BaseMessage[] = [
      new HumanMessage('Hi'),
      new AIMessage({
        content: 'thinking',
        tool_calls: [{ name: 't1', args: {}, id: 'tool1' }],
      }),
      new ToolMessage({ content: 'r1', tool_call_id: 'tool1' }),
    ];
    const out = sanitizeOrphanToolBlocks(messages);
    expect(out).toBe(messages);
  });

  it('pops trailing AI messages whose tool_use blocks were all stripped', () => {
    const messages: BaseMessage[] = [
      new HumanMessage('Hi'),
      new AIMessage({
        content: [
          { type: 'tool_use', id: 'orphan_trail', name: 'x', input: {} },
        ],
        tool_calls: [{ name: 'x', args: {}, id: 'orphan_trail' }],
      }),
    ];

    const out = sanitizeOrphanToolBlocks(messages);
    // Trailing stripped-AI is popped per Bedrock/Anthropic rule: the
    // conversation must not end with an empty assistant exchange.
    expect(out).toHaveLength(1);
    expect(out[0].getType()).toBe('human');
  });

  it('ignores Anthropic server-tool ids (srvtoolu_*)', () => {
    const messages: BaseMessage[] = [
      new HumanMessage('Hi'),
      new AIMessage({
        content: [
          {
            type: 'tool_use',
            id: 'srvtoolu_anthropic_internal',
            name: 'web_search',
            input: {},
          },
        ],
        tool_calls: [
          { name: 'web_search', args: {}, id: 'srvtoolu_anthropic_internal' },
        ],
      }),
    ];

    const out = sanitizeOrphanToolBlocks(messages);
    // No orphans detected — the server-tool prefix is exempt.
    expect(out).toBe(messages);
  });
});