#include "wav_io.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

#define TEST_FILE "test_output.wav"
#define SAMPLE_RATE 44100
#define EPSILON 1e-6f

static int tests_passed = 0;
static int tests_failed = 0;

#define TEST(name, expr) do { \
    if (expr) { \
        tests_passed++; \
        printf("  PASS: %s\n", name); \
    } else { \
        tests_failed++; \
        printf("  FAIL: %s\n", name); \
    } \
} while(0)

static int float_close(float a, float b) {
    return fabsf(a - b) < EPSILON;
}

// Test 1: Save and load a simple float buffer (round-trip)
static void test_round_trip_float(void) {
    printf("Test: Round-trip float WAV\n");
    
    size_t num_samples = 1000;
    float *original = (float *)malloc(num_samples * sizeof(float));
    for (size_t i = 0; i < num_samples; i++) {
        original[i] = sinf(2.0f * 3.14159f * 440.0f * i / SAMPLE_RATE);
    }
    
    int ret = save_wav_float(TEST_FILE, original, num_samples, SAMPLE_RATE);
    TEST("save_wav_float returns 0", ret == 0);
    
    float *loaded = NULL;
    size_t loaded_samples = 0;
    unsigned int loaded_rate = 0;
    ret = load_wav_float(TEST_FILE, &loaded, &loaded_samples, &loaded_rate);
    TEST("load_wav_float returns 0", ret == 0);
    TEST("loaded sample count matches", loaded_samples == num_samples);
    TEST("loaded sample rate matches", loaded_rate == SAMPLE_RATE);
    
    int samples_match = 1;
    for (size_t i = 0; i < num_samples; i++) {
        if (!float_close(original[i], loaded[i])) {
            samples_match = 0;
            break;
        }
    }
    TEST("all samples match after round-trip", samples_match);
    
    free(original);
    free(loaded);
    remove(TEST_FILE);
}

// Test 2: Save with NULL filename
static void test_save_null_filename(void) {
    printf("Test: save_wav_float with NULL filename\n");
    float buf[10] = {0};
    int ret = save_wav_float(NULL, buf, 10, SAMPLE_RATE);
    TEST("returns -1", ret == -1);
}

// Test 3: Save with NULL buffer
static void test_save_null_buffer(void) {
    printf("Test: save_wav_float with NULL buffer\n");
    int ret = save_wav_float("test.wav", NULL, 10, SAMPLE_RATE);
    TEST("returns -1", ret == -1);
}

// Test 4: Save with zero samples
static void test_save_zero_samples(void) {
    printf("Test: save_wav_float with zero samples\n");
    float buf[10] = {0};
    int ret = save_wav_float("test.wav", buf, 0, SAMPLE_RATE);
    TEST("returns -1", ret == -1);
}

// Test 5: Load with NULL filename
static void test_load_null_filename(void) {
    printf("Test: load_wav_float with NULL filename\n");
    float *buf = NULL;
    size_t samples = 0;
    unsigned int rate = 0;
    int ret = load_wav_float(NULL, &buf, &samples, &rate);
    TEST("returns -1", ret == -1);
}

// Test 6: Load with NULL buffer pointer
static void test_load_null_buffer(void) {
    printf("Test: load_wav_float with NULL buffer pointer\n");
    size_t samples = 0;
    unsigned int rate = 0;
    int ret = load_wav_float("test.wav", NULL, &samples, &rate);
    TEST("returns -1", ret == -1);
}

// Test 7: Load with NULL num_samples
static void test_load_null_num_samples(void) {
    printf("Test: load_wav_float with NULL num_samples\n");
    float *buf = NULL;
    unsigned int rate = 0;
    int ret = load_wav_float("test.wav", &buf, NULL, &rate);
    TEST("returns -1", ret == -1);
}

// Test 8: Load with NULL sample_rate
static void test_load_null_sample_rate(void) {
    printf("Test: load_wav_float with NULL sample_rate\n");
    float *buf = NULL;
    size_t samples = 0;
    int ret = load_wav_float("test.wav", &buf, &samples, NULL);
    TEST("returns -1", ret == -1);
}

// Test 9: Load non-existent file
static void test_load_nonexistent(void) {
    printf("Test: load_wav_float with non-existent file\n");
    float *buf = NULL;
    size_t samples = 0;
    unsigned int rate = 0;
    int ret = load_wav_float("nonexistent.wav", &buf, &samples, &rate);
    TEST("returns -1", ret == -1);
}

// Test 10: Save and load a large buffer
static void test_large_buffer(void) {
    printf("Test: Large buffer round-trip\n");
    
    size_t num_samples = 100000;
    float *original = (float *)malloc(num_samples * sizeof(float));
    for (size_t i = 0; i < num_samples; i++) {
        original[i] = (float)i / num_samples;
    }
    
    int ret = save_wav_float(TEST_FILE, original, num_samples, 48000);
    TEST("save returns 0", ret == 0);
    
    float *loaded = NULL;
    size_t loaded_samples = 0;
    unsigned int loaded_rate = 0;
    ret = load_wav_float(TEST_FILE, &loaded, &loaded_samples, &loaded_rate);
    TEST("load returns 0", ret == 0);
    TEST("sample count matches", loaded_samples == num_samples);
    TEST("sample rate matches", loaded_rate == 48000);
    
    int samples_match = 1;
    for (size_t i = 0; i < num_samples; i++) {
        if (!float_close(original[i], loaded[i])) {
            samples_match = 0;
            break;
        }
    }
    TEST("all samples match", samples_match);
    
    free(original);
    free(loaded);
    remove(TEST_FILE);
}

// Test 11: Save and load with different sample rates
static void test_different_sample_rates(void) {
    printf("Test: Different sample rates\n");
    
    float buf[] = {0.0f, 0.5f, 1.0f, -0.5f, -1.0f};
    size_t num_samples = 5;
    unsigned int rates[] = {8000, 22050, 44100, 48000, 96000};
    
    for (int i = 0; i < 5; i++) {
        char filename[64];
        snprintf(filename, sizeof(filename), "test_rate_%u.wav", rates[i]);
        
        int ret = save_wav_float(filename, buf, num_samples, rates[i]);
        TEST("save returns 0", ret == 0);
        
        float *loaded = NULL;
        size_t loaded_samples = 0;
        unsigned int loaded_rate = 0;
        ret = load_wav_float(filename, &loaded, &loaded_samples, &loaded_rate);
        TEST("load returns 0", ret == 0);
        TEST("sample rate matches", loaded_rate == rates[i]);
        TEST("sample count matches", loaded_samples == num_samples);
        
        int samples_match = 1;
        for (size_t j = 0; j < num_samples; j++) {
            if (!float_close(buf[j], loaded[j])) {
                samples_match = 0;
                break;
            }
        }
        TEST("samples match", samples_match);
        
        free(loaded);
        remove(filename);
    }
}

// Test 12: Save and load with negative values
static void test_negative_values(void) {
    printf("Test: Negative values\n");
    
    float buf[] = {-1.0f, -0.75f, -0.5f, -0.25f, 0.0f, 0.25f, 0.5f, 0.75f, 1.0f};
    size_t num_samples = 9;
    
    int ret = save_wav_float(TEST_FILE, buf, num_samples, SAMPLE_RATE);
    TEST("save returns 0", ret == 0);
    
    float *loaded = NULL;
    size_t loaded_samples = 0;
    unsigned int loaded_rate = 0;
    ret = load_wav_float(TEST_FILE, &loaded, &loaded_samples, &loaded_rate);
    TEST("load returns 0", ret == 0);
    TEST("sample count matches", loaded_samples == num_samples);
    
    int samples_match = 1;
    for (size_t i = 0; i < num_samples; i++) {
        if (!float_close(buf[i], loaded[i])) {
            samples_match = 0;
            break;
        }
    }
    TEST("all samples match", samples_match);
    
    free(loaded);
    remove(TEST_FILE);
}

// Test 13: Save and load with silence (all zeros)
static void test_silence(void) {
    printf("Test: Silence (all zeros)\n");
    
    size_t num_samples = 100;
    float *buf = (float *)calloc(num_samples, sizeof(float));
    
    int ret = save_wav_float(TEST_FILE, buf, num_samples, SAMPLE_RATE);
    TEST("save returns 0", ret == 0);
    
    float *loaded = NULL;
    size_t loaded_samples = 0;
    unsigned int loaded_rate = 0;
    ret = load_wav_float(TEST_FILE, &loaded, &loaded_samples, &loaded_rate);
    TEST("load returns 0", ret == 0);
    TEST("sample count matches", loaded_samples == num_samples);
    
    int all_zeros = 1;
    for (size_t i = 0; i < num_samples; i++) {
        if (loaded[i] != 0.0f) {
            all_zeros = 0;
            break;
        }
    }
    TEST("all samples are zero", all_zeros);
    
    free(buf);
    free(loaded);
    remove(TEST_FILE);
}

// Test 14: Load a manually crafted 16-bit PCM WAV file
static void test_load_16bit_pcm(void) {
    printf("Test: Load 16-bit PCM WAV\n");
    
    // Create a minimal 16-bit PCM WAV file manually
    FILE *f = fopen(TEST_FILE, "wb");
    if (!f) {
        printf("  SKIP: Could not create test file\n");
        return;
    }
    
    int16_t pcm_data[] = {0, 16384, 32767, -16384, -32768};
    size_t num_samples = 5;
    uint32_t data_size = num_samples * sizeof(int16_t);
    uint32_t chunk_size = 36 + data_size;
    
    fwrite("RIFF", 1, 4, f);
    unsigned char buf[4];
    buf[0] = chunk_size & 0xFF;
    buf[1] = (chunk_size >> 8) & 0xFF;
    buf[2] = (chunk_size >> 16) & 0xFF;
    buf[3] = (chunk_size >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    fwrite("WAVE", 1, 4, f);
    
    fwrite("fmt ", 1, 4, f);
    uint32_t fmt_size = 16;
    buf[0] = fmt_size & 0xFF;
    buf[1] = (fmt_size >> 8) & 0xFF;
    buf[2] = (fmt_size >> 16) & 0xFF;
    buf[3] = (fmt_size >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    
    uint16_t audio_format = 1; // PCM
    buf[0] = audio_format & 0xFF;
    buf[1] = (audio_format >> 8) & 0xFF;
    fwrite(buf, 1, 2, f);
    
    uint16_t channels = 1;
    buf[0] = channels & 0xFF;
    buf[1] = (channels >> 8) & 0xFF;
    fwrite(buf, 1, 2, f);
    
    uint32_t sample_rate = 44100;
    buf[0] = sample_rate & 0xFF;
    buf[1] = (sample_rate >> 8) & 0xFF;
    buf[2] = (sample_rate >> 16) & 0xFF;
    buf[3] = (sample_rate >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    
    uint32_t byte_rate = sample_rate * sizeof(int16_t);
    buf[0] = byte_rate & 0xFF;
    buf[1] = (byte_rate >> 8) & 0xFF;
    buf[2] = (byte_rate >> 16) & 0xFF;
    buf[3] = (byte_rate >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    
    uint16_t block_align = sizeof(int16_t);
    buf[0] = block_align & 0xFF;
    buf[1] = (block_align >> 8) & 0xFF;
    fwrite(buf, 1, 2, f);
    
    uint16_t bits_per_sample = 16;
    buf[0] = bits_per_sample & 0xFF;
    buf[1] = (bits_per_sample >> 8) & 0xFF;
    fwrite(buf, 1, 2, f);
    
    fwrite("data", 1, 4, f);
    buf[0] = data_size & 0xFF;
    buf[1] = (data_size >> 8) & 0xFF;
    buf[2] = (data_size >> 16) & 0xFF;
    buf[3] = (data_size >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    
    fwrite(pcm_data, sizeof(int16_t), num_samples, f);
    fclose(f);
    
    float *loaded = NULL;
    size_t loaded_samples = 0;
    unsigned int loaded_rate = 0;
    int ret = load_wav_float(TEST_FILE, &loaded, &loaded_samples, &loaded_rate);
    TEST("load returns 0", ret == 0);
    TEST("sample count matches", loaded_samples == num_samples);
    TEST("sample rate matches", loaded_rate == 44100);
    
    // Expected float values: 0/32768, 16384/32768, 32767/32768, -16384/32768, -32768/32768
    float expected[] = {0.0f, 0.5f, 0.999969f, -0.5f, -1.0f};
    int samples_match = 1;
    for (size_t i = 0; i < num_samples; i++) {
        if (!float_close(expected[i], loaded[i])) {
            samples_match = 0;
            break;
        }
    }
    TEST("samples match expected values", samples_match);
    
    free(loaded);
    remove(TEST_FILE);
}

// Test 15: Load a stereo WAV file (should mix to mono)
static void test_load_stereo(void) {
    printf("Test: Load stereo WAV (mixes to mono)\n");
    
    // Create a stereo 16-bit PCM WAV file
    FILE *f = fopen(TEST_FILE, "wb");
    if (!f) {
        printf("  SKIP: Could not create test file\n");
        return;
    }
    
    int16_t pcm_data[] = {10000, 20000, -10000, -20000, 0, 0}; // L,R pairs
    size_t num_frames = 3;
    uint32_t data_size = num_frames * 2 * sizeof(int16_t);
    uint32_t chunk_size = 36 + data_size;
    
    fwrite("RIFF", 1, 4, f);
    unsigned char buf[4];
    buf[0] = chunk_size & 0xFF;
    buf[1] = (chunk_size >> 8) & 0xFF;
    buf[2] = (chunk_size >> 16) & 0xFF;
    buf[3] = (chunk_size >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    fwrite("WAVE", 1, 4, f);
    
    fwrite("fmt ", 1, 4, f);
    uint32_t fmt_size = 16;
    buf[0] = fmt_size & 0xFF;
    buf[1] = (fmt_size >> 8) & 0xFF;
    buf[2] = (fmt_size >> 16) & 0xFF;
    buf[3] = (fmt_size >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    
    uint16_t audio_format = 1;
    buf[0] = audio_format & 0xFF;
    buf[1] = (audio_format >> 8) & 0xFF;
    fwrite(buf, 1, 2, f);
    
    uint16_t channels = 2;
    buf[0] = channels & 0xFF;
    buf[1] = (channels >> 8) & 0xFF;
    fwrite(buf, 1, 2, f);
    
    uint32_t sample_rate = 44100;
    buf[0] = sample_rate & 0xFF;
    buf[1] = (sample_rate >> 8) & 0xFF;
    buf[2] = (sample_rate >> 16) & 0xFF;
    buf[3] = (sample_rate >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    
    uint32_t byte_rate = sample_rate * 2 * sizeof(int16_t);
    buf[0] = byte_rate & 0xFF;
    buf[1] = (byte_rate >> 8) & 0xFF;
    buf[2] = (byte_rate >> 16) & 0xFF;
    buf[3] = (byte_rate >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    
    uint16_t block_align = 2 * sizeof(int16_t);
    buf[0] = block_align & 0xFF;
    buf[1] = (block_align >> 8) & 0xFF;
    fwrite(buf, 1, 2, f);
    
    uint16_t bits_per_sample = 16;
    buf[0] = bits_per_sample & 0xFF;
    buf[1] = (bits_per_sample >> 8) & 0xFF;
    fwrite(buf, 1, 2, f);
    
    fwrite("data", 1, 4, f);
    buf[0] = data_size & 0xFF;
    buf[1] = (data_size >> 8) & 0xFF;
    buf[2] = (data_size >> 16) & 0xFF;
    buf[3] = (data_size >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    
    fwrite(pcm_data, sizeof(int16_t), num_frames * 2, f);
    fclose(f);
    
    float *loaded = NULL;
    size_t loaded_samples = 0;
    unsigned int loaded_rate = 0;
    int ret = load_wav_float(TEST_FILE, &loaded, &loaded_samples, &loaded_rate);
    TEST("load returns 0", ret == 0);
    TEST("sample count matches frames", loaded_samples == num_frames);
    
    // Expected: (10000 + 20000)/2 / 32768 = 0.457764, (-10000 + -20000)/2 / 32768 = -0.457764, (0+0)/2 = 0
    float expected[] = {0.457764f, -0.457764f, 0.0f};
    int samples_match = 1;
    for (size_t i = 0; i < num_frames; i++) {
        if (!float_close(expected[i], loaded[i])) {
            samples_match = 0;
            break;
        }
    }
    TEST("stereo mixed to mono correctly", samples_match);
    
    free(loaded);
    remove(TEST_FILE);
}

// Test 16: Load a file with extra chunks (like LIST)
static void test_load_with_extra_chunks(void) {
    printf("Test: Load WAV with extra chunks\n");
    
    // Create a WAV with a LIST chunk before data
    FILE *f = fopen(TEST_FILE, "wb");
    if (!f) {
        printf("  SKIP: Could not create test file\n");
        return;
    }
    
    float samples[] = {0.5f, -0.5f, 0.25f, -0.25f};
    size_t num_samples = 4;
    uint32_t data_size = num_samples * sizeof(float);
    uint32_t list_size = 8; // "INFO" + padding
    uint32_t chunk_size = 36 + list_size + data_size;
    
    fwrite("RIFF", 1, 4, f);
    unsigned char buf[4];
    buf[0] = chunk_size & 0xFF;
    buf[1] = (chunk_size >> 8) & 0xFF;
    buf[2] = (chunk_size >> 16) & 0xFF;
    buf[3] = (chunk_size >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    fwrite("WAVE", 1, 4, f);
    
    // fmt chunk
    fwrite("fmt ", 1, 4, f);
    uint32_t fmt_size = 16;
    buf[0] = fmt_size & 0xFF;
    buf[1] = (fmt_size >> 8) & 0xFF;
    buf[2] = (fmt_size >> 16) & 0xFF;
    buf[3] = (fmt_size >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    
    uint16_t audio_format = 3;
    buf[0] = audio_format & 0xFF;
    buf[1] = (audio_format >> 8) & 0xFF;
    fwrite(buf, 1, 2, f);
    
    uint16_t channels = 1;
    buf[0] = channels & 0xFF;
    buf[1] = (channels >> 8) & 0xFF;
    fwrite(buf, 1, 2, f);
    
    uint32_t sample_rate = 44100;
    buf[0] = sample_rate & 0xFF;
    buf[1] = (sample_rate >> 8) & 0xFF;
    buf[2] = (sample_rate >> 16) & 0xFF;
    buf[3] = (sample_rate >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    
    uint32_t byte_rate = sample_rate * sizeof(float);
    buf[0] = byte_rate & 0xFF;
    buf[1] = (byte_rate >> 8) & 0xFF;
    buf[2] = (byte_rate >> 16) & 0xFF;
    buf[3] = (byte_rate >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    
    uint16_t block_align = sizeof(float);
    buf[0] = block_align & 0xFF;
    buf[1] = (block_align >> 8) & 0xFF;
    fwrite(buf, 1, 2, f);
    
    uint16_t bits_per_sample = 32;
    buf[0] = bits_per_sample & 0xFF;
    buf[1] = (bits_per_sample >> 8) & 0xFF;
    fwrite(buf, 1, 2, f);
    
    // LIST chunk
    fwrite("LIST", 1, 4, f);
    buf[0] = list_size & 0xFF;
    buf[1] = (list_size >> 8) & 0xFF;
    buf[2] = (list_size >> 16) & 0xFF;
    buf[3] = (list_size >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    fwrite("INFO", 1, 4, f);
    fwrite("\0\0\0\0", 1, 4, f); // padding
    
    // data chunk
    fwrite("data", 1, 4, f);
    buf[0] = data_size & 0xFF;
    buf[1] = (data_size >> 8) & 0xFF;
    buf[2] = (data_size >> 16) & 0xFF;
    buf[3] = (data_size >> 24) & 0xFF;
    fwrite(buf, 1, 4, f);
    
    fwrite(samples, sizeof(float), num_samples, f);
    fclose(f);
    
    float *loaded = NULL;
    size_t loaded_samples = 0;
    unsigned int loaded_rate = 0;
    int ret = load_wav_float(TEST_FILE, &loaded, &loaded_samples, &loaded_rate);
    TEST("load returns 0", ret == 0);
    TEST("sample count matches", loaded_samples == num_samples);
    
    int samples_match = 1;
    for (size_t i = 0; i < num_samples; i++) {
        if (!float_close(samples[i], loaded[i])) {
            samples_match = 0;
            break;
        }
    }
    TEST("samples match despite extra chunks", samples_match);
    
    free(loaded);
    remove(TEST_FILE);
}

int main(void) {
    printf("=== WAV I/O Tests ===\n\n");
    
    test_round_trip_float();
    printf("\n");
    test_save_null_filename();
    test_save_null_buffer();
    test_save_zero_samples();
    printf("\n");
    test_load_null_filename();
    test_load_null_buffer();
    test_load_null_num_samples();
    test_load_null_sample_rate();
    test_load_nonexistent();
    printf("\n");
    test_large_buffer();
    printf("\n");
    test_different_sample_rates();
    printf("\n");
    test_negative_values();
    printf("\n");
    test_silence();
    printf("\n");
    test_load_16bit_pcm();
    printf("\n");
    test_load_stereo();
    printf("\n");
    test_load_with_extra_chunks();
    
    printf("\n=== Results ===\n");
    printf("Passed: %d\n", tests_passed);
    printf("Failed: %d\n", tests_failed);
    
    return tests_failed > 0 ? 1 : 0;
}