import {
  estimateTokenCount,
  StreamingTokenCounter,
  validateInputTokens,
  createTokenAwareStream
} from '../TokenValidation.js';

describe('TokenValidation', () => {
  describe('estimateTokenCount', () => {
    it('should estimate tokens from string length', () => {
      const text = 'Hello world'; // ~11 chars
      const tokens = estimateTokenCount(text);
      expect(tokens).toBeGreaterThan(0);
      expect(typeof tokens).toBe('number');
    });

    it('should estimate tokens from character count', () => {
      const charCount = 100;
      const tokens = estimateTokenCount(charCount);
      expect(tokens).toBeGreaterThan(0);
      expect(typeof tokens).toBe('number');
    });

    it('should return 0 for empty string', () => {
      expect(estimateTokenCount('')).toBe(0);
    });

    it('should return 0 for invalid number input', () => {
      expect(estimateTokenCount(0)).toBe(0);
      expect(estimateTokenCount(-10)).toBe(0);
      expect(estimateTokenCount(Infinity)).toBe(0);
      expect(estimateTokenCount(NaN)).toBe(0);
    });

    it('should handle large text', () => {
      const largeText = 'a'.repeat(10000);
      const tokens = estimateTokenCount(largeText);
      expect(tokens).toBeGreaterThan(1000);
    });

    it('should return 0 for null/undefined', () => {
      expect(estimateTokenCount(null as any)).toBe(0);
      expect(estimateTokenCount(undefined as any)).toBe(0);
    });
  });

  describe('StreamingTokenCounter', () => {
    it('should initialize with correct defaults', () => {
      const counter = new StreamingTokenCounter(1000, 'test-provider');
      const stats = counter.getStats();

      expect(stats.tokenCount).toBe(0);
      expect(stats.charCount).toBe(0);
      expect(stats.maxTokens).toBe(1000);
      expect(stats.utilization).toBe(0);
    });

    it('should update token count when adding chunks', () => {
      const counter = new StreamingTokenCounter(1000, 'test');
      const result = counter.addChunk('Hello world');

      expect(result.tokenCount).toBeGreaterThan(0);
      expect(result.exceeded).toBe(false);
      expect(result.shouldStop).toBe(false);
    });

    it('should detect when approaching limit (80%)', () => {
      const counter = new StreamingTokenCounter(100, 'test');

      // Add enough text to reach ~80% of limit
      const chunk = 'a'.repeat(320); // ~80 tokens at 4 chars/token
      counter.addChunk(chunk);

      expect(counter.isApproachingLimit()).toBe(true);
    });

    it('should return shouldStop at 95% of limit', () => {
      const counter = new StreamingTokenCounter(100, 'test');

      // Add enough text to reach 95% of limit
      const chunk = 'a'.repeat(380); // ~95 tokens at 4 chars/token
      const result = counter.addChunk(chunk);

      expect(result.shouldStop).toBe(true);
    });

    it('should mark as exceeded when over limit', () => {
      const counter = new StreamingTokenCounter(10, 'test');

      // Add text exceeding limit
      const chunk = 'a'.repeat(100); // Way over 10 tokens
      const result = counter.addChunk(chunk);

      expect(result.exceeded).toBe(true);
    });

    it('should reset counter state', () => {
      const counter = new StreamingTokenCounter(1000, 'test');
      counter.addChunk('Hello world');

      counter.reset();
      const stats = counter.getStats();

      expect(stats.tokenCount).toBe(0);
      expect(stats.charCount).toBe(0);
    });

    it('should handle empty/invalid chunks gracefully', () => {
      const counter = new StreamingTokenCounter(1000, 'test');

      const result1 = counter.addChunk('');
      const result2 = counter.addChunk(null as any);
      const result3 = counter.addChunk(undefined as any);

      expect(result1.tokenCount).toBe(0);
      expect(result2.tokenCount).toBe(0);
      expect(result3.tokenCount).toBe(0);
    });

    it('should provide warning message when approaching limit', () => {
      const counter = new StreamingTokenCounter(100, 'TestProvider');

      // Add enough to approach limit
      counter.addChunk('a'.repeat(320));
      const warning = counter.getWarningMessage();

      expect(warning).toContain('TestProvider');
      expect(warning).toContain('approaching');
    });

    it('should provide warning message when exceeded', () => {
      const counter = new StreamingTokenCounter(10, 'TestProvider');

      counter.addChunk('a'.repeat(100));
      const warning = counter.getWarningMessage();

      expect(warning).toContain('TestProvider');
      expect(warning).toContain('exceeded');
    });

    it('should return null warning when under threshold', () => {
      const counter = new StreamingTokenCounter(1000, 'test');
      counter.addChunk('Short text');

      expect(counter.getWarningMessage()).toBeNull();
    });

    it('should calculate utilization correctly', () => {
      const counter = new StreamingTokenCounter(100, 'test');
      counter.addChunk('a'.repeat(200)); // ~50 tokens

      const stats = counter.getStats();
      expect(stats.utilization).toBeGreaterThan(0);
      expect(stats.utilization).toBeLessThanOrEqual(1);
    });

    it('should enforce minimum maxTokens of 1', () => {
      const counter = new StreamingTokenCounter(0, 'test');
      const stats = counter.getStats();

      expect(stats.maxTokens).toBe(1);
    });
  });

  describe('validateInputTokens', () => {
    it('should validate prompts within token limits', () => {
      const result = validateInputTokens('Hello world', 1000, 'test-provider');

      expect(result.valid).toBe(true);
      expect(result.tokenCount).toBeGreaterThan(0);
      expect(result.error).toBeUndefined();
    });

    it('should reject prompts exceeding token limits', () => {
      const longPrompt = 'a'.repeat(10000);
      const result = validateInputTokens(longPrompt, 100, 'test-provider');

      expect(result.valid).toBe(false);
      expect(result.tokenCount).toBeGreaterThan(100);
      expect(result.error).toContain('Prompt too long');
      expect(result.error).toContain('test-provider');
    });

    it('should reject empty prompts', () => {
      const result = validateInputTokens('', 1000, 'test');

      expect(result.valid).toBe(false);
      expect(result.tokenCount).toBe(0);
      expect(result.error).toContain('cannot be empty');
    });

    it('should reject null/undefined prompts', () => {
      const result1 = validateInputTokens(null as any, 1000, 'test');
      const result2 = validateInputTokens(undefined as any, 1000, 'test');

      expect(result1.valid).toBe(false);
      expect(result2.valid).toBe(false);
    });

    it('should handle zero maxInputTokens', () => {
      const result = validateInputTokens('Hello', 0, 'test');

      // When maxInputTokens is 0, validation should pass (no limit)
      expect(result.valid).toBe(true);
    });

    it('should use default provider name', () => {
      const result = validateInputTokens('a'.repeat(10000), 10);

      expect(result.error).toContain('unknown');
    });

    it('should include token counts in error messages', () => {
      const result = validateInputTokens('a'.repeat(1000), 10, 'TestProvider');

      expect(result.error).toMatch(/~\d+ tokens/);
      expect(result.error).toContain('10');
    });
  });

  describe('createTokenAwareStream', () => {
    async function* createMockStream(chunks: string[]): AsyncGenerator<string> {
      for (const chunk of chunks) {
        yield chunk;
      }
    }

    it('should yield all chunks when under limit', async () => {
      const chunks = ['Hello', ' ', 'world'];
      const stream = createTokenAwareStream(createMockStream(chunks), 1000, 'test');

      const received: string[] = [];
      for await (const chunk of stream) {
        received.push(chunk);
      }

      expect(received).toEqual(chunks);
    });

    it('should stop streaming when approaching limit', async () => {
      const chunks = Array(100).fill('a'.repeat(100)); // Many large chunks
      const stream = createTokenAwareStream(createMockStream(chunks), 100, 'test');

      const received: string[] = [];
      for await (const chunk of stream) {
        received.push(chunk);
      }

      // Should have stopped before processing all chunks
      expect(received.length).toBeLessThan(chunks.length);
    });

    it('should handle non-string chunks', async () => {
      async function* mixedStream(): AsyncGenerator<any> {
        yield 'text';
        yield { type: 'object' };
        yield 123;
        yield 'more text';
      }

      const stream = createTokenAwareStream(mixedStream(), 1000, 'test');
      const received: any[] = [];

      for await (const chunk of stream) {
        received.push(chunk);
      }

      expect(received).toHaveLength(4);
      expect(received[1]).toEqual({ type: 'object' });
      expect(received[2]).toBe(123);
    });

    it('should handle empty streams', async () => {
      async function* emptyStream(): AsyncGenerator<string> {
        // Yields nothing
      }

      const stream = createTokenAwareStream(emptyStream(), 1000, 'test');
      const received: string[] = [];

      for await (const chunk of stream) {
        received.push(chunk);
      }

      expect(received).toHaveLength(0);
    });

    it('should propagate errors from source stream', async () => {
      async function* errorStream(): AsyncGenerator<string> {
        yield 'chunk1';
        throw new Error('Stream error');
      }

      const stream = createTokenAwareStream(errorStream(), 1000, 'test');

      await expect(async () => {
        for await (const _ of stream) {
          // Consume stream
        }
      }).rejects.toThrow('Stream error');
    });

    it('should use console.warn for warnings', async () => {
      const warnSpy = jest.spyOn(console, 'warn').mockImplementation();
      const infoSpy = jest.spyOn(console, 'info').mockImplementation();

      try {
        const chunks = Array(50).fill('a'.repeat(100)); // Enough to trigger warnings
        const stream = createTokenAwareStream(createMockStream(chunks), 100, 'TestProvider');

        for await (const _ of stream) {
          // Consume stream
        }

        // Should have logged warnings
        expect(warnSpy).toHaveBeenCalled();
        expect(infoSpy).toHaveBeenCalled();
      } finally {
        warnSpy.mockRestore();
        infoSpy.mockRestore();
      }
    });

    it('should log provider name in messages', async () => {
      const infoSpy = jest.spyOn(console, 'info').mockImplementation();

      try {
        const chunks = ['test'];
        const stream = createTokenAwareStream(createMockStream(chunks), 1000, 'CustomProvider');

        for await (const _ of stream) {
          // Consume stream
        }

        expect(infoSpy).toHaveBeenCalledWith(expect.stringContaining('CustomProvider'));
      } finally {
        infoSpy.mockRestore();
      }
    });

    it('should only log warning once when approaching limit', async () => {
      const warnSpy = jest.spyOn(console, 'warn').mockImplementation();

      try {
        // Create chunks that will approach limit multiple times
        const chunks = Array(20).fill('a'.repeat(100));
        const stream = createTokenAwareStream(createMockStream(chunks), 100, 'test');

        for await (const _ of stream) {
          // Consume stream
        }

        // Filter for "approaching" warnings (not "Stopping" warnings)
        const approachingWarnings = warnSpy.mock.calls.filter(
          call => call[0]?.toString().includes('approaching')
        );

        // Should only warn once about approaching limit
        expect(approachingWarnings.length).toBeLessThanOrEqual(1);
      } finally {
        warnSpy.mockRestore();
      }
    });
  });
});