import { describe, it, expect, beforeEach, vi } from 'vitest';
import { CREPEModel, detectPitchCREPE, CREPEUtils } from './crepe';

// Mock TFJSModelManager
vi.mock('./tfjs', () => ({
  TFJSModelManager: vi.fn().mockImplementation(() => ({
    initialize: vi.fn().mockResolvedValue(undefined),
    loadModel: vi.fn().mockResolvedValue(undefined),
    predictArray: vi.fn().mockResolvedValue(new Float32Array(360).fill(0.1)),
    dispose: vi.fn(),
  })),
  TensorUtils: {},
}));

describe('CREPEModel', () => {
  let model: CREPEModel;

  beforeEach(() => {
    model = new CREPEModel();
  });

  describe('constructor', () => {
    it('should create model with default config', () => {
      const config = model.getConfig();
      
      expect(config.modelUrl).toBe('/models/crepe/model.json');
      expect(config.modelSize).toBe('small');
      expect(config.viterbi).toBe(false);
      expect(config.stepSize).toBe(10);
    });

    it('should accept custom config', () => {
      const customModel = new CREPEModel({
        modelUrl: '/custom/model.json',
        modelSize: 'large',
        viterbi: true,
        stepSize: 20,
      });

      const config = customModel.getConfig();
      expect(config.modelUrl).toBe('/custom/model.json');
      expect(config.modelSize).toBe('large');
      expect(config.viterbi).toBe(true);
      expect(config.stepSize).toBe(20);
    });
  });

  describe('initialize', () => {
    it('should initialize model', async () => {
      await model.initialize();
      expect(model.isInitialized()).toBe(true);
    });

    it('should not reinitialize', async () => {
      await model.initialize();
      await model.initialize(); // Should be ignored
      
      expect(model.isInitialized()).toBe(true);
    });
  });

  describe('predict', () => {
    it('should predict pitch from audio', async () => {
      await model.initialize();
      
      // Create a test audio frame (1024 samples at 16kHz)
      const audioFrame = new Float32Array(1024);
      for (let i = 0; i < audioFrame.length; i++) {
        audioFrame[i] = Math.sin((2 * Math.PI * 440 * i) / 16000);
      }

      const result = await model.predict(audioFrame, 16000);
      
      expect(result).toHaveProperty('frequency');
      expect(result).toHaveProperty('confidence');
      expect(result.frequency).toBeGreaterThan(0);
      expect(result.confidence).toBeGreaterThanOrEqual(0);
      expect(result.confidence).toBeLessThanOrEqual(1);
    });

    it('should throw if not initialized', async () => {
      const audioFrame = new Float32Array(1024);
      
      await expect(model.predict(audioFrame)).rejects.toThrow('not initialized');
    });

    it('should handle different sample rates', async () => {
      await model.initialize();
      
      // Test with 44.1kHz audio
      const audioFrame = new Float32Array(2205); // 50ms at 44.1kHz
      for (let i = 0; i < audioFrame.length; i++) {
        audioFrame[i] = Math.sin((2 * Math.PI * 440 * i) / 44100);
      }

      const result = await model.predict(audioFrame, 44100);
      
      expect(result.frequency).toBeGreaterThan(0);
    });

    it('should handle short audio frames', async () => {
      await model.initialize();
      
      const shortFrame = new Float32Array(512); // Shorter than 1024
      const result = await model.predict(shortFrame, 16000);
      
      expect(result).toBeDefined();
    });

    it('should handle long audio frames', async () => {
      await model.initialize();
      
      const longFrame = new Float32Array(2048); // Longer than 1024
      const result = await model.predict(longFrame, 16000);
      
      expect(result).toBeDefined();
    });
  });

  describe('predictBatch', () => {
    it('should predict multiple frames', async () => {
      await model.initialize();
      
      const frames = [
        new Float32Array(1024),
        new Float32Array(1024),
        new Float32Array(1024),
      ];

      const results = await model.predictBatch(frames, 16000);
      
      expect(results).toHaveLength(3);
      results.forEach(result => {
        expect(result).toHaveProperty('frequency');
        expect(result).toHaveProperty('confidence');
      });
    });
  });

  describe('viterbiSmooth', () => {
    it('should return unchanged if viterbi disabled', () => {
      const predictions = [
        { frequency: 440, confidence: 0.9 },
        { frequency: 450, confidence: 0.8 },
        { frequency: 460, confidence: 0.85 },
      ];

      const smoothed = model.viterbiSmooth(predictions);
      expect(smoothed).toEqual(predictions);
    });

    it('should smooth predictions when enabled', () => {
      const viterbiModel = new CREPEModel({ viterbi: true });
      
      const predictions = [
        { frequency: 440, confidence: 0.9 },
        { frequency: 880, confidence: 0.8 }, // Large jump
        { frequency: 450, confidence: 0.85 },
      ];

      const smoothed = viterbiModel.viterbiSmooth(predictions);
      
      // Second prediction should have reduced confidence due to large jump
      expect(smoothed[1].confidence).toBeLessThan(predictions[1].confidence);
    });

    it('should handle single prediction', () => {
      const viterbiModel = new CREPEModel({ viterbi: true });
      const predictions = [{ frequency: 440, confidence: 0.9 }];

      const smoothed = viterbiModel.viterbiSmooth(predictions);
      expect(smoothed).toEqual(predictions);
    });
  });

  describe('frequencyToBin', () => {
    it('should convert frequency to bin', () => {
      const bin = model.frequencyToBin(440);
      
      expect(bin).toBeGreaterThanOrEqual(0);
      expect(bin).toBeLessThan(360);
    });

    it('should handle different frequencies', () => {
      const lowBin = model.frequencyToBin(100);
      const highBin = model.frequencyToBin(1000);
      
      expect(highBin).toBeGreaterThan(lowBin);
    });
  });

  describe('dispose', () => {
    it('should dispose model', async () => {
      await model.initialize();
      
      model.dispose();
      
      expect(model.isInitialized()).toBe(false);
    });
  });
});

describe('detectPitchCREPE utility', () => {
  it('should detect pitch from audio', async () => {
    // Mock a high-confidence prediction
    const mockPredict = vi.fn().mockResolvedValue({
      frequency: 440,
      confidence: 0.9,
    });

    vi.doMock('./tfjs', () => ({
      TFJSModelManager: vi.fn().mockImplementation(() => ({
        initialize: vi.fn().mockResolvedValue(undefined),
        loadModel: vi.fn().mockResolvedValue(undefined),
        predictArray: vi.fn().mockResolvedValue(new Float32Array(360).fill(0.1)),
        dispose: vi.fn(),
      })),
    }));

    const audioFrame = new Float32Array(1024);
    const frequency = await detectPitchCREPE(audioFrame, 16000);
    
    // Should return a frequency or null
    expect(frequency === null || typeof frequency === 'number').toBe(true);
  });
});

describe('CREPEUtils', () => {
  describe('constants', () => {
    it('should have correct constants', () => {
      expect(CREPEUtils.MIN_FREQ).toBeCloseTo(32.70, 2);
      expect(CREPEUtils.SAMPLE_RATE).toBe(16000);
      expect(CREPEUtils.FRAME_SIZE).toBe(1024);
      expect(CREPEUtils.BINS).toBe(360);
    });
  });

  describe('isInRange', () => {
    it('should check if frequency is in CREPE range', () => {
      expect(CREPEUtils.isInRange(440)).toBe(true);
      expect(CREPEUtils.isInRange(20)).toBe(false);
      expect(CREPEUtils.isInRange(3000)).toBe(false);
    });
  });

  describe('clamp', () => {
    it('should clamp frequency to CREPE range', () => {
      expect(CREPEUtils.clamp(20)).toBeCloseTo(CREPEUtils.MIN_FREQ);
      expect(CREPEUtils.clamp(3000)).toBeCloseTo(CREPEUtils.MAX_FREQ);
      expect(CREPEUtils.clamp(440)).toBe(440);
    });
  });

  describe('MIDI conversion', () => {
    it('should convert MIDI to bin', () => {
      const bin = CREPEUtils.midiToBin(69); // A4 = 440 Hz
      expect(bin).toBeGreaterThanOrEqual(0);
      expect(bin).toBeLessThan(360);
    });

    it('should convert bin to MIDI', () => {
      const bin = CREPEUtils.midiToBin(69);
      const midi = CREPEUtils.binToMidi(bin);
      expect(midi).toBeCloseTo(69, 0);
    });

    it('should handle round-trip conversion', () => {
      const originalMidi = 60; // C4
      const bin = CREPEUtils.midiToBin(originalMidi);
      const convertedMidi = CREPEUtils.binToMidi(bin);
      expect(convertedMidi).toBeCloseTo(originalMidi, 0);
    });
  });
});