// src/specs/emergency-prune.test.ts
/**
 * Tests for the emergency pruning feature that handles "input too long" errors
 * by retrying with more aggressive message pruning and adding a context notice.
 */
import {
  HumanMessage,
  AIMessage,
  SystemMessage,
  BaseMessage,
} from '@langchain/core/messages';
import type * as t from '@/types';
import { createPruneMessages } from '@/messages/prune';
import { Providers } from '@/common';

// Simple token counter for testing (1 character = 1 token)
const createTestTokenCounter = (): t.TokenCounter => {
  return (message: BaseMessage): number => {
    const content = message.content as
      | string
      | Array<t.MessageContentComplex | string>
      | undefined;
    if (typeof content === 'string') {
      return content.length;
    }
    if (Array.isArray(content)) {
      return content.reduce((total, item) => {
        if (typeof item === 'string') return total + item.length;
        if (
          typeof item === 'object' &&
          'text' in item &&
          typeof item.text === 'string'
        ) {
          return total + item.text.length;
        }
        return total;
      }, 0);
    }
    return 0;
  };
};

// Helper to create test messages
const createTestMessages = (
  count: number,
  tokensPer: number
): BaseMessage[] => {
  const messages: BaseMessage[] = [
    new SystemMessage('You are a helpful assistant.'),
  ];

  for (let i = 0; i < count; i++) {
    const content = 'x'.repeat(tokensPer);
    if (i % 2 === 0) {
      messages.push(new HumanMessage(content));
    } else {
      messages.push(new AIMessage(content));
    }
  }

  return messages;
};

// Helper to build indexTokenCountMap
const buildIndexTokenCountMap = (
  messages: BaseMessage[],
  tokenCounter: t.TokenCounter
): Record<string, number> => {
  const map: Record<string, number> = {};
  messages.forEach((msg, index) => {
    map[index] = tokenCounter(msg);
  });
  return map;
};

/**
 * Estimates a human-friendly description of the conversation timeframe based on message count.
 * This mirrors the implementation in Graph.ts
 */
const getContextTimeframeDescription = (messageCount: number): string => {
  if (messageCount <= 5) {
    return 'just the last few exchanges';
  } else if (messageCount <= 15) {
    return 'the last several minutes';
  } else if (messageCount <= 30) {
    return 'roughly the past hour';
  } else if (messageCount <= 60) {
    return 'the past couple of hours';
  } else if (messageCount <= 150) {
    return 'the past few hours';
  } else if (messageCount <= 300) {
    return 'roughly a day\'s worth';
  } else if (messageCount <= 700) {
    return 'the past few days';
  } else {
    return 'about a week or more';
  }
};

describe('Emergency Pruning Feature', () => {
  const tokenCounter = createTestTokenCounter();

  describe('Normal Pruning vs Emergency Pruning', () => {
    it('should prune more aggressively with 50% reduced context', () => {
      // Create 20 messages, each with 100 tokens = 2000 tokens total (excluding system)
      const messages = createTestMessages(20, 100);
      const indexTokenCountMap = buildIndexTokenCountMap(
        messages,
        tokenCounter
      );

      // Normal prune with 1500 token limit
      const normalMaxTokens = 1500;
      const normalPrune = createPruneMessages({
        startIndex: 0,
        provider: Providers.BEDROCK,
        tokenCounter,
        maxTokens: normalMaxTokens,
        thinkingEnabled: false,
        indexTokenCountMap,
      });

      const { context: normalContext } = normalPrune({ messages });

      // Emergency prune with 50% (750 tokens)
      const emergencyMaxTokens = Math.floor(normalMaxTokens * 0.5);
      const emergencyPrune = createPruneMessages({
        startIndex: 0,
        provider: Providers.BEDROCK,
        tokenCounter,
        maxTokens: emergencyMaxTokens,
        thinkingEnabled: false,
        indexTokenCountMap,
      });

      const { context: emergencyContext } = emergencyPrune({ messages });

      // Emergency should have fewer messages
      expect(emergencyContext.length).toBeLessThan(normalContext.length);
      console.log(
        `Normal prune: ${normalContext.length} messages, Emergency prune: ${emergencyContext.length} messages`
      );
    });

    it('should preserve system message and latest user message after emergency prune', () => {
      const messages = createTestMessages(10, 200);
      const indexTokenCountMap = buildIndexTokenCountMap(
        messages,
        tokenCounter
      );

      // Very aggressive prune - only 300 tokens
      const emergencyPrune = createPruneMessages({
        startIndex: 0,
        provider: Providers.BEDROCK,
        tokenCounter,
        maxTokens: 300,
        thinkingEnabled: false,
        indexTokenCountMap,
      });

      const { context } = emergencyPrune({ messages });

      // Should still have system message if it fits
      if (context.length > 0) {
        // Check that we have at least the most recent messages
        const lastMessage = context[context.length - 1];
        expect(lastMessage).toBeDefined();
      }
    });
  });

  describe('Pruning Notice Message Injection', () => {
    it('should calculate correct number of pruned messages', () => {
      const originalCount = 20;
      const messages = createTestMessages(originalCount, 100);
      const indexTokenCountMap = buildIndexTokenCountMap(
        messages,
        tokenCounter
      );

      const emergencyPrune = createPruneMessages({
        startIndex: 0,
        provider: Providers.BEDROCK,
        tokenCounter,
        maxTokens: 500, // Very small to force aggressive pruning
        thinkingEnabled: false,
        indexTokenCountMap,
      });

      const { context: reducedMessages } = emergencyPrune({ messages });

      // Calculate how many were pruned (this is what we inject in the notice)
      const prunedCount = messages.length - reducedMessages.length;

      expect(prunedCount).toBeGreaterThan(0);
      console.log(
        `Original: ${messages.length}, After prune: ${reducedMessages.length}, Pruned: ${prunedCount}`
      );
    });

    it('should inject personalized notice message after system message', () => {
      const messages = createTestMessages(10, 100);
      const indexTokenCountMap = buildIndexTokenCountMap(
        messages,
        tokenCounter
      );

      const emergencyPrune = createPruneMessages({
        startIndex: 0,
        provider: Providers.BEDROCK,
        tokenCounter,
        maxTokens: 800,
        thinkingEnabled: false,
        indexTokenCountMap,
      });

      const { context: reducedMessages } = emergencyPrune({ messages });

      // Simulate the notice injection logic from Graph.ts
      const prunedCount = messages.length - reducedMessages.length;
      const remainingCount = reducedMessages.length;
      const estimatedContextDescription =
        getContextTimeframeDescription(remainingCount);

      const pruneNoticeMessage = new HumanMessage({
        content: `[CONTEXT NOTICE]
Our conversation has grown quite long, so I've focused on ${estimatedContextDescription} of our chat (${remainingCount} recent messages). ${prunedCount} earlier messages are no longer in my immediate memory.

If I seem to be missing something we discussed earlier, just give me a quick reminder and I'll pick right back up! I'm still fully engaged and ready to help with whatever you need.`,
      });

      // Insert after system message
      const hasSystemMessage = reducedMessages[0]?.getType() === 'system';
      const insertIndex = hasSystemMessage ? 1 : 0;

      const messagesWithNotice = [
        ...reducedMessages.slice(0, insertIndex),
        pruneNoticeMessage,
        ...reducedMessages.slice(insertIndex),
      ];

      // Verify notice is in correct position
      if (hasSystemMessage) {
        expect(messagesWithNotice[0].getType()).toBe('system');
        expect(messagesWithNotice[1].getType()).toBe('human');
        expect(messagesWithNotice[1].content as string).toContain(
          '[CONTEXT NOTICE]'
        );
        expect(messagesWithNotice[1].content as string).toContain(
          'recent messages'
        );
        expect(messagesWithNotice[1].content as string).toContain(
          'quick reminder'
        );
      } else {
        expect(messagesWithNotice[0].getType()).toBe('human');
        expect(messagesWithNotice[0].content as string).toContain(
          '[CONTEXT NOTICE]'
        );
      }

      // Total messages should be reduced + 1 notice
      expect(messagesWithNotice.length).toBe(reducedMessages.length + 1);

      console.log(
        `Notice preview:\n${(pruneNoticeMessage.content as string).substring(0, 200)}...`
      );
    });
  });

  describe('Context Timeframe Description', () => {
    it('should return appropriate descriptions for different message counts', () => {
      expect(getContextTimeframeDescription(3)).toBe(
        'just the last few exchanges'
      );
      expect(getContextTimeframeDescription(10)).toBe(
        'the last several minutes'
      );
      expect(getContextTimeframeDescription(25)).toBe('roughly the past hour');
      expect(getContextTimeframeDescription(45)).toBe(
        'the past couple of hours'
      );
      expect(getContextTimeframeDescription(100)).toBe('the past few hours');
      expect(getContextTimeframeDescription(200)).toBe('roughly a day\'s worth');
      expect(getContextTimeframeDescription(500)).toBe('the past few days');
      expect(getContextTimeframeDescription(1000)).toBe('about a week or more');
    });
  });

  describe('Error Detection Patterns', () => {
    const errorPatterns = [
      'Input is too long for the model',
      'context length exceeded',
      'maximum context length',
      'ValidationException: Input is too long',
      'prompt is too long for this model',
      'The input is too long',
    ];

    it('should detect various "input too long" error patterns', () => {
      const isInputTooLongError = (errorMessage: string): boolean => {
        const lowerMessage = errorMessage.toLowerCase();
        return (
          lowerMessage.includes('too long') ||
          lowerMessage.includes('input is too long') ||
          lowerMessage.includes('context length') ||
          lowerMessage.includes('maximum context') ||
          lowerMessage.includes('validationexception') ||
          lowerMessage.includes('prompt is too long')
        );
      };

      for (const pattern of errorPatterns) {
        expect(isInputTooLongError(pattern)).toBe(true);
        console.log(`✓ Detected: "${pattern}"`);
      }

      // Should not match unrelated errors
      expect(isInputTooLongError('Network timeout')).toBe(false);
      expect(isInputTooLongError('Invalid API key')).toBe(false);
      expect(isInputTooLongError('Rate limit exceeded')).toBe(false);
    });
  });

  describe('Edge Cases', () => {
    it('should handle empty messages after pruning', () => {
      // Single very long message that exceeds the limit
      const messages: BaseMessage[] = [
        new SystemMessage('System prompt'),
        new HumanMessage('x'.repeat(10000)), // Way too long
      ];
      const indexTokenCountMap = buildIndexTokenCountMap(
        messages,
        tokenCounter
      );

      const emergencyPrune = createPruneMessages({
        startIndex: 0,
        provider: Providers.BEDROCK,
        tokenCounter,
        maxTokens: 100, // Very small limit
        thinkingEnabled: false,
        indexTokenCountMap,
      });

      const { context } = emergencyPrune({ messages });

      // Should have at least tried to keep something or be empty
      // The key is it shouldn't throw
      expect(Array.isArray(context)).toBe(true);
    });

    it('should work with only system message and one user message', () => {
      const messages: BaseMessage[] = [
        new SystemMessage('You are helpful.'),
        new HumanMessage('Hello'),
      ];
      const indexTokenCountMap = buildIndexTokenCountMap(
        messages,
        tokenCounter
      );

      const emergencyPrune = createPruneMessages({
        startIndex: 0,
        provider: Providers.BEDROCK,
        tokenCounter,
        maxTokens: 500,
        thinkingEnabled: false,
        indexTokenCountMap,
      });

      const { context } = emergencyPrune({ messages });

      expect(context.length).toBe(2);
      expect(context[0].getType()).toBe('system');
      expect(context[1].getType()).toBe('human');
    });

    it('should handle conversation without system message', () => {
      const messages: BaseMessage[] = [
        new HumanMessage('Hello'),
        new AIMessage('Hi there!'),
        new HumanMessage('How are you?'),
      ];
      const indexTokenCountMap = buildIndexTokenCountMap(
        messages,
        tokenCounter
      );

      const emergencyPrune = createPruneMessages({
        startIndex: 0,
        provider: Providers.BEDROCK,
        tokenCounter,
        maxTokens: 100,
        thinkingEnabled: false,
        indexTokenCountMap,
      });

      const { context } = emergencyPrune({ messages });

      // Should keep the most recent messages that fit
      expect(context.length).toBeGreaterThan(0);
      expect(context[0].getType()).not.toBe('system');
    });
  });
});