"""Tests for PitchShifter."""

import numpy as np

from tuimble.modulator import PitchShifter

SAMPLE_RATE = 48000
FRAME_SIZE = 960  # 20ms at 48kHz


def _sine_pcm(freq: float, n_samples: int = FRAME_SIZE) -> bytes:
    """Generate a single-frequency int16 PCM frame."""
    t = np.arange(n_samples) / SAMPLE_RATE
    samples = (np.sin(2 * np.pi * freq * t) * 16000).astype(np.int16)
    return samples.tobytes()


def _dominant_freq(pcm: bytes) -> float:
    """Return the dominant frequency in an int16 PCM buffer."""
    samples = np.frombuffer(pcm, dtype=np.int16).astype(np.float32)
    fft = np.abs(np.fft.rfft(samples))
    freqs = np.fft.rfftfreq(len(samples), 1.0 / SAMPLE_RATE)
    return freqs[np.argmax(fft)]


def test_zero_semitones_passthrough():
    ps = PitchShifter(SAMPLE_RATE)
    pcm = _sine_pcm(440.0)
    assert ps.process(pcm) == pcm


def test_pitch_shift_changes_output():
    ps = PitchShifter(SAMPLE_RATE)
    pcm = _sine_pcm(440.0)
    ps.semitones = 3.0
    result = ps.process(pcm)
    assert result != pcm


def test_output_length_preserved():
    ps = PitchShifter(SAMPLE_RATE)
    pcm = _sine_pcm(440.0)
    ps.semitones = 5.0
    result = ps.process(pcm)
    assert len(result) == len(pcm)


def test_semitones_clamping():
    ps = PitchShifter()
    ps.semitones = 20.0
    assert ps.semitones == 12.0
    ps.semitones = -20.0
    assert ps.semitones == -12.0
    ps.semitones = 5.0
    assert ps.semitones == 5.0


def test_empty_input():
    ps = PitchShifter()
    ps.semitones = 3.0
    assert ps.process(b"") == b""
    assert ps.process(b"\x00") == b"\x00"


def test_pitch_up_frequency_increases():
    ps = PitchShifter(SAMPLE_RATE)
    pcm = _sine_pcm(440.0)
    ps.semitones = 4.0
    result = ps.process(pcm)
    assert _dominant_freq(result) > _dominant_freq(pcm)


def test_pitch_down_frequency_decreases():
    ps = PitchShifter(SAMPLE_RATE)
    pcm = _sine_pcm(440.0)
    ps.semitones = -4.0
    result = ps.process(pcm)
    assert _dominant_freq(result) < _dominant_freq(pcm)