import { describe, it, expect, beforeEach, jest } from '@jest/globals';
import { TokenLimiterPlugin } from '../../../src/guardrails/plugins/token-limiter';
import { createGuardrailContext } from '../../../src/guardrails/context';

// Mock logger to avoid console noise during tests
jest.mock('../../../src/utils/logger');

describe('TokenLimiterPlugin', () => {
  let plugin: TokenLimiterPlugin;

  beforeEach(async () => {
    plugin = new TokenLimiterPlugin();
  });

  describe('initialization', () => {
    it('should initialize with default values', async () => {
      await plugin.initialize({ enabled: true });

      expect(plugin.enabled).toBe(true);
      expect(plugin.name).toBe('token_limiter');
      expect(plugin.phases).toContain('pre_request');
      expect(plugin.getLimits().maxInputTokens).toBe(8192);
    });

    it('should initialize with custom config', async () => {
      await plugin.initialize({
        enabled: true,
        priority: 15,
        max_input_tokens: 4096,
        max_output_tokens: 2048,
        warn_at_percentage: 90,
      });

      expect(plugin.priority).toBe(15);
      const limits = plugin.getLimits();
      expect(limits.maxInputTokens).toBe(4096);
      expect(limits.maxOutputTokens).toBe(2048);
    });
  });

  describe('token estimation', () => {
    it('should estimate tokens from text', async () => {
      await plugin.initialize({ enabled: true });

      // ~4 chars per token
      expect(plugin.estimateTokenCount('Hello')).toBeGreaterThan(0);
      expect(plugin.estimateTokenCount('Hello world')).toBeGreaterThan(plugin.estimateTokenCount('Hello'));
    });

    it('should return 0 for empty text', async () => {
      await plugin.initialize({ enabled: true });

      expect(plugin.estimateTokenCount('')).toBe(0);
    });
  });

  describe('token limiting', () => {
    it('should allow prompts under the limit', async () => {
      await plugin.initialize({
        enabled: true,
        max_input_tokens: 1000,
      });

      const shortPrompt = 'Hello world'; // ~7 tokens
      const context = createGuardrailContext({ prompt: shortPrompt });
      const result = await plugin.execute('pre_request', context);

      expect(result.action).toBe('allow');
    });

    it('should block prompts over the limit', async () => {
      await plugin.initialize({
        enabled: true,
        max_input_tokens: 10, // Very small limit
      });

      // Create a prompt that exceeds the limit (~4 chars per token)
      const longPrompt = 'This is a longer prompt that will definitely exceed our tiny token limit';
      const context = createGuardrailContext({ prompt: longPrompt });
      const result = await plugin.execute('pre_request', context);

      expect(result.action).toBe('block');
      expect(result.blockedBy).toBe('token_limiter');
      expect(result.blockReason).toContain('Token limit exceeded');
    });

    it('should include messages in token count', async () => {
      await plugin.initialize({
        enabled: true,
        max_input_tokens: 20,
      });

      const context = createGuardrailContext({
        prompt: 'Short',
        messages: [
          { role: 'user', content: 'This is a message with many tokens that should push us over the limit', timestamp: new Date() },
        ],
      });

      const result = await plugin.execute('pre_request', context);
      expect(result.action).toBe('block');
    });
  });

  describe('warnings', () => {
    it('should add warning when approaching limit', async () => {
      await plugin.initialize({
        enabled: true,
        max_input_tokens: 100,
        warn_at_percentage: 50, // Warn at 50%
      });

      // Create prompt that's ~60-80% of limit
      const prompt = 'A'.repeat(300); // ~75 tokens
      const context = createGuardrailContext({ prompt });
      await plugin.execute('pre_request', context);

      const warnings = context.violations.filter((v) => v.severity === 'warning');
      expect(warnings.length).toBeGreaterThan(0);
    });

    it('should not warn when well under limit', async () => {
      await plugin.initialize({
        enabled: true,
        max_input_tokens: 1000,
        warn_at_percentage: 80,
      });

      const shortPrompt = 'Hello';
      const context = createGuardrailContext({ prompt: shortPrompt });
      await plugin.execute('pre_request', context);

      const warnings = context.violations.filter((v) => v.severity === 'warning');
      expect(warnings.length).toBe(0);
    });
  });

  describe('phase handling', () => {
    it('should process pre_request and post_response phases', async () => {
      await plugin.initialize({ enabled: true });

      expect(plugin.phases).toContain('pre_request');
      expect(plugin.phases).toContain('post_response');
    });

    it('should allow other phases', async () => {
      await plugin.initialize({ enabled: true });

      const context = createGuardrailContext({ prompt: 'test' });

      // pre_tool_input should pass through
      const result = await plugin.execute('pre_tool_input', context);
      expect(result.action).toBe('allow');
    });
  });

  describe('output token limiting', () => {
    it('should allow responses under the limit', async () => {
      await plugin.initialize({
        enabled: true,
        max_output_tokens: 1000,
      });

      const context = createGuardrailContext({ prompt: 'test' });
      context.response = 'Short response';
      const result = await plugin.execute('post_response', context);

      expect(result.action).toBe('allow');
    });

    it('should block responses over the limit', async () => {
      await plugin.initialize({
        enabled: true,
        max_output_tokens: 10, // Very small limit
      });

      const context = createGuardrailContext({ prompt: 'test' });
      context.response = 'This is a longer response that will definitely exceed our tiny token limit for output';
      const result = await plugin.execute('post_response', context);

      expect(result.action).toBe('block');
      expect(result.blockedBy).toBe('token_limiter');
      expect(result.blockReason).toContain('Output token limit exceeded');
    });

    it('should skip output check if no max_output_tokens configured', async () => {
      await plugin.initialize({
        enabled: true,
        // No max_output_tokens
      });

      const context = createGuardrailContext({ prompt: 'test' });
      context.response = 'A'.repeat(10000); // Very long response
      const result = await plugin.execute('post_response', context);

      expect(result.action).toBe('allow');
    });

    it('should add warning when response approaches output limit', async () => {
      await plugin.initialize({
        enabled: true,
        max_output_tokens: 100,
        warn_at_percentage: 50, // Warn at 50%
      });

      // Create response that's ~60-80% of limit
      const context = createGuardrailContext({ prompt: 'test' });
      context.response = 'A'.repeat(300); // ~75 tokens
      await plugin.execute('post_response', context);

      const warnings = context.violations.filter((v) => v.severity === 'warning');
      expect(warnings.length).toBeGreaterThan(0);
      expect(warnings[0].rule).toBe('max_output_tokens_warning');
    });
  });
});