import { describe, it, expect } from 'vitest';
import { checkContextBudget } from '../../src/utils/context-budget.js';

const defaultConfig = {
  maxInputChars: 30000,
  maxFileChars: 12000,
  maxFiles: 5,
  maxLogChars: 10000,
};

function makeMinimalInput(question) {
  return { question };
}

describe('checkContextBudget', () => {
  it('returns ok true with no warnings for minimal input within limits', () => {
    const result = checkContextBudget(makeMinimalInput('Hello'), defaultConfig);
    expect(result.ok).toBe(true);
    expect(result.warnings).toEqual([]);
    expect(result.input.question).toBe('Hello');
  });

  it('throws when question is missing', () => {
    expect(() => checkContextBudget({}, defaultConfig)).toThrow('question is required');
  });

  it('throws when question is empty string', () => {
    expect(() => checkContextBudget({ question: '' }, defaultConfig)).toThrow('question is required');
  });

  it('throws when question is null', () => {
    expect(() => checkContextBudget({ question: null }, defaultConfig)).toThrow();
  });

  it('drops excess files and adds a warning', () => {
    const files = Array.from({ length: 7 }, (_, i) => ({
      path: `src/file${i}.js`,
      content: 'x',
    }));
    const result = checkContextBudget(
      { question: 'ok', relevantFiles: files },
      defaultConfig,
    );
    expect(result.ok).toBe(true);
    expect(result.input.relevantFiles.length).toBe(5);
    expect(result.warnings[0]).toContain('first 5');
    expect(result.warnings[0]).toContain('maxFiles=5');
  });

  it('truncates a single file that exceeds maxFileChars', () => {
    const longContent = 'x'.repeat(15000);
    const result = checkContextBudget(
      { question: 'ok', relevantFiles: [{ path: 'src/long.js', content: longContent }] },
      defaultConfig,
    );
    expect(result.ok).toBe(true);
    expect(result.input.relevantFiles[0].content.length).toBe(12000);
    expect(result.warnings[0]).toContain('src/long.js');
    expect(result.warnings[0]).toContain('truncated to 12000');
  });

  it('truncates logs that exceed maxLogChars', () => {
    const longLogs = 'x'.repeat(15000);
    const result = checkContextBudget({ question: 'ok', logs: longLogs }, defaultConfig);
    expect(result.ok).toBe(true);
    expect(result.input.logs.length).toBe(10000);
    expect(result.warnings[0]).toContain('Logs were truncated to 10000');
  });

  it('applies all hard limits and fails when total still exceeds budget', () => {
    // After hard limits: 1 + 5*12000 + 10000 = 70001 > 30000, no droppable content
    const files = Array.from({ length: 7 }, (_, i) => ({
      path: `src/file${i}.js`,
      content: 'y'.repeat(15000),
    }));
    const result = checkContextBudget(
      { question: 'q', relevantFiles: files, logs: 'z'.repeat(15000) },
      defaultConfig,
    );
    expect(result.ok).toBe(false);
    expect(result.input).toBeNull();
    // Should have truncation warnings even though final result is failure
    expect(result.warnings.length).toBeGreaterThanOrEqual(2);
  });

  it('keeps context intact when total already within budget', () => {
    const context = 'c'.repeat(20000);
    const question = 'q';
    const result = checkContextBudget({ question, context }, defaultConfig);
    expect(result.ok).toBe(true);
    expect(result.input.context.length).toBe(20000); // no trimming needed
    expect(result.warnings).toEqual([]);
  });

  it('trims context when total over budget but projectSummary fits exactly after', () => {
    const context = 'c'.repeat(25000);
    const projectSummary = 'p'.repeat(10000);
    // Total: 1 + 25000 + 10000 = 35001 > 30000, need to trim
    const result = checkContextBudget({ question: 'q', context, projectSummary }, defaultConfig);
    expect(result.ok).toBe(true);
    // Context trimmed first: allowed = max(0, 30000 - 35001 + 25000) = 19999
    expect(result.input.context.length).toBe(19999);
    expect(result.warnings.some((w) => w.includes('Context was trimmed'))).toBe(true);
    // After context trim: remaining = 1 + 19999 + 10000 = 30000 — exactly at limit.
    // Project summary is NOT trimmed (fits exactly).
    expect(result.input.projectSummary.length).toBe(10000);
    expect(result.warnings.some((w) => w.includes('Project summary was trimmed'))).toBe(false);
  });

  it('returns ok false when question alone exceeds maxInputChars', () => {
    const hugeQuestion = 'q'.repeat(35001);
    const result = checkContextBudget({ question: hugeQuestion }, defaultConfig);
    expect(result.ok).toBe(false);
    expect(result.input).toBeNull();
  });

  it('returns ok false when hard-limited content alone exceeds budget with no droppable sections', () => {
    // Each file oversized to trigger truncation; after trunc: 5 * 12000 + 100(question) = 60100 > 30000
    const result = checkContextBudget(
      { question: 'q'.repeat(100), relevantFiles: Array.from({ length: 5 }, () => ({ path: 'a.js', content: 'x'.repeat(20000) })) },
      { ...defaultConfig, maxFileChars: 12000, maxInputChars: 30000 },
    );
    expect(result.ok).toBe(false);
    expect(result.input).toBeNull();
  });

  it('excludes empty optional fields from output', () => {
    const result = checkContextBudget({ question: 'hi' }, defaultConfig);
    expect(result.input.context).toBeUndefined();
    expect(result.input.projectSummary).toBeUndefined();
    expect(result.input.taskSummary).toBeUndefined();
    expect(result.input.relevantFiles).toBeUndefined();
    expect(result.input.logs).toBeUndefined();
  });

  it('preserves non-empty optional fields in output', () => {
    const result = checkContextBudget(
      { question: 'hi', context: 'ctx', logs: 'log' },
      defaultConfig,
    );
    expect(result.input.context).toBe('ctx');
    expect(result.input.logs).toBe('log');
  });

  it('handles tiny config limits forcing aggressive trimming', () => {
    const tinyConfig = { maxInputChars: 50, maxFileChars: 10, maxFiles: 1, maxLogChars: 5 };
    const result = checkContextBudget(
      { question: 'x', context: 'cccccc', relevantFiles: [{ path: 'a.js', content: '1234567890' }], logs: 'lllllll' },
      tinyConfig,
    );
    expect(result.ok).toBe(true);
  });

  it('returns ok false after exhausting all droppable content', () => {
    const result = checkContextBudget(
      { question: 'q'.repeat(100), context: 'c'.repeat(20000), projectSummary: 'p'.repeat(20000) },
      { ...defaultConfig, maxInputChars: 50 },
    );
    expect(result.ok).toBe(false);
  });

  it('does not mutate primitive strings but truncates file objects in shallow copy', () => {
    const files = [{ path: 'a.js', content: 'x'.repeat(20000) }];
    const logs = 'l'.repeat(20000);
    const context = 'c'.repeat(40000);
    const input = { question: 'q', context, relevantFiles: files, logs };
    checkContextBudget(input, defaultConfig);
    // File objects in shallow copy are mutated (content truncated)
    expect(files[0].content.length).toBe(12000);
    // Primitive strings are not mutated (only local variable rebinding)
    expect(context.length).toBe(40000);
    expect(logs.length).toBe(20000);
  });

  it('adds warning for file truncation with all files needing truncation', () => {
    const files = [{ path: 'src/a.js', content: 'x'.repeat(15000) }, { path: 'src/b.js', content: 'y'.repeat(15000) }];
    const result = checkContextBudget({ question: 'q', relevantFiles: files }, defaultConfig);
    expect(result.ok).toBe(true);
    // Should have 2 file truncation warnings (one per file over limit)
    const fileWarnings = result.warnings.filter((w) => w.includes('truncated'));
    expect(fileWarnings.length).toBe(2);
  });

  it('adds all expected warnings when multiple trimming actions occur', () => {
    const files = Array.from({ length: 6 }, (_, i) => ({ path: `src/file${i}.js`, content: 'x' }));
    const logs = 'l'.repeat(15000);
    const context = 'c'.repeat(25000);
    const result = checkContextBudget(
      { question: 'q', context, relevantFiles: files, logs },
      defaultConfig,
    );
    expect(result.ok).toBe(true);
    expect(result.warnings.some((w) => w.includes('Only the first'))).toBe(true);
    expect(result.warnings.some((w) => w.includes('Logs were truncated'))).toBe(true);
    expect(result.warnings.some((w) => w.includes('Context was trimmed'))).toBe(true);
  });

  it('trims context and projectSummary when both contribute to budget overflow', () => {
    const context = 'c'.repeat(28000);
    const projectSummary = 'p'.repeat(3000);
    const taskSummary = 't'.repeat(1500);
    // Total: 1 + 28000 + 3000 + 1500 = 32501 > 30000
    // Context trimmed to max(0, 30000 - 32501 + 28000) = 25499
    const result = checkContextBudget({ question: 'q', context, projectSummary, taskSummary }, defaultConfig);
    expect(result.ok).toBe(true);
    expect(result.input.context.length).toBe(25499);
    expect(result.warnings.some((w) => w.includes('Context was trimmed'))).toBe(true);
    // After: remaining = 1 + 25499 + 3000 + 1500 = 30000 — exactly at limit. No more trimming.
    expect(result.input.projectSummary.length).toBe(3000);
    expect(result.warnings.some((w) => w.includes('Project summary was trimmed'))).toBe(false);
    expect(result.input.taskSummary).toEqual(taskSummary);
  });

  it('trims all three droppable sections when they are all large', () => {
    const context = 'c'.repeat(20000);
    const projectSummary = 'p'.repeat(10000);
    const taskSummary = 't'.repeat(10000);
    // Total: 1 + 20000 + 10000 + 10000 = 40001 > 30000
    const result = checkContextBudget({ question: 'q', context, projectSummary, taskSummary }, defaultConfig);
    expect(result.ok).toBe(true);
    // Context trimmed to max(0, 30000 - 40001 + 20000) = 9999
    expect(result.input.context.length).toBe(9999);
    expect(result.warnings.some((w) => w.includes('Context was trimmed'))).toBe(true);
    // After: remaining = 1 + 9999 + 10000 + 10000 = 30000 — exactly at limit. No more trimming.
    expect(result.input.projectSummary.length).toBe(10000);
    expect(result.input.taskSummary.length).toBe(10000);
  });

  it('trims projectSummary when context is within budget but total still over', () => {
    const context = 'c'.repeat(5000);
    const projectSummary = 'p'.repeat(28000);
    // Total: 1 + 5000 + 28000 = 33001 > 30000
    // Context allowed = max(0, 30000 - 33001 + 5000) = 1999
    const result = checkContextBudget({ question: 'q', context, projectSummary }, defaultConfig);
    expect(result.ok).toBe(true);
    expect(result.input.context.length).toBe(1999);
    // After: remaining = 1 + 1999 + 28000 = 30000 — exactly at limit. No more trimming.
    expect(result.input.projectSummary.length).toBe(28000);
    expect(result.warnings.some((w) => w.includes('Context was trimmed'))).toBe(true);
    expect(result.warnings.some((w) => w.includes('Project summary was trimmed'))).toBe(false);
  });

  it('question + truncated logs exceeds budget with no droppable sections', () => {
    const longLogs = 'l'.repeat(30000);
    const result = checkContextBudget(
      { question: 'q'.repeat(100), logs: longLogs },
      { ...defaultConfig, maxLogChars: 10000 },
    );
    // After truncation: 100 + 10000 = 10100 ≤ 30000 → should pass (no droppable needed)
    expect(result.ok).toBe(true);
  });

  it('question + truncated logs exceeds budget with no other content', () => {
    const longLogs = 'l'.repeat(29950);
    const result = checkContextBudget(
      { question: 'q'.repeat(100), logs: longLogs },
      { ...defaultConfig, maxLogChars: 10000, maxInputChars: 30000 },
    );
    // After truncation: 100 + 10000 = 10100 ≤ 30000 → should pass
    expect(result.ok).toBe(true);
  });

  it('returns ok false when question + truncated logs exceeds tiny budget', () => {
    const longLogs = 'l'.repeat(29950);
    const result = checkContextBudget(
      { question: 'q'.repeat(100), logs: longLogs },
      { ...defaultConfig, maxLogChars: 10000, maxInputChars: 5000 },
    );
    // After truncation: 100 + 10000 = 10100 > 5000 → fails
    expect(result.ok).toBe(false);
    expect(result.input).toBeNull();
  });
});