/* watermark.c
 *
 * DSSS-style audio watermark embedder/extractor prototype in C.
 * Dependencies: libsndfile, fftw3f, math
 *
 * Build:
 *   gcc -O2 -o audio_dsss watermark.c -lsndfile -lfftw3f -lm
 *
 * Modes:
 *   embed  in.wav out.wav message.txt seed
 *   extract in.wav out.bin  seed
 *   test   in.wav message.txt seed   (creates watermarked.wav + extracted.bin and compares)
 *
 * Notes:
 *  - Mono conversion: stereo averaged to mono.
 *  - Adds 4-byte little-endian length prefix to payload.
 *  - Uses r2c/c2r FFT for correct real-frame handling and proper normalization.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sndfile.h>
#include <fftw3.h>
#include <math.h>
#include <stdint.h>

#define FRAME_SIZE 2048
#define HOP 512
#define MIN_FREQ 1000.0
#define MAX_FREQ 5000.0
#define ALPHA 1.0f
#define FRAMES_PER_BIT 4
#define MAX_BITS 32768

static float hann_window(int n, int i) {
    return 0.5f * (1.0f - cosf(2.0f * M_PI * i / (n - 1)));
}

static unsigned char *load_file(const char *path, size_t *out_size) {
    FILE *f = fopen(path, "rb");
    if (!f) { perror("fopen message"); return NULL; }
    fseek(f, 0, SEEK_END);
    long sz = ftell(f);
    fseek(f, 0, SEEK_SET);
    if (sz <= 0) { fclose(f); *out_size = 0; return NULL; }
    unsigned char *buf = malloc(sz);
    if (!buf) { fclose(f); return NULL; }
    if (fread(buf, 1, sz, f) != (size_t)sz) { perror("fread"); free(buf); fclose(f); return NULL; }
    fclose(f);
    *out_size = sz;
    return buf;
}

static int write_file(const char *path, const unsigned char *buf, size_t sz) {
    FILE *f = fopen(path, "wb");
    if (!f) { perror("fopen out"); return -1; }
    if (fwrite(buf, 1, sz, f) != sz) { perror("fwrite"); fclose(f); return -1; }
    fclose(f);
    return 0;
}

/* bytes -> bits LSB-first */
static int bytes_to_bits(const unsigned char *bytes, size_t nbytes, int *bits, size_t *nbits) {
    size_t maxbits = nbytes * 8;
    if (maxbits > MAX_BITS) return -1;
    for (size_t i = 0; i < nbytes; ++i) {
        for (int b = 0; b < 8; ++b) bits[i*8 + b] = (bytes[i] >> b) & 1;
    }
    *nbits = maxbits;
    return 0;
}

/* bits -> bytes LSB-first */
static size_t bits_to_bytes(const int *bits, size_t nbits, unsigned char *out) {
    size_t nbytes = (nbits + 7) / 8;
    memset(out, 0, nbytes);
    for (size_t i = 0; i < nbits; ++i) if (bits[i]) out[i/8] |= (1u << (i%8));
    return nbytes;
}

/* xorshift32 */
static uint32_t prng_next(uint32_t *state) {
    uint32_t x = *state;
    x ^= x << 13;
    x ^= x >> 17;
    x ^= x << 5;
    *state = x;
    return x;
}

/* PN chips ±1 */
static void generate_pn(int *chips, size_t len, uint32_t *seed) {
    for (size_t i = 0; i < len; ++i) chips[i] = (prng_next(seed) & 1) ? 1 : -1;
}

/* frames count */
static int frames_count(int signal_len, int frame_size, int hop) {
    if (signal_len < frame_size) return 0;
    return 1 + (signal_len - frame_size) / hop;
}

/* mono conversion */
static float *to_mono(const float *in, sf_count_t frames, int channels) {
    float *mono = malloc(sizeof(float) * frames);
    if (!mono) return NULL;
    if (channels == 1) {
        for (sf_count_t i = 0; i < frames; ++i) mono[i] = in[i];
    } else {
        for (sf_count_t i = 0; i < frames; ++i) {
            float s = 0.0f;
            for (int c = 0; c < channels; ++c) s += in[i*channels + c];
            mono[i] = s / channels;
        }
    }
    return mono;
}

/* prepend 32-bit little-endian length to message buffer */
static unsigned char *prepend_length(const unsigned char *msg, size_t msgsz, size_t *out_sz) {
    if (!msg || !out_sz) return NULL;         // sanity check
    size_t total = msgsz + 4;                 // 4 bytes for length
    unsigned char *buf = malloc(total);
    if (!buf) return NULL;

    // store message length in little-endian format
    uint32_t len32 = (uint32_t)msgsz;
    buf[0] = (unsigned char)(len32 & 0xFF);
    buf[1] = (unsigned char)((len32 >> 8) & 0xFF);
    buf[2] = (unsigned char)((len32 >> 16) & 0xFF);
    buf[3] = (unsigned char)((len32 >> 24) & 0xFF);

    memcpy(buf + 4, msg, msgsz);             // copy message after prefix
    *out_sz = total;
    return buf;
}


/* ---------- Embed ---------- */
static int embed(const char *in_wav, const char *out_wav, const char *message_path, uint32_t seed) {
    SF_INFO sfinfo;
    memset(&sfinfo, 0, sizeof(sfinfo));
    SNDFILE *infile = sf_open(in_wav, SFM_READ, &sfinfo);
    if (!infile) { fprintf(stderr, "sf_open input failed: %s\n", in_wav); return 1; }

    sf_count_t nframes = sfinfo.frames;
    int channels = sfinfo.channels;
    float *interleaved = malloc(sizeof(float) * nframes * channels);
    if (!interleaved) { sf_close(infile); return 1; }
    if (sf_readf_float(infile, interleaved, nframes) != nframes) {
        fprintf(stderr, "sf_readf_float failed\n"); free(interleaved); sf_close(infile); return 1;
    }
    sf_close(infile);

    float *audio = to_mono(interleaved, nframes, channels);
    free(interleaved);
    if (!audio) { fprintf(stderr, "to_mono failed\n"); return 1; }

    size_t msgsz;
    unsigned char *msg = load_file(message_path, &msgsz);
    if (!msg) { free(audio); fprintf(stderr, "load_file failed\n"); return 1; }

    size_t msg_with_len_sz;
    unsigned char *msg_with_len = prepend_length(msg, msgsz, &msg_with_len_sz);
    free(msg);
    if (!msg_with_len) { free(audio); fprintf(stderr, "prepend_length failed\n"); return 1; }

    int *bits = malloc(MAX_BITS * sizeof(int));
    size_t nbits;
    if (bytes_to_bits(msg_with_len, msg_with_len_sz, bits, &nbits) != 0) {
        fprintf(stderr, "message too large\n"); free(msg_with_len); free(audio); free(bits); return 1;
    }
    free(msg_with_len);

    int frame_size = FRAME_SIZE;
    int hop = HOP;
    int nfft = frame_size;
    int nframes_stft = frames_count((int)nframes, frame_size, hop);
    if (nframes_stft <= 0) {
        fprintf(stderr, "audio too short for frame_size\n"); free(bits); free(audio); return 1;
    }

    float *window = malloc(sizeof(float) * frame_size);
    for (int i = 0; i < frame_size; ++i) window[i] = hann_window(frame_size, i);

    /* r2c: output bins = nfft/2 + 1 */
    int nbins = nfft/2 + 1;

    fftwf_plan plan_fwd = fftwf_plan_dft_r2c_1d(nfft, NULL, NULL, FFTW_ESTIMATE);
    fftwf_plan plan_inv = fftwf_plan_dft_c2r_1d(nfft, NULL, NULL, FFTW_ESTIMATE);
    /* We'll create per-frame buffers */
    float *frame_in = fftwf_alloc_real(nfft);
    fftwf_complex *spec = fftwf_alloc_complex(nbins);
    float *ifft_out = fftwf_alloc_real(nfft);

    /* Frequency bin indices for chosen band */
    double sr = sfinfo.samplerate;
    int bin_min = (int)floor(MIN_FREQ / sr * nfft);
    int bin_max = (int)ceil(MAX_FREQ / sr * nfft);
    if (bin_min < 1) bin_min = 1;
    if (bin_max > nfft/2) bin_max = nfft/2;
    int target_bins = bin_max - bin_min + 1;
    if (target_bins <= 0) { fprintf(stderr, "bad freq band selection\n"); goto cleanup_embed; }

    float *out_audio = calloc(nframes, sizeof(float));
    if (!out_audio) { fprintf(stderr, "alloc out_audio failed\n"); goto cleanup_embed; }

    int max_bits_possible = nframes_stft / FRAMES_PER_BIT;
    if ((int)nbits > max_bits_possible) {
        fprintf(stderr, "Warning: message (%zu bits) longer than capacity (%d bits). Truncating.\n", nbits, max_bits_possible);
        nbits = max_bits_possible;
    }

    uint32_t prng = seed;
    int sample_idx = 0;
    int frame_no = 0;

    /* create actual forward and inverse plans bound to buffers (since we need to execute them) */
    fftwf_destroy_plan(plan_fwd);
    fftwf_destroy_plan(plan_inv);
    plan_fwd = fftwf_plan_dft_r2c_1d(nfft, frame_in, spec, FFTW_MEASURE);
    plan_inv = fftwf_plan_dft_c2r_1d(nfft, spec, ifft_out, FFTW_MEASURE);

    while (sample_idx + frame_size <= (int)nframes) {
        /* prepare input */
        for (int i = 0; i < nfft; ++i) {
            float v = 0.0f;
            if (i < frame_size) v = audio[sample_idx + i] * window[i];
            frame_in[i] = v;
        }
        fftwf_execute(plan_fwd);

        int which_bit = frame_no / FRAMES_PER_BIT;
        if (which_bit < (int)nbits) {
            int b = bits[which_bit];
            int *chips = malloc(sizeof(int) * target_bins);
            generate_pn(chips, target_bins, &prng);
            for (int k = 0; k < target_bins; ++k) {
                int bin = bin_min + k;
                /* spec[bin] is complex: [real, imag] */
                float re = spec[bin][0];
                float im = spec[bin][1];
                float mag = sqrtf(re*re + im*im);
                float scale = ALPHA * (1.0f + mag);
                int sign = chips[k] * (b ? 1 : -1);
                /* small complex perturbation */
                spec[bin][0] += sign * scale * 0.5f;
                spec[bin][1] += sign * scale * 0.5f;
            }
            free(chips);
        }

        /* inverse */
        fftwf_execute(plan_inv);

        /* normalize and overlap-add */
        for (int i = 0; i < frame_size; ++i) {
            float v = ifft_out[i] / (float)nfft; /* normalization */
            out_audio[sample_idx + i] += v * window[i];
        }

        sample_idx += hop;
        frame_no++;
    }

    /* Fill any uncovered samples from original */
    for (int i = 0; i < (int)nframes; ++i) {
        if (!isfinite(out_audio[i])) out_audio[i] = audio[i];
    }

    /* Write WAV */
    SF_INFO outinfo;
    memset(&outinfo, 0, sizeof(outinfo));
    outinfo.samplerate = sfinfo.samplerate;
    outinfo.channels = 1;
    outinfo.format = sfinfo.format;
    SNDFILE *outfile = sf_open(out_wav, SFM_WRITE, &outinfo);
    if (!outfile) { fprintf(stderr, "sf_open out failed\n"); free(out_audio); goto cleanup_embed; }
    if (sf_writef_float(outfile, out_audio, nframes) != nframes) {
        fprintf(stderr, "sf_writef_float failed\n");
    }
    sf_close(outfile);

    printf("Embedding done. Embedded %zu bits (approx %zu bytes).\n", nbits, nbits/8);

    free(out_audio);
    cleanup_embed:
    free(bits);
    fftwf_free(frame_in);
    fftwf_free(spec);
    fftwf_free(ifft_out);
    fftwf_destroy_plan(plan_fwd);
    fftwf_destroy_plan(plan_inv);
    free(window);
    free(audio);
    return 0;
}

/* ---------- Extract ---------- */
static int extract(const char *in_wav, const char *out_msg_path, uint32_t seed) {
    SF_INFO sfinfo;
    memset(&sfinfo, 0, sizeof(sfinfo));
    SNDFILE *infile = sf_open(in_wav, SFM_READ, &sfinfo);
    if (!infile) { fprintf(stderr, "sf_open input failed: %s\n", in_wav); return 1; }

    sf_count_t nframes = sfinfo.frames;
    int channels = sfinfo.channels;
    float *interleaved = malloc(sizeof(float) * nframes * channels);
    if (!interleaved) { sf_close(infile); return 1; }
    if (sf_readf_float(infile, interleaved, nframes) != nframes) {
        fprintf(stderr, "sf_readf_float failed\n"); free(interleaved); sf_close(infile); return 1;
    }
    sf_close(infile);

    float *audio = to_mono(interleaved, nframes, channels);
    free(interleaved);
    if (!audio) { fprintf(stderr, "to_mono failed\n"); return 1; }

    int frame_size = FRAME_SIZE;
    int hop = HOP;
    int nfft = frame_size;
    int nframes_stft = frames_count((int)nframes, frame_size, hop);
    if (nframes_stft <= 0) { fprintf(stderr, "audio too short for frame size\n"); free(audio); return 1; }

    float *window = malloc(sizeof(float) * frame_size);
    for (int i = 0; i < frame_size; ++i) window[i] = hann_window(frame_size, i);

    int nbins = nfft/2 + 1;
    fftwf_plan plan_fwd = fftwf_plan_dft_r2c_1d(nfft, NULL, NULL, FFTW_ESTIMATE);
    float *frame_in = fftwf_alloc_real(nfft);
    fftwf_complex *spec = fftwf_alloc_complex(nbins);

    fftwf_destroy_plan(plan_fwd);
    plan_fwd = fftwf_plan_dft_r2c_1d(nfft, frame_in, spec, FFTW_MEASURE);

    double sr = sfinfo.samplerate;
    int bin_min = (int)floor(MIN_FREQ / sr * nfft);
    int bin_max = (int)ceil(MAX_FREQ / sr * nfft);
    if (bin_min < 1) bin_min = 1;
    if (bin_max > nfft/2) bin_max = nfft/2;
    int target_bins = bin_max - bin_min + 1;
    if (target_bins <= 0) { fprintf(stderr, "bad freq band\n"); goto cleanup_extract; }

    float *frame_corr = calloc(nframes_stft, sizeof(float));
    if (!frame_corr) { fprintf(stderr, "alloc frame_corr failed\n"); goto cleanup_extract; }

    uint32_t prng = seed;
    int sample_idx = 0;
    int frame_no = 0;
    while (sample_idx + frame_size <= (int)nframes) {
        for (int i = 0; i < nfft; ++i) frame_in[i] = (i < frame_size) ? (audio[sample_idx + i] * window[i]) : 0.0f;
        fftwf_execute(plan_fwd);

        int *chips = malloc(sizeof(int) * target_bins);
        generate_pn(chips, target_bins, &prng);

        double corr = 0.0;
        for (int k = 0; k < target_bins; ++k) {
            int bin = bin_min + k;
            float re = spec[bin][0];
            float im = spec[bin][1];
            float mag = sqrtf(re*re + im*im) + 1e-9f;
            /* correlate normalized real part with chip */
            corr += (re / mag) * (double)chips[k];
        }
        frame_corr[frame_no] = (float)corr;
        free(chips);

        sample_idx += hop;
        frame_no++;
    }

    int max_possible_bits = nframes_stft / FRAMES_PER_BIT;
    int *recovered_bits = malloc(sizeof(int) * max_possible_bits);
    if (!recovered_bits) { fprintf(stderr, "alloc recovered_bits failed\n"); goto cleanup_extract; }

    for (int b = 0; b < max_possible_bits; ++b) {
        double sum = 0.0;
        for (int f = 0; f < FRAMES_PER_BIT; ++f) {
            int fi = b * FRAMES_PER_BIT + f;
            if (fi >= nframes_stft) break;
            sum += frame_corr[fi];
        }
        recovered_bits[b] = (sum > 0.0) ? 1 : 0;
    }

    /* Convert to bytes */
    size_t nbits_out = max_possible_bits;
    unsigned char *outbytes = malloc((nbits_out/8 + 8));
    memset(outbytes, 0, (nbits_out/8 + 8));
    size_t nbytes = bits_to_bytes(recovered_bits, nbits_out, outbytes);

    /* Read first 4 bytes as length prefix (little-endian) */
    if (nbytes < 4) {
        fprintf(stderr, "extracted data too small to contain length prefix\n");
        write_file(out_msg_path, outbytes, nbytes);
        goto cleanup_extract;
    }
    uint32_t claimed_len = (uint32_t)outbytes[0] | ((uint32_t)outbytes[1] << 8) | ((uint32_t)outbytes[2] << 16) | ((uint32_t)outbytes[3] << 24);
    if (claimed_len > nbytes - 4) {
        fprintf(stderr, "warning: claimed length %u larger than extracted bytes %zu (truncated)\n", claimed_len, nbytes - 4);
        claimed_len = (uint32_t)(nbytes - 4);
    }

    if (write_file(out_msg_path, outbytes + 4, claimed_len) != 0) {
        fprintf(stderr, "write_file failed\n");
    } else {
        printf("Extraction produced %u bytes (after length prefix).\n", claimed_len);
    }

    cleanup_extract:
    if (frame_corr) free(frame_corr);
    if (recovered_bits) free(recovered_bits);
    if (outbytes) free(outbytes);
    fftwf_free(frame_in);
    fftwf_free(spec);
    fftwf_destroy_plan(plan_fwd);
    free(window);
    free(audio);
    return 0;
}

/* ---------- main ---------- */
int main(int argc, char **argv) {
    if (argc < 2) {
        fprintf(stderr, "Usage:\n  %s embed in.wav out.wav message.txt seed\n  %s extract in.wav out.bin seed\n  %s test in.wav message.txt seed\n", argv[0], argv[0], argv[0]);
        return 1;
    }
    if (strcmp(argv[1], "embed") == 0) {
        if (argc < 6) { fprintf(stderr, "embed needs in out message seed\n"); return 1; }
        uint32_t seed = (uint32_t)atoi(argv[5]);
        return embed(argv[2], argv[3], argv[4], seed);
    } else if (strcmp(argv[1], "extract") == 0) {
        if (argc < 5) { fprintf(stderr, "extract needs in out seed\n"); return 1; }
        uint32_t seed = (uint32_t)atoi(argv[4]);
        return extract(argv[2], argv[3], seed);
    } else if (strcmp(argv[1], "test") == 0) {
        if (argc < 5) { fprintf(stderr, "test needs in message seed\n"); return 1; }
        const char *inwav = argv[2];
        const char *msg = argv[3];
        uint32_t seed = (uint32_t)atoi(argv[4]);
        const char *watermarked = "watermarked.wav";
        const char *extracted = "extracted.bin";
        printf("Test mode: embedding %s into %s with seed %u -> %s\n", msg, inwav, seed, watermarked);
        int r = embed(inwav, watermarked, msg, seed);
        if (r != 0) { fprintf(stderr, "embed failed\n"); return r; }
        printf("Now extracting from %s -> %s\n", watermarked, extracted);
        r = extract(watermarked, extracted, seed);
        if (r != 0) { fprintf(stderr, "extract failed\n"); return r; }
        /* compare original message and extracted */
        size_t orig_sz;
        unsigned char *orig = load_file(msg, &orig_sz);
        size_t ext_sz;
        unsigned char *ext = load_file(extracted, &ext_sz);
        if (!orig || !ext) {
            fprintf(stderr, "couldn't read original or extracted for comparison\n");
            free(orig); free(ext);
            return 0;
        }
        if (orig_sz == ext_sz && memcmp(orig, ext, orig_sz) == 0) {
            printf("SUCCESS: extracted matches original (%zu bytes)\n", orig_sz);
        } else {
            printf("MISMATCH: original %zu bytes vs extracted %zu bytes\n", orig_sz, ext_sz);
            /* helpful hex diff of first 32 bytes */
            size_t show = (orig_sz < ext_sz ? orig_sz : ext_sz);
            if (show > 32) show = 32;
            printf("Original first %zu bytes: ", show);
            for (size_t i = 0; i < show; ++i) printf("%02X ", orig[i]);
            printf("\nExtracted first %zu bytes: ", show);
            for (size_t i = 0; i < show; ++i) printf("%02X ", ext[i]);
            printf("\n");
        }
        free(orig); free(ext);
        return 0;
    } else {
        fprintf(stderr, "unknown mode: %s\n", argv[1]);
        return 1;
    }
}