/**
 * Tests for long-task/messages-compressor.ts (Phase 4.6 of agent-fix sprint)
 * + integration tests for runRetrievalAgentLoop's messagesContextManager hook.
 *
 * Coverage:
 *   - Unit: maybeCompressMessages
 *     - No-op when below threshold
 *     - Triggers compression when above threshold
 *     - Preserves system + protected head verbatim
 *     - Preserves protected tail verbatim
 *     - Replaces middle with single summary message
 *     - LLM call invoked when configured
 *     - Falls back to placeholder when no LLM
 *     - shouldCompressMessages predicate matches behavior
 *     - Custom estimateTokensFn honored
 *     - onCompressionEvent fires with correct payload
 *   - Integration: runRetrievalAgentLoop with messagesContextManager
 *     - Backwards compat: omitted field → behavior unchanged
 *     - Compression fires at threshold during multi-turn loop
 *     - Emits 'messages_compressed' event
 *     - Compression doesn't break retrieve / finalize / loop_exhausted paths
 *     - Multiple compressions in one loop run
 *     - Final answer reachable after compression
 */

import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import * as fsp from 'node:fs/promises';
import * as path from 'node:path';
import * as os from 'node:os';

import {
  maybeCompressMessages,
  shouldCompressMessages,
  type MessagesContextManagerConfig,
  type MessagesCompressionEvent,
} from '../src/long-task/messages-compressor.js';
import {
  runRetrievalAgentLoop,
  type LlmCallFn,
  type LlmCallResult,
  type RetrievalSearchFn,
  type MultiStepAgentRunConfig,
  type AgentRunProgressEvent,
} from '../src/retrieval-agent-loop.js';

// ─────────────────────────────────────────────────────────────────────────
// Fixtures
// ─────────────────────────────────────────────────────────────────────────

let tmpRoot: string;

beforeEach(async () => {
  tmpRoot = await fsp.mkdtemp(path.join(os.tmpdir(), 'waggle-msgcomp-'));
});

afterEach(async () => {
  await fsp.rm(tmpRoot, { recursive: true, force: true });
});

function makeMessages(count: number, fillerChars = 100): Array<{ role: string; content: string }> {
  const out: Array<{ role: string; content: string }> = [
    { role: 'system', content: 'You are a research assistant.' },
    { role: 'user', content: 'Survey 30 historical events and rank themes.' },
  ];
  for (let i = 0; i < count; i += 1) {
    out.push({ role: 'assistant', content: `{"action":"retrieve","query":"q${i}"}` });
    out.push({ role: 'user', content: `[results for q${i}]\n` + 'X'.repeat(fillerChars) });
  }
  return out;
}

function makeFakeLlm(content: string, costUsd = 0.001): LlmCallFn {
  return async () => ({
    content,
    inTokens: 100,
    outTokens: 50,
    costUsd,
    latencyMs: 50,
  });
}

// ─────────────────────────────────────────────────────────────────────────
// Unit: maybeCompressMessages — threshold gating
// ─────────────────────────────────────────────────────────────────────────

describe('maybeCompressMessages — threshold gating', () => {
  it('no-op when below threshold', async () => {
    const messages = makeMessages(2, 10);
    const r = await maybeCompressMessages(messages, { budgetTokens: 100000 }, null);
    expect(r.compressed).toBe(false);
    expect(r.messages.length).toBe(messages.length);
    expect(r.summary).toBeNull();
  });

  it('triggers compression when above threshold (no LLM, fallback summary)', async () => {
    const messages = makeMessages(20, 200);
    const r = await maybeCompressMessages(messages, { budgetTokens: 1000, threshold: 0.5 }, null);
    expect(r.compressed).toBe(true);
    expect(r.messages.length).toBeLessThan(messages.length);
  });

  it('shouldCompressMessages predicate matches behavior', () => {
    const small = makeMessages(2, 10);
    const big = makeMessages(20, 200);
    expect(shouldCompressMessages(small, { budgetTokens: 100000 })).toBe(false);
    expect(shouldCompressMessages(big, { budgetTokens: 1000, threshold: 0.5 })).toBe(true);
  });

  it('custom estimateTokensFn honored', () => {
    const msgs = makeMessages(2, 10);
    const alwaysHigh = (() => 999999) as (m: ReadonlyArray<{ role: string; content: string }>) => number;
    expect(shouldCompressMessages(msgs, {
      budgetTokens: 1000,
      estimateTokensFn: alwaysHigh,
    })).toBe(true);
  });

  it('default threshold is 0.7', () => {
    const msgs = makeMessages(20, 200);
    // estimateTokens(msgs) > budgetTokens * 0.7?  Use the same default.
    const should = shouldCompressMessages(msgs, { budgetTokens: 1000 });
    // Conservative: just verify it can return either bool deterministically.
    expect(typeof should).toBe('boolean');
  });
});

// ─────────────────────────────────────────────────────────────────────────
// Unit: maybeCompressMessages — preservation guarantees
// ─────────────────────────────────────────────────────────────────────────

describe('maybeCompressMessages — preservation', () => {
  it('preserves system message + protected head verbatim', async () => {
    const messages = makeMessages(20, 200);
    const r = await maybeCompressMessages(messages, {
      budgetTokens: 1000,
      threshold: 0.3,
      protectedHeadMessages: 1,
    }, null);
    expect(r.compressed).toBe(true);
    expect(r.messages[0]).toEqual(messages[0]);    // system retained
    expect(r.messages[1]).toEqual(messages[1]);    // first user (kickoff) retained
  });

  it('preserves protected tail verbatim (last messages)', async () => {
    const messages = makeMessages(20, 200);
    const r = await maybeCompressMessages(messages, {
      budgetTokens: 1000,
      threshold: 0.3,
      protectedTailTokens: 200,
    }, null);
    expect(r.compressed).toBe(true);
    // The very last message in the result should match the very last in the input.
    expect(r.messages[r.messages.length - 1]).toEqual(messages[messages.length - 1]);
  });

  it('replaces middle with a single summary message', async () => {
    const messages = makeMessages(20, 200);
    const r = await maybeCompressMessages(messages, { budgetTokens: 1000, threshold: 0.3 }, null);
    expect(r.compressed).toBe(true);
    // Find the summary message (added between head and tail).
    const summaryMsg = r.messages.find(m => m.role === 'system' && m.content.startsWith('[Compressed:'));
    expect(summaryMsg).toBeDefined();
  });

  it('reduces token count after compression', async () => {
    const messages = makeMessages(30, 300);
    const r = await maybeCompressMessages(messages, { budgetTokens: 1000, threshold: 0.3 }, null);
    // No direct token count exposed here, but message count should drop.
    expect(r.messages.length).toBeLessThan(messages.length);
  });
});

// ─────────────────────────────────────────────────────────────────────────
// Unit: maybeCompressMessages — LLM summarizer
// ─────────────────────────────────────────────────────────────────────────

describe('maybeCompressMessages — LLM summarizer', () => {
  it('invokes llmCall when configured', async () => {
    let calls = 0;
    const llm: LlmCallFn = async () => {
      calls += 1;
      return { content: 'CONDENSED_FACTS', inTokens: 100, outTokens: 30, costUsd: 0.002, latencyMs: 80 };
    };
    const messages = makeMessages(20, 200);
    const r = await maybeCompressMessages(messages, {
      budgetTokens: 1000,
      threshold: 0.3,
      llmCall: llm,
      summarizationModel: 'budget-model',
    }, null);
    expect(calls).toBe(1);
    expect(r.compressed).toBe(true);
    expect(r.summary).toBe('CONDENSED_FACTS');
    const summaryMsg = r.messages.find(m => m.content.includes('CONDENSED_FACTS'));
    expect(summaryMsg).toBeDefined();
  });

  it('does NOT invoke llmCall when not configured (uses fallback)', async () => {
    const messages = makeMessages(20, 200);
    const r = await maybeCompressMessages(messages, { budgetTokens: 1000, threshold: 0.3 }, null);
    expect(r.compressed).toBe(true);
    expect(r.summary).toContain('Compressed region');  // fallback marker
  });

  it('falls back to placeholder when llmCall returns error', async () => {
    const llm: LlmCallFn = async () => ({
      content: '',
      inTokens: 0, outTokens: 0, costUsd: 0, latencyMs: 0,
      error: 'rate-limited',
    });
    const messages = makeMessages(20, 200);
    const r = await maybeCompressMessages(messages, {
      budgetTokens: 1000,
      threshold: 0.3,
      llmCall: llm,
      summarizationModel: 'budget-model',
    }, null);
    expect(r.compressed).toBe(true);
    expect(r.summary).toContain('Compressed region');
  });

  it('previousSummary carried forward into summarizer prompt', async () => {
    let receivedMessages: Array<{ role: string; content: string }> = [];
    const llm: LlmCallFn = async (input) => {
      receivedMessages = input.messages;
      return { content: 'NEW_SUMMARY', inTokens: 50, outTokens: 30, costUsd: 0.001, latencyMs: 50 };
    };
    const messages = makeMessages(20, 200);
    await maybeCompressMessages(messages, {
      budgetTokens: 1000,
      threshold: 0.3,
      llmCall: llm,
      summarizationModel: 'budget-model',
    }, 'PRIOR_SUMMARY_TEXT');
    // The summarizer should have received the previousSummary in its system prompt.
    const sysMsg = receivedMessages.find(m => m.role === 'system');
    expect(sysMsg?.content).toContain('PRIOR_SUMMARY_TEXT');
  });
});

// ─────────────────────────────────────────────────────────────────────────
// Unit: maybeCompressMessages — events
// ─────────────────────────────────────────────────────────────────────────

describe('maybeCompressMessages — events', () => {
  it('onCompressionEvent fires with correct payload when compressed', async () => {
    const events: MessagesCompressionEvent[] = [];
    const messages = makeMessages(20, 200);
    await maybeCompressMessages(messages, {
      budgetTokens: 1000,
      threshold: 0.3,
      onCompressionEvent: (e) => events.push(e),
    }, null);
    expect(events.length).toBe(1);
    expect(events[0]?.type).toBe('messages_compressed');
    expect(events[0]?.before_tokens).toBeGreaterThan(0);
    expect(events[0]?.after_tokens).toBeLessThan(events[0]!.before_tokens);
    expect(events[0]?.messages_before).toBe(messages.length);
    expect(events[0]?.cost_usd).toBe(0); // no llmCall, so 0
  });

  it('onCompressionEvent does NOT fire when no compression', async () => {
    const events: MessagesCompressionEvent[] = [];
    await maybeCompressMessages(makeMessages(2, 10), {
      budgetTokens: 100000,
      onCompressionEvent: (e) => events.push(e),
    }, null);
    expect(events.length).toBe(0);
  });

  it('onCompressionEvent records llmCall cost', async () => {
    const events: MessagesCompressionEvent[] = [];
    const llm: LlmCallFn = async () => ({
      content: 'GIST', inTokens: 100, outTokens: 30, costUsd: 0.0042, latencyMs: 50,
    });
    await maybeCompressMessages(makeMessages(20, 200), {
      budgetTokens: 1000,
      threshold: 0.3,
      llmCall: llm,
      summarizationModel: 'b',
      onCompressionEvent: (e) => events.push(e),
    }, null);
    expect(events[0]?.cost_usd).toBe(0.0042);
    expect(events[0]?.summary_generated).toBe(true);
  });
});

// ─────────────────────────────────────────────────────────────────────────
// Integration: runRetrievalAgentLoop + messagesContextManager
// ─────────────────────────────────────────────────────────────────────────

const fakeSearch: RetrievalSearchFn = async ({ query }) => ({
  formattedResults: `[result for "${query}"]\n` + 'Y'.repeat(500),
  resultCount: 1,
});

interface ScriptedReply {
  content: string;
  costUsd?: number;
}

function makeScriptedLlm(replies: ScriptedReply[]): LlmCallFn {
  let idx = 0;
  return async () => {
    const reply = replies[idx];
    idx += 1;
    if (!reply) {
      return { content: '{"action":"finalize","response":"(end)"}', inTokens: 10, outTokens: 10, costUsd: 0, latencyMs: 10 };
    }
    const r: LlmCallResult = {
      content: reply.content,
      inTokens: 100, outTokens: 30,
      costUsd: reply.costUsd ?? 0.001,
      latencyMs: 30,
    };
    return r;
  };
}

function baseConfig(overrides: Partial<MultiStepAgentRunConfig> = {}): MultiStepAgentRunConfig {
  return {
    modelAlias: 'claude-opus-4-7',
    persona: 'You are a test agent.',
    question: 'rank themes',
    llmCall: makeScriptedLlm([{ content: '{"action":"finalize","response":"done"}' }]),
    search: fakeSearch,
    maxSteps: 10,
    ...overrides,
  };
}

describe('runRetrievalAgentLoop — backwards compat with messages-compression hook', () => {
  it('omitted messagesContextManager → no compression triggered', async () => {
    const events: AgentRunProgressEvent[] = [];
    const llm = makeScriptedLlm([
      { content: '{"action":"retrieve","query":"q1"}' },
      { content: '{"action":"retrieve","query":"q2"}' },
      { content: '{"action":"finalize","response":"done"}' },
    ]);
    const result = await runRetrievalAgentLoop(baseConfig({
      llmCall: llm,
      onProgress: (e) => events.push(e),
    }));
    expect(result.rawResponse).toBe('done');
    expect(events.some(e => e.type === 'messages_compressed')).toBe(false);
  });
});

describe('runRetrievalAgentLoop — messagesContextManager active', () => {
  it('emits messages_compressed event when threshold crossed mid-loop', async () => {
    const events: AgentRunProgressEvent[] = [];
    // Loop with several retrievals; budget tight enough that compression fires.
    const llm = makeScriptedLlm([
      { content: '{"action":"retrieve","query":"q1"}' },
      { content: '{"action":"retrieve","query":"q2"}' },
      { content: '{"action":"retrieve","query":"q3"}' },
      { content: '{"action":"retrieve","query":"q4"}' },
      { content: '{"action":"finalize","response":"answer"}' },
    ]);
    const result = await runRetrievalAgentLoop(baseConfig({
      llmCall: llm,
      onProgress: (e) => events.push(e),
      messagesContextManager: {
        budgetTokens: 100,
        threshold: 0.4,
        retainRecentTurns: 2,
      },
      maxSteps: 6,
    }));
    expect(result.rawResponse).toBe('answer');
    expect(events.some(e => e.type === 'messages_compressed')).toBe(true);
  });

  it('messages_compressed event has step_index + count payloads', async () => {
    const events: AgentRunProgressEvent[] = [];
    const llm = makeScriptedLlm([
      { content: '{"action":"retrieve","query":"q1"}' },
      { content: '{"action":"retrieve","query":"q2"}' },
      { content: '{"action":"retrieve","query":"q3"}' },
      { content: '{"action":"retrieve","query":"q4"}' },
      { content: '{"action":"finalize","response":"a"}' },
    ]);
    await runRetrievalAgentLoop(baseConfig({
      llmCall: llm,
      onProgress: (e) => events.push(e),
      messagesContextManager: { budgetTokens: 100, threshold: 0.4, retainRecentTurns: 2 },
      maxSteps: 6,
    }));
    const compressEvt = events.find(e => e.type === 'messages_compressed');
    expect(compressEvt?.step_index).toBeGreaterThanOrEqual(0);
    expect(compressEvt?.messages_before_count).toBeGreaterThan(0);
    expect(compressEvt?.messages_after_count).toBeLessThan(compressEvt!.messages_before_count!);
  });

  it('compression does not break the retrieve → finalize flow', async () => {
    const llm = makeScriptedLlm([
      { content: '{"action":"retrieve","query":"q1"}' },
      { content: '{"action":"retrieve","query":"q2"}' },
      { content: '{"action":"retrieve","query":"q3"}' },
      { content: '{"action":"retrieve","query":"q4"}' },
      { content: '{"action":"retrieve","query":"q5"}' },
      { content: '{"action":"finalize","response":"all-good"}' },
    ]);
    const result = await runRetrievalAgentLoop(baseConfig({
      llmCall: llm,
      messagesContextManager: { budgetTokens: 500, threshold: 0.4 },
      maxSteps: 8,
    }));
    expect(result.rawResponse).toBe('all-good');
    expect(result.retrievalCalls).toBe(5);
  });

  it('multiple compressions over a long run', async () => {
    const events: AgentRunProgressEvent[] = [];
    // 8 retrieves before finalize — should trigger 2+ compressions on tight budget.
    const replies: ScriptedReply[] = [];
    for (let i = 0; i < 8; i += 1) {
      replies.push({ content: `{"action":"retrieve","query":"q${i}"}` });
    }
    replies.push({ content: '{"action":"finalize","response":"end"}' });
    const llm = makeScriptedLlm(replies);
    await runRetrievalAgentLoop(baseConfig({
      llmCall: llm,
      onProgress: (e) => events.push(e),
      messagesContextManager: { budgetTokens: 100, threshold: 0.4, retainRecentTurns: 2 },
      maxSteps: 10,
    }));
    const compressEvents = events.filter(e => e.type === 'messages_compressed');
    expect(compressEvents.length).toBeGreaterThanOrEqual(1);
  });

  it('compression onCompressionEvent (config-level) ALSO fires', async () => {
    const cfgEvents: MessagesCompressionEvent[] = [];
    const llm = makeScriptedLlm([
      { content: '{"action":"retrieve","query":"q1"}' },
      { content: '{"action":"retrieve","query":"q2"}' },
      { content: '{"action":"retrieve","query":"q3"}' },
      { content: '{"action":"finalize","response":"a"}' },
    ]);
    await runRetrievalAgentLoop(baseConfig({
      llmCall: llm,
      messagesContextManager: {
        budgetTokens: 300,
        threshold: 0.4,
        retainRecentTurns: 2,
        onCompressionEvent: (e) => cfgEvents.push(e),
      },
      maxSteps: 5,
    }));
    expect(cfgEvents.length).toBeGreaterThanOrEqual(1);
    expect(cfgEvents[0]?.type).toBe('messages_compressed');
  });

  it('compression engaged + LLM summarizer end-to-end', async () => {
    const events: AgentRunProgressEvent[] = [];
    const summarizerLlm = makeFakeLlm('SUMMARY_TEXT', 0.005);
    const llm = makeScriptedLlm([
      { content: '{"action":"retrieve","query":"q1"}' },
      { content: '{"action":"retrieve","query":"q2"}' },
      { content: '{"action":"retrieve","query":"q3"}' },
      { content: '{"action":"finalize","response":"r"}' },
    ]);
    const cfgEvents: MessagesCompressionEvent[] = [];
    await runRetrievalAgentLoop(baseConfig({
      llmCall: llm,
      onProgress: (e) => events.push(e),
      messagesContextManager: {
        budgetTokens: 300,
        threshold: 0.4,
        retainRecentTurns: 2,
        llmCall: summarizerLlm,
        summarizationModel: 'budget-model',
        onCompressionEvent: (e) => cfgEvents.push(e),
      },
      maxSteps: 5,
    }));
    expect(cfgEvents.some(e => e.summary_generated === true)).toBe(true);
    expect(cfgEvents.some(e => e.cost_usd > 0)).toBe(true);
  });

  it('replay determinism: two runs with identical inputs → identical compression decisions', async () => {
    const buildLlm = () => makeScriptedLlm([
      { content: '{"action":"retrieve","query":"q1"}' },
      { content: '{"action":"retrieve","query":"q2"}' },
      { content: '{"action":"retrieve","query":"q3"}' },
      { content: '{"action":"retrieve","query":"q4"}' },
      { content: '{"action":"finalize","response":"deterministic"}' },
    ]);
    const events1: AgentRunProgressEvent[] = [];
    const events2: AgentRunProgressEvent[] = [];
    const result1 = await runRetrievalAgentLoop(baseConfig({
      llmCall: buildLlm(),
      onProgress: (e) => events1.push(e),
      messagesContextManager: { budgetTokens: 100, threshold: 0.4, retainRecentTurns: 2 },
      maxSteps: 6,
    }));
    const result2 = await runRetrievalAgentLoop(baseConfig({
      llmCall: buildLlm(),
      onProgress: (e) => events2.push(e),
      messagesContextManager: { budgetTokens: 100, threshold: 0.4, retainRecentTurns: 2 },
      maxSteps: 6,
    }));
    expect(result1.rawResponse).toBe(result2.rawResponse);
    expect(events1.filter(e => e.type === 'messages_compressed').length)
      .toBe(events2.filter(e => e.type === 'messages_compressed').length);
  });
});