import {
  HumanMessage,
  AIMessage,
  SystemMessage,
  ToolMessage,
  BaseMessage,
} from '@langchain/core/messages';
import type { TokenCounter } from '@/types/run';
import {
  formatMessagesForSummary,
  buildFullSummaryPrompt,
  buildSimpleSummaryPrompt,
  createEmergencySummary,
  validateSummarySize,
  summarize,
} from './summarize';

/**
 * Simple token counter that approximates 1 token per 4 characters.
 * Used across all tests that need a TokenCounter.
 */
const simpleTokenCounter: TokenCounter = (msg: BaseMessage): number => {
  const content =
    typeof msg.content === 'string' ? msg.content : JSON.stringify(msg.content);
  return Math.ceil(content.length / 4);
};

describe('formatMessagesForSummary', () => {
  it('formats messages as "role: content" pairs', () => {
    const messages = [
      new HumanMessage('Hello there'),
      new AIMessage('Hi! How can I help?'),
      new SystemMessage('You are a helpful assistant'),
    ];

    const result = formatMessagesForSummary(messages);

    expect(result).toBe(
      'human: Hello there\nai: Hi! How can I help?\nsystem: You are a helpful assistant'
    );
  });

  it('truncates long content (>4000 chars) with ellipsis', () => {
    const longContent = 'x'.repeat(5000);
    const messages = [new HumanMessage(longContent)];

    const result = formatMessagesForSummary(messages);

    // 4000 chars of content + "..." + "human: " prefix
    expect(result).toBe(`human: ${'x'.repeat(4000)}...`);
    expect(result.length).toBe(4000 + 3 + 7); // content + ellipsis + "human: "
  });

  it('does not truncate content at exactly 4000 chars', () => {
    const exactContent = 'y'.repeat(4000);
    const messages = [new HumanMessage(exactContent)];

    const result = formatMessagesForSummary(messages);

    expect(result).toBe(`human: ${exactContent}`);
    expect(result).not.toContain('...');
  });

  it('handles empty messages array', () => {
    const result = formatMessagesForSummary([]);
    expect(result).toBe('');
  });

  it('handles messages with complex (array) content', () => {
    const messages = [
      new HumanMessage({
        content: [
          { type: 'text', text: 'Look at this image' },
          {
            type: 'image_url',
            image_url: { url: 'data:image/png;base64,abc123' },
          },
        ],
      }),
    ];

    const result = formatMessagesForSummary(messages);

    // Array content gets JSON.stringified
    expect(result).toContain('human: ');
    expect(result).toContain('Look at this image');
    expect(result).toContain('image_url');
  });
});

describe('buildFullSummaryPrompt', () => {
  const conversation = 'human: Hello\nai: Hi there!';

  it('contains all 10 section headers', () => {
    const prompt = buildFullSummaryPrompt(conversation);

    const expectedSections = [
      '## 1. Task Overview & User Intent',
      '## 2. Technical Context',
      '## 3. Files & Code State',
      '## 4. Problem Resolution History',
      '## 5. Agent Workflow State',
      '## 6. Tool Results Summary',
      '## 7. Progress Tracking',
      '## 8. Active Working State',
      '## 9. Continuation Plan',
      '## 10. Critical Context',
    ];

    for (const section of expectedSections) {
      expect(prompt).toContain(section);
    }
  });

  it('includes conversation text', () => {
    const prompt = buildFullSummaryPrompt(conversation);
    expect(prompt).toContain(conversation);
  });

  it('includes agent workflow state when isMultiAgent=true', () => {
    const prompt = buildFullSummaryPrompt(conversation, {
      isMultiAgent: true,
      agentWorkflowState: {
        currentAgentId: 'coder',
        agentChain: ['planner', 'coder'],
        pendingAgents: ['reviewer'],
      },
    });

    expect(prompt).toContain('Current agent: coder');
    expect(prompt).toContain('Agent chain: planner -> coder');
    expect(prompt).toContain('Pending agents: reviewer');
    expect(prompt).not.toContain('N/A (single-agent conversation)');
  });

  it('shows "none" for empty pendingAgents in multi-agent mode', () => {
    const prompt = buildFullSummaryPrompt(conversation, {
      isMultiAgent: true,
      agentWorkflowState: {
        currentAgentId: 'coder',
        agentChain: ['coder'],
        pendingAgents: [],
      },
    });

    expect(prompt).toContain('Pending agents: none');
  });

  it('omits agent workflow when not multi-agent', () => {
    const prompt = buildFullSummaryPrompt(conversation);
    expect(prompt).toContain('N/A (single-agent conversation)');
  });

  it('omits agent workflow when isMultiAgent is true but no state provided', () => {
    const prompt = buildFullSummaryPrompt(conversation, { isMultiAgent: true });
    expect(prompt).toContain('N/A (single-agent conversation)');
  });
});

describe('buildSimpleSummaryPrompt', () => {
  it('contains "Summarize this conversation" directive', () => {
    const prompt = buildSimpleSummaryPrompt('human: test');
    expect(prompt).toContain('Summarize this conversation');
  });

  it('includes conversation text', () => {
    const conversation = 'human: What is 2+2?\nai: 4';
    const prompt = buildSimpleSummaryPrompt(conversation);
    expect(prompt).toContain(conversation);
  });
});

describe('createEmergencySummary', () => {
  it('extracts first user message', () => {
    const messages = [
      new HumanMessage('First question'),
      new HumanMessage('Second question'),
    ];

    const result = createEmergencySummary(messages);

    expect(result).toContain('Original request: First question');
  });

  it('extracts last AI message', () => {
    const messages = [
      new AIMessage('First answer'),
      new AIMessage('Second answer'),
    ];

    const result = createEmergencySummary(messages);

    // Should contain the LAST AI message, not the first
    expect(result).toContain('Last response: Second answer');
  });

  it('extracts tool names from tool messages', () => {
    const messages = [
      new ToolMessage({
        content: 'result1',
        tool_call_id: 'tc1',
        name: 'search',
      }),
      new ToolMessage({
        content: 'result2',
        tool_call_id: 'tc2',
        name: 'calculator',
      }),
      new ToolMessage({
        content: 'result3',
        tool_call_id: 'tc3',
        name: 'search',
      }), // duplicate
    ];

    const result = createEmergencySummary(messages);

    expect(result).toContain('Tools used: search, calculator');
  });

  it('reports correct message count', () => {
    const messages = [
      new HumanMessage('q1'),
      new AIMessage('a1'),
      new HumanMessage('q2'),
      new AIMessage('a2'),
      new ToolMessage({ content: 'r', tool_call_id: 'tc1', name: 'tool1' }),
    ];

    const result = createEmergencySummary(messages);

    expect(result).toContain('Messages compacted: 5');
  });

  it('handles empty array', () => {
    const result = createEmergencySummary([]);

    expect(result).toContain('[Emergency Context Summary]');
    expect(result).toContain('Original request: ');
    expect(result).toContain('Last response: ');
    expect(result).toContain('Tools used: none');
    expect(result).toContain('Messages compacted: 0');
  });

  it('handles messages with no user or AI messages', () => {
    const messages = [
      new SystemMessage('You are helpful'),
      new ToolMessage({ content: 'ok', tool_call_id: 'tc1', name: 'ping' }),
    ];

    const result = createEmergencySummary(messages);

    expect(result).toContain('Original request: ');
    expect(result).toContain('Last response: ');
    expect(result).toContain('Tools used: ping');
    expect(result).toContain('Messages compacted: 2');
  });

  it('truncates long content to 500 chars', () => {
    const longContent = 'a'.repeat(1000);
    const messages = [new HumanMessage(longContent)];

    const result = createEmergencySummary(messages);

    // The original request should be truncated to 500 chars
    expect(result).toContain(`Original request: ${'a'.repeat(500)}`);
    expect(result).not.toContain('a'.repeat(501));
  });
});

describe('validateSummarySize', () => {
  it('returns true when summary fits in budget', () => {
    const summary = 'Short summary';
    // simpleTokenCounter: ceil(13/4) = 4 tokens
    const result = validateSummarySize(summary, 100, simpleTokenCounter);
    expect(result).toBe(true);
  });

  it('returns false when summary exceeds budget', () => {
    const summary = 'This is a summary that should exceed a tiny budget';
    // simpleTokenCounter: ceil(50/4) = 13 tokens
    const result = validateSummarySize(summary, 1, simpleTokenCounter);
    expect(result).toBe(false);
  });

  it('uses the provided tokenCounter', () => {
    const customCounter = jest.fn().mockReturnValue(50);
    validateSummarySize('any text', 100, customCounter);

    expect(customCounter).toHaveBeenCalledTimes(1);
    // Should be called with a HumanMessage wrapping the summary text
    const arg = customCounter.mock.calls[0][0];
    expect(arg).toBeInstanceOf(HumanMessage);
    expect(arg.content).toBe('any text');
  });
});

describe('summarize', () => {
  const sampleMessages = [
    new HumanMessage('What is TypeScript?'),
    new AIMessage('TypeScript is a typed superset of JavaScript.'),
    new ToolMessage({
      content: 'tool result',
      tool_call_id: 'tc1',
      name: 'search',
    }),
  ];

  it('Tier 1 success: callback returns valid summary -> result.tier === "full"', async () => {
    const mockCallback = jest
      .fn()
      .mockResolvedValueOnce('## Summary\nThis is a valid summary');

    const result = await summarize(sampleMessages, mockCallback);

    expect(result.tier).toBe('full');
    expect(result.summary).toBe('## Summary\nThis is a valid summary');
    expect(mockCallback).toHaveBeenCalledTimes(1);
  });

  it('Tier 1 fails -> Tier 2 success: first callback throws, second succeeds', async () => {
    const mockCallback = jest
      .fn()
      .mockRejectedValueOnce(new Error('LLM overloaded'))
      .mockResolvedValueOnce('Simple summary');

    const result = await summarize(sampleMessages, mockCallback);

    expect(result.tier).toBe('simple');
    expect(result.summary).toBe('Simple summary');
    expect(mockCallback).toHaveBeenCalledTimes(2);
  });

  it('Tier 1 + 2 fail -> Tier 3: callback always throws', async () => {
    const mockCallback = jest.fn().mockRejectedValue(new Error('LLM down'));

    const result = await summarize(sampleMessages, mockCallback);

    expect(result.tier).toBe('emergency');
    expect(result.summary).toContain('[Emergency Context Summary]');
    expect(result.summary).toContain('What is TypeScript?');
    expect(mockCallback).toHaveBeenCalledTimes(2);
  });

  it('Callback returns undefined -> falls to Tier 2', async () => {
    const mockCallback = jest
      .fn()
      .mockResolvedValueOnce(undefined)
      .mockResolvedValueOnce('Tier 2 summary');

    const result = await summarize(sampleMessages, mockCallback);

    expect(result.tier).toBe('simple');
    expect(result.summary).toBe('Tier 2 summary');
    expect(mockCallback).toHaveBeenCalledTimes(2);
  });

  it('Callback returns empty string -> falls to Tier 2', async () => {
    const mockCallback = jest
      .fn()
      .mockResolvedValueOnce('')
      .mockResolvedValueOnce('Tier 2 result');

    const result = await summarize(sampleMessages, mockCallback);

    expect(result.tier).toBe('simple');
    expect(result.summary).toBe('Tier 2 result');
  });

  it('Empty messages: returns emergency summary immediately', async () => {
    const mockCallback = jest.fn().mockResolvedValue('should not matter');

    const result = await summarize([], mockCallback);

    // With empty messages, tier 1 still runs (formatMessagesForSummary returns ''),
    // but callback returns content so it should succeed as tier 1
    // unless the callback result is validated against budget.
    // Actually, empty messages still go through the normal flow.
    // The callback returns a valid string, so it should be 'full'.
    expect(result.messagesCompacted).toBe(0);
  });

  it('Token budget validation: summary exceeds budget -> falls to next tier', async () => {
    const longSummary = 'x'.repeat(2000);
    const mockCallback = jest
      .fn()
      .mockResolvedValueOnce(longSummary) // Tier 1: exceeds budget
      .mockResolvedValueOnce('Short tier 2 summary'); // Tier 2: no budget check

    const result = await summarize(sampleMessages, mockCallback, {
      summaryBudget: 10, // Very small budget: 10 tokens
      tokenCounter: simpleTokenCounter,
    });

    expect(result.tier).toBe('simple');
    expect(result.summary).toBe('Short tier 2 summary');
    expect(mockCallback).toHaveBeenCalledTimes(2);
  });

  it('Token budget validation: summary fits within budget -> stays at Tier 1', async () => {
    const mockCallback = jest.fn().mockResolvedValueOnce('Short summary');

    const result = await summarize(sampleMessages, mockCallback, {
      summaryBudget: 1000,
      tokenCounter: simpleTokenCounter,
    });

    expect(result.tier).toBe('full');
    expect(result.summary).toBe('Short summary');
    expect(mockCallback).toHaveBeenCalledTimes(1);
  });

  it('messagesCompacted count matches input length', async () => {
    const mockCallback = jest.fn().mockResolvedValueOnce('Valid summary');

    const messages = [
      new HumanMessage('q1'),
      new AIMessage('a1'),
      new HumanMessage('q2'),
      new AIMessage('a2'),
      new HumanMessage('q3'),
    ];

    const result = await summarize(messages, mockCallback);

    expect(result.messagesCompacted).toBe(5);
  });

  it('messagesCompacted count is correct even for emergency tier', async () => {
    const mockCallback = jest.fn().mockRejectedValue(new Error('fail'));

    const result = await summarize(sampleMessages, mockCallback);

    expect(result.tier).toBe('emergency');
    expect(result.messagesCompacted).toBe(3);
  });

  it('passes maxOutputTokens to Tier 1 callback', async () => {
    const mockCallback = jest.fn().mockResolvedValueOnce('Summary');

    await summarize(sampleMessages, mockCallback, { maxOutputTokens: 2048 });

    expect(mockCallback).toHaveBeenCalledWith(expect.any(String), 2048);
  });

  it('defaults maxOutputTokens to 1024 for Tier 1', async () => {
    const mockCallback = jest.fn().mockResolvedValueOnce('Summary');

    await summarize(sampleMessages, mockCallback);

    expect(mockCallback).toHaveBeenCalledWith(expect.any(String), 1024);
  });

  it('passes 512 as maxTokens for Tier 2 (simple) callback', async () => {
    const mockCallback = jest
      .fn()
      .mockRejectedValueOnce(new Error('Tier 1 fail'))
      .mockResolvedValueOnce('Simple summary');

    await summarize(sampleMessages, mockCallback);

    // Second call (Tier 2) should use SIMPLE_MAX_TOKENS = 512
    expect(mockCallback).toHaveBeenCalledTimes(2);
    expect(mockCallback.mock.calls[1][1]).toBe(512);
  });
});