#include "acoustic_analyzer/io/wav_file_source.h"
#include <cstring>
#include <algorithm>
#include <cmath>
#include <cstdint>

namespace acoustic {

struct WavHeader {
    char riff[4];
    uint32_t file_size;
    char wave[4];
    char fmt[4];
    uint32_t fmt_size;
    uint16_t audio_format;
    uint16_t num_channels;
    uint32_t sample_rate;
    uint32_t byte_rate;
    uint16_t block_align;
    uint16_t bits_per_sample;
};

WavFileSource::WavFileSource(const std::string& path, int target_sample_rate)
    : path_(path), target_sample_rate_(target_sample_rate), fp_(nullptr) {}

WavFileSource::~WavFileSource() {
    close();
}

void WavFileSource::close() {
    if (fp_) {
        fclose(fp_);
        fp_ = nullptr;
    }
}

bool WavFileSource::open() {
    close();
    fp_ = fopen(path_.c_str(), "rb");
    if (!fp_) return false;
    if (!parse_wav_header()) {
        close();
        return false;
    }
    return true;
}

bool WavFileSource::parse_wav_header() {
    WavHeader hdr;
    if (fread(&hdr, sizeof(hdr), 1, fp_) != 1) return false;
    if (std::memcmp(hdr.riff, "RIFF", 4) != 0 || std::memcmp(hdr.wave, "WAVE", 4) != 0) return false;
    if (hdr.audio_format != 1 && hdr.audio_format != 3) return false; // PCM or IEEE float

    num_channels_ = hdr.num_channels;
    file_sample_rate_ = hdr.sample_rate;
    bits_per_sample_ = hdr.bits_per_sample;

    // Find data chunk
    char chunk_id[4];
    uint32_t chunk_size;
    while (fread(chunk_id, 4, 1, fp_) == 1) {
        if (fread(&chunk_size, 4, 1, fp_) != 1) return false;
        if (std::memcmp(chunk_id, "data", 4) == 0) {
            data_start_ = ftell(fp_);
            int bytes_per_sample = bits_per_sample_ / 8;
            if (bytes_per_sample == 0) bytes_per_sample = 1;
            total_samples_ = chunk_size / (num_channels_ * bytes_per_sample);
            read_pos_ = 0;
            return true;
        }
        fseek(fp_, chunk_size, SEEK_CUR);
    }
    return false;
}

namespace {

float convert_sample(const uint8_t* raw, int bytes, bool is_float) {
    if (is_float) {
        if (bytes == 4) {
            float val;
            std::memcpy(&val, raw, 4);
            return val;
        }
        return 0.0f;
    }
    // Integer PCM
    if (bytes == 1) {
        int32_t val = static_cast<int32_t>(raw[0]) - 128;
        return val / 128.0f;
    } else if (bytes == 2) {
        int16_t val;
        std::memcpy(&val, raw, 2);
        return val / 32768.0f;
    } else if (bytes == 3) {
        int32_t val = (static_cast<int32_t>(raw[2]) << 16)
                    | (static_cast<int32_t>(raw[1]) << 8)
                    | static_cast<int32_t>(raw[0]);
        // Sign-extend 24-bit to 32-bit
        if (val & 0x800000) val |= 0xFF000000;
        return val / 8388608.0f;
    } else if (bytes >= 4) {
        int32_t val;
        std::memcpy(&val, raw, 4);
        return std::max(-1.0f, std::min(1.0f, val / 2147483648.0f));
    }
    return 0.0f;
}

} // anonymous namespace

void WavFileSource::resample_if_needed(std::vector<float>& mono, int src_rate, int dst_rate) {
    if (src_rate == dst_rate || dst_rate <= 0) return;
    double ratio = static_cast<double>(dst_rate) / src_rate;
    std::size_t new_len = static_cast<std::size_t>(mono.size() * ratio);
    std::vector<float> resampled(new_len);
    for (std::size_t i = 0; i < new_len; ++i) {
        double src_idx = i / ratio;
        std::size_t idx0 = static_cast<std::size_t>(src_idx);
        std::size_t idx1 = std::min(idx0 + 1, mono.size() - 1);
        double frac = src_idx - idx0;
        resampled[i] = static_cast<float>(mono[idx0] * (1.0 - frac) + mono[idx1] * frac);
    }
    mono.swap(resampled);
}

std::size_t WavFileSource::read(std::vector<std::vector<float>>& out, std::size_t max_samples) {
    if (!fp_) return 0;

    std::size_t samples_to_read = std::min(max_samples, total_samples_ - read_pos_);
    if (samples_to_read == 0) return 0;

    int bytes_per_sample = bits_per_sample_ / 8;
    if (bytes_per_sample == 0) bytes_per_sample = 1;
    int block_align = num_channels_ * bytes_per_sample;
    bool is_float = (bits_per_sample_ == 32); // Assume 32-bit is float if audio_format==3

    std::vector<uint8_t> raw(samples_to_read * block_align);
    fseek(fp_, static_cast<long>(data_start_ + read_pos_ * block_align), SEEK_SET);
    std::size_t read_blocks = fread(raw.data(), block_align, samples_to_read, fp_);
    if (read_blocks == 0) return 0;

    out.resize(num_channels_);
    for (int ch = 0; ch < num_channels_; ++ch) {
        out[ch].resize(read_blocks);
    }

    for (std::size_t i = 0; i < read_blocks; ++i) {
        for (int ch = 0; ch < num_channels_; ++ch) {
            out[ch][i] = convert_sample(&raw[i * block_align + ch * bytes_per_sample],
                                         bytes_per_sample, is_float);
        }
    }

    read_pos_ += read_blocks;

    // Resample if needed
    if (target_sample_rate_ > 0 && target_sample_rate_ != file_sample_rate_) {
        for (int ch = 0; ch < num_channels_; ++ch) {
            resample_if_needed(out[ch], file_sample_rate_, target_sample_rate_);
        }
        return out[0].size();
    }
    return read_blocks;
}

} // namespace acoustic
