import {
  AIMessage,
  BaseMessage,
  HumanMessage,
  SystemMessage,
  ToolMessage,
} from '@langchain/core/messages';
import type { TokenCounter } from '@/types/run';
import { applyContextPruning } from '../contextPruning';

const counter: TokenCounter = (msg: BaseMessage) => {
  if (typeof msg.content === 'string') return Math.max(1, msg.content.length);
  return 1;
};

describe('applyContextPruning', () => {
  it('returns zero counts when disabled', () => {
    const messages: BaseMessage[] = [
      new HumanMessage('hi'),
      new AIMessage('thinking'),
      new ToolMessage({
        content: 'a'.repeat(10000),
        tool_call_id: 't1',
      }),
    ];
    const map = { 0: 1, 1: 1, 2: 10000 };
    const out = applyContextPruning({
      messages,
      indexTokenCountMap: map,
      tokenCounter: counter,
      config: { enabled: false },
    });
    expect(out.softTrimmed).toBe(0);
    expect(out.hardCleared).toBe(0);
    // Message untouched
    expect((messages[2] as ToolMessage).content).toHaveLength(10000);
  });

  it('soft-trims an older tool result with content over the threshold', () => {
    // Build a sequence with an old tool result outside the protected zone.
    const messages: BaseMessage[] = [
      new SystemMessage('sys'),
      new HumanMessage('q1'),
      new AIMessage({
        content: 'thinking',
        tool_calls: [{ name: 't', args: {}, id: 'old' }],
      }),
      new ToolMessage({
        // Long string — should soft-trim.
        content: 'OLD_TOOL_RESULT' + 'a'.repeat(20000),
        tool_call_id: 'old',
      }),
      new HumanMessage('q2'),
      new AIMessage('answer 1'),
      new HumanMessage('q3'),
      new AIMessage('answer 2'),
      new HumanMessage('q4'),
      new AIMessage('answer 3'),
      new HumanMessage('q5'),
      new AIMessage('answer 4'),
    ];
    const map: Record<string, number> = {};
    for (let i = 0; i < messages.length; i++) {
      map[i] = counter(messages[i]);
    }

    const result = applyContextPruning({
      messages,
      indexTokenCountMap: map,
      tokenCounter: counter,
      config: {
        enabled: true,
        keepLastAssistants: 2,
        softTrimRatio: 0.3,
        hardClearRatio: 0.95,
        minPrunableToolChars: 100,
        softTrim: { maxChars: 200, headChars: 50, tailChars: 50 },
        hardClear: { enabled: true, placeholder: '[evicted]' },
      },
    });

    // The old tool result at index 3 should be soft-trimmed.
    expect(result.softTrimmed).toBeGreaterThan(0);
    const oldToolMsg = messages[3] as ToolMessage;
    expect((oldToolMsg.content as string).length).toBeLessThan(20000);
    expect(oldToolMsg.content).toContain('soft-trimmed');
  });

  it('hard-clears very old tool results when above hardClearRatio', () => {
    // Position-age = (totalMessages - i) / totalMessages; with totalMessages=12
    // and i=0, ratio=1.0 (oldest). But system at i=0 is in the protected zone.
    // i=2 (the AI/tool pair) gives ratio=10/12 ≈ 0.83.
    const messages: BaseMessage[] = [
      new SystemMessage('sys'),
      new HumanMessage('q1'),
      new AIMessage({
        content: 'thinking',
        tool_calls: [{ name: 't', args: {}, id: 'oldest' }],
      }),
      new ToolMessage({
        content: 'OLDEST' + 'x'.repeat(5000),
        tool_call_id: 'oldest',
      }),
    ];
    for (let i = 0; i < 8; i++) {
      messages.push(new HumanMessage(`recent ${i}`));
      messages.push(new AIMessage(`reply ${i}`));
    }
    const map: Record<string, number> = {};
    for (let i = 0; i < messages.length; i++) {
      map[i] = counter(messages[i]);
    }

    const result = applyContextPruning({
      messages,
      indexTokenCountMap: map,
      tokenCounter: counter,
      config: {
        enabled: true,
        keepLastAssistants: 2,
        softTrimRatio: 0.3,
        hardClearRatio: 0.7,
        minPrunableToolChars: 100,
        softTrim: { maxChars: 200, headChars: 50, tailChars: 50 },
        hardClear: { enabled: true, placeholder: '[evicted]' },
      },
    });

    expect(result.hardCleared).toBeGreaterThan(0);
    const oldToolMsg = messages[3] as ToolMessage;
    expect(oldToolMsg.content).toBe('[evicted]');
  });

  it('does not touch tool results inside the protected zone (last N assistant turns)', () => {
    const messages: BaseMessage[] = [
      new SystemMessage('sys'),
      new HumanMessage('q1'),
      new AIMessage({
        content: 'thinking',
        tool_calls: [{ name: 't', args: {}, id: 'recent' }],
      }),
      new ToolMessage({
        content: 'RECENT' + 'y'.repeat(20000),
        tool_call_id: 'recent',
      }),
    ];
    const map: Record<string, number> = {};
    for (let i = 0; i < messages.length; i++) {
      map[i] = counter(messages[i]);
    }

    const result = applyContextPruning({
      messages,
      indexTokenCountMap: map,
      tokenCounter: counter,
      config: {
        enabled: true,
        keepLastAssistants: 4,
        softTrimRatio: 0.0,
        hardClearRatio: 0.0,
        minPrunableToolChars: 100,
        softTrim: { maxChars: 200, headChars: 50, tailChars: 50 },
        hardClear: { enabled: true, placeholder: '[evicted]' },
      },
    });

    expect(result.softTrimmed).toBe(0);
    expect(result.hardCleared).toBe(0);
    expect((messages[3] as ToolMessage).content).toContain('y'.repeat(100));
  });

  it('skips messages with image content even when out of protected zone', () => {
    const imageBlock = {
      type: 'image_url',
      image_url: { url: 'data:image/png;base64,iVBORw0K...' },
    };
    const messages: BaseMessage[] = [
      new SystemMessage('sys'),
      new HumanMessage('q'),
      new AIMessage({
        content: 'thinking',
        tool_calls: [{ name: 't', args: {}, id: 'img' }],
      }),
      new ToolMessage({
        content: [imageBlock] as never, // structured content not pruned
        tool_call_id: 'img',
      }),
    ];
    for (let i = 0; i < 8; i++) {
      messages.push(new HumanMessage(`recent ${i}`));
      messages.push(new AIMessage(`reply ${i}`));
    }
    const map: Record<string, number> = {};
    for (let i = 0; i < messages.length; i++) {
      map[i] = 1;
    }

    const result = applyContextPruning({
      messages,
      indexTokenCountMap: map,
      tokenCounter: counter,
      config: {
        enabled: true,
        keepLastAssistants: 2,
        softTrimRatio: 0.0,
        hardClearRatio: 0.0,
        minPrunableToolChars: 0,
        softTrim: { maxChars: 200, headChars: 50, tailChars: 50 },
        hardClear: { enabled: true, placeholder: '[evicted]' },
      },
    });

    // Image-containing tool message must not have been touched.
    expect(result.softTrimmed).toBe(0);
    expect(Array.isArray((messages[3] as ToolMessage).content)).toBe(true);
  });

  it('returns zero counts on empty message array', () => {
    const out = applyContextPruning({
      messages: [],
      indexTokenCountMap: {},
      tokenCounter: counter,
      config: { enabled: true },
    });
    expect(out.softTrimmed).toBe(0);
    expect(out.hardCleared).toBe(0);
  });
});