import { describe, it, expect } from 'vitest';
import {
  guardPromptBeforeSend,
  compactOversizedToolResults,
  parsePromptSafeLimitTokens,
  buildPromptLimitAgentInstruction,
  looksLikeLargeEncodedPayload,
  LARGE_TOOL_RESULT_TOKENS,
} from './prompt-guard.js';
import { ContextManager } from '../context-manager.js';
import type { Message, ToolDef } from '../../llm/openai-compat.js';

function readPair(callId: string, filePath: string, content: string): Message[] {
  return [
    {
      role: 'assistant',
      tool_calls: [{
        id: callId,
        type: 'function',
        function: { name: 'Read', arguments: JSON.stringify({ file_path: filePath }) },
      }],
    },
    { role: 'tool', tool_call_id: callId, content },
  ];
}

function bashTurn(callId: string, output: string): Message[] {
  return [
    {
      role: 'assistant',
      tool_calls: [{
        id: callId,
        type: 'function',
        function: { name: 'Bash', arguments: JSON.stringify({ command: 'noop' }) },
      }],
    },
    { role: 'tool', tool_call_id: callId, content: output },
  ];
}

const NO_TOOLS: ToolDef[] = [];

function asciiApproxTokens(tokens: number): string {
  return 'A'.repeat(Math.ceil(tokens * 3.6));
}

describe('parsePromptSafeLimitTokens', () => {
  it('parses comma-formatted token counts', () => {
    expect(parsePromptSafeLimitTokens('LLM request blocked before send: ... safe limit 24,000 tokens ...'))
      .toBe(24_000);
  });
  it('parses plain token counts', () => {
    expect(parsePromptSafeLimitTokens('safe limit 1000 tokens')).toBe(1000);
  });
  it('returns null when not present', () => {
    expect(parsePromptSafeLimitTokens('some other error')).toBeNull();
  });
});

describe('looksLikeLargeEncodedPayload', () => {
  it('returns false for short text', () => {
    expect(looksLikeLargeEncodedPayload('A'.repeat(7_000))).toBe(false);
  });
  it('detects base64 data URLs (above the 8k length threshold)', () => {
    const payload = '<html><img src="data:image/png;base64,' + 'A'.repeat(9_000) + '">';
    expect(looksLikeLargeEncodedPayload(payload)).toBe(true);
  });
  it('detects loose base64 markers (above the 8k length threshold)', () => {
    const payload = 'prefix '.repeat(200) + 'base64: ' + 'X'.repeat(9_000);
    expect(looksLikeLargeEncodedPayload(payload)).toBe(true);
  });
});

describe('buildPromptLimitAgentInstruction', () => {
  it('mentions estimated and max tokens with locale formatting', () => {
    const msg = buildPromptLimitAgentInstruction(120_000, 100_000);
    expect(msg).toContain('120,000');
    expect(msg).toContain('100,000');
    expect(msg).toMatch(/Read\(offset\/limit\)/);
  });
});

describe('compactOversizedToolResults', () => {
  it('replaces only tool messages above LARGE_TOOL_RESULT_TOKENS', () => {
    const small = 'X'.repeat(100);
    const big = asciiApproxTokens(LARGE_TOOL_RESULT_TOKENS + 1);
    const messages: Message[] = [
      { role: 'system', content: 'sys' },
      { role: 'user', content: 'task' },
      { role: 'tool', tool_call_id: 't1', content: small },
      { role: 'tool', tool_call_id: 't2', content: big },
    ];
    const result = compactOversizedToolResults(messages, 0, 1_000);
    expect(result.changed).toBe(true);
    expect(result.omittedCount).toBe(1);
    expect(messages[2]!.content).toBe(small); // small untouched
    expect(messages[3]!.content).toMatch(/Tool result omitted/);
  });

  it('returns unchanged when prompt already fits', () => {
    const messages: Message[] = [
      { role: 'tool', tool_call_id: 't', content: 'X'.repeat(LARGE_TOOL_RESULT_TOKENS + 100) },
    ];
    const result = compactOversizedToolResults(messages, 0, 1_000_000);
    expect(result.changed).toBe(false);
    expect(messages[0]!.content).toMatch(/^X+$/);
  });

  it('processes candidates largest-first to converge faster', () => {
    // Total is just over 50k tokens. Replacing the largest candidate alone
    // drops the prompt under budget, so the smaller candidate stays intact.
    const huge = asciiApproxTokens(42_000);
    const big = 'B'.repeat(Math.ceil(17_000 * 3.6));
    const messages: Message[] = [
      { role: 'tool', tool_call_id: 'a', content: big },
      { role: 'tool', tool_call_id: 'b', content: huge },
    ];
    compactOversizedToolResults(messages, 0, 50_000);
    expect(messages[1]!.content).toMatch(/Tool result omitted/);
    expect(messages[0]!.content).toMatch(/^B+$/); // untouched
  });

  it('annotates encoded payloads with a base64 hint', () => {
    const encoded = 'data:image/png;base64,' + asciiApproxTokens(LARGE_TOOL_RESULT_TOKENS + 1);
    const messages: Message[] = [
      { role: 'tool', tool_call_id: 'e', content: encoded },
    ];
    compactOversizedToolResults(messages, 0, 100);
    expect(messages[0]!.content).toContain('base64/data URLs');
  });
});

describe('guardPromptBeforeSend', () => {
  it('returns ok without work when there is no contextManager', async () => {
    const messages: Message[] = [{ role: 'user', content: 'hi' }];
    const result = await guardPromptBeforeSend(messages, NO_TOOLS, undefined);
    expect(result.ok).toBe(true);
    if (result.ok) {
      expect(result.deduped).toBe(false);
      expect(result.compacted).toBe(false);
      expect(result.summarized).toBe(false);
    }
  });

  it('returns ok with no stages run when prompt is well under budget', async () => {
    const cm = new ContextManager({ limitTokens: 100_000 });
    const messages: Message[] = [
      { role: 'system', content: 'sys' },
      { role: 'user', content: 'short task' },
    ];
    const result = await guardPromptBeforeSend(messages, NO_TOOLS, cm);
    expect(result.ok).toBe(true);
    if (result.ok) {
      expect(result.deduped).toBe(false);
      expect(result.compacted).toBe(false);
      expect(result.summarized).toBe(false);
    }
  });

  it('stage 1: dedup alone resolves overflow when only file-reads accumulated', async () => {
    const cm = new ContextManager({ limitTokens: 30_000 });
    const big = asciiApproxTokens(14_000);
    const messages: Message[] = [
      { role: 'system', content: 'sys' },
      { role: 'user', content: 'task' },
      ...readPair('r1', '/dup.ts', big),
      ...readPair('r2', '/dup.ts', big),
    ];
    const result = await guardPromptBeforeSend(messages, NO_TOOLS, cm);
    expect(result.ok).toBe(true);
    if (result.ok) {
      expect(result.deduped).toBe(true);
      expect(result.compacted).toBe(false);
      expect(result.summarized).toBe(false);
    }
  });

  it('stage 2: compaction kicks in when tool result is large but distinct', async () => {
    const cm = new ContextManager({ limitTokens: 20_000 });
    const huge = asciiApproxTokens(20_000);
    const messages: Message[] = [
      { role: 'system', content: 'sys' },
      { role: 'user', content: 'task' },
      ...bashTurn('b1', huge),
    ];
    const result = await guardPromptBeforeSend(messages, NO_TOOLS, cm);
    expect(result.ok).toBe(true);
    if (result.ok) {
      expect(result.compacted).toBe(true);
      expect(result.summarized).toBe(false);
    }
    // Bash result was replaced with placeholder
    expect((messages.find(m => m.role === 'tool')?.content as string)).toMatch(/Tool result omitted/);
  });

  // Helpers for the stage-3 overflow scenarios: each individual turn is below
  // LARGE_TOOL_RESULT_TOKENS so compaction cannot help, but the accumulated
  // history exceeds the 24k guard at limitTokens=30_000.
  function makeOverflowingHistory(): Message[] {
    const moderate = asciiApproxTokens(6_000);
    const turns: Message[] = [];
    for (let i = 1; i <= 8; i++) turns.push(...bashTurn(`b${i}`, moderate));
    return [
      { role: 'system', content: 'sys' },
      { role: 'user', content: 'task' },
      ...turns,
    ];
  }

  it('stage 3: summarization runs when dedup+compact still leave overflow', async () => {
    const cm = new ContextManager({ limitTokens: 30_000 });
    const messages = makeOverflowingHistory();
    const result = await guardPromptBeforeSend(messages, NO_TOOLS, cm, {
      runIsolatedLlm: async () => '## ゴール\nrun bash\n## 進捗\nDone: 8',
    });
    expect(result.ok).toBe(true);
    if (result.ok) expect(result.summarized).toBe(true);
  });

  it('stage 3 is skipped when historySummarization.enabled is false', async () => {
    const cm = new ContextManager({ limitTokens: 30_000 });
    const messages = makeOverflowingHistory();
    let llmCalled = false;
    const result = await guardPromptBeforeSend(messages, NO_TOOLS, cm, {
      historySummarization: { enabled: false },
      runIsolatedLlm: async () => { llmCalled = true; return 'never'; },
    });
    expect(llmCalled).toBe(false);
    expect(result.ok).toBe(false);
  });

  it('stage 3 is skipped when runIsolatedLlm is missing', async () => {
    const cm = new ContextManager({ limitTokens: 30_000 });
    const messages = makeOverflowingHistory();
    const result = await guardPromptBeforeSend(messages, NO_TOOLS, cm);
    expect(result.ok).toBe(false);
  });

  it('returns ok:false with structured message when all stages fail', async () => {
    const cm = new ContextManager({ limitTokens: 1_000 });
    // Original user task itself is too big — none of the stages can shrink it.
    const messages: Message[] = [
      { role: 'system', content: 'sys' },
      { role: 'user', content: 'X'.repeat(10_000) },
    ];
    const result = await guardPromptBeforeSend(messages, NO_TOOLS, cm);
    expect(result.ok).toBe(false);
    if (!result.ok) {
      expect(result.message).toContain('LLM request blocked before send');
      expect(result.message).toContain('safe limit');
      expect(result.limitTokens).toBe(1_000);
    }
  });

  it('honors a custom promptGuardRatio', async () => {
    const cm = new ContextManager({ limitTokens: 100_000 });
    // Under 80% (80k) but over 60% (60k).
    const messages: Message[] = [
      { role: 'system', content: 'sys' },
      { role: 'user', content: asciiApproxTokens(70_000) },
    ];
    const okAt08 = await guardPromptBeforeSend(messages, NO_TOOLS, cm, { promptGuardRatio: 0.8 });
    expect(okAt08.ok).toBe(true);
    const failAt06 = await guardPromptBeforeSend(messages, NO_TOOLS, cm, { promptGuardRatio: 0.6 });
    expect(failAt06.ok).toBe(false);
  });

  it('falls through to ok:false when summarizer LLM throws', async () => {
    const cm = new ContextManager({ limitTokens: 30_000 });
    const messages = makeOverflowingHistory();
    const result = await guardPromptBeforeSend(messages, NO_TOOLS, cm, {
      runIsolatedLlm: async () => { throw new Error('LLM down'); },
    });
    expect(result.ok).toBe(false);
  });
});
