#include "engine/framework/audio/dsp.h"
#include "engine/framework/audio/probability_segmenter.h"
#include "engine/framework/audio/resampling.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstdint>
#include <iostream>
#include <random>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>

namespace {

constexpr double kPi = 3.14159265358979323846264338327950288;

enum class WaveformMode {
    Fixed,
    PerRun,
};

struct WaveformProfile {
    double envelope_bias = 0.0;
    double envelope_rate_scale = 1.0;
    double envelope_phase = 0.0;
    double voiced_rate_scale = 1.0;
    double fricative_rate_scale = 1.0;
    double pause_period = 2.4;
    double pause_start = 1.95;
    double pause_end = 2.20;
};

const WaveformProfile & get_per_run_waveform_profile() {
    static const WaveformProfile profile = []() {
        const auto ticks = static_cast<uint64_t>(
            std::chrono::steady_clock::now().time_since_epoch().count());
        std::seed_seq seed{
            static_cast<uint32_t>(ticks),
            static_cast<uint32_t>(ticks >> 32),
            0x6d2b79f5u,
            0x9e3779b9u,
        };
        std::mt19937 rng(seed);
        std::uniform_real_distribution<double> envelope_bias(-0.08, 0.08);
        std::uniform_real_distribution<double> envelope_rate_scale(0.85, 1.18);
        std::uniform_real_distribution<double> envelope_phase(0.0, 2.0 * kPi);
        std::uniform_real_distribution<double> voiced_rate_scale(0.92, 1.10);
        std::uniform_real_distribution<double> fricative_rate_scale(0.88, 1.15);
        std::uniform_real_distribution<double> pause_period(2.0, 2.8);
        std::uniform_real_distribution<double> pause_start(1.6, 2.1);
        std::uniform_real_distribution<double> pause_width(0.14, 0.32);
        WaveformProfile value;
        value.envelope_bias = envelope_bias(rng);
        value.envelope_rate_scale = envelope_rate_scale(rng);
        value.envelope_phase = envelope_phase(rng);
        value.voiced_rate_scale = voiced_rate_scale(rng);
        value.fricative_rate_scale = fricative_rate_scale(rng);
        value.pause_period = pause_period(rng);
        value.pause_start = pause_start(rng);
        value.pause_end = std::min(value.pause_period - 0.05, value.pause_start + pause_width(rng));
        return value;
    }();
    return profile;
}

void require(bool condition, const std::string & message) {
    if (!condition) {
        throw std::runtime_error(message);
    }
}

void require_shape_equal(
    const std::vector<int64_t> & actual,
    const std::vector<int64_t> & expected,
    const std::string & label) {
    if (actual != expected) {
        std::ostringstream oss;
        oss << label << " shape mismatch";
        throw std::runtime_error(oss.str());
    }
}

void require_close(
    const std::vector<float> & actual,
    const std::vector<float> & expected,
    float max_allowed,
    double mean_allowed,
    const std::string & label) {
    require(actual.size() == expected.size(), label + " size mismatch");
    float max_diff = 0.0f;
    size_t max_index = 0;
    double mean_diff = 0.0;
    for (size_t i = 0; i < actual.size(); ++i) {
        const float diff = std::fabs(actual[i] - expected[i]);
        mean_diff += diff;
        if (diff > max_diff) {
            max_diff = diff;
            max_index = i;
        }
    }
    mean_diff /= static_cast<double>(actual.size());
    if (max_diff > max_allowed || mean_diff > mean_allowed) {
        std::ostringstream oss;
        oss << label << " drift too large: max_diff=" << max_diff
            << " mean_diff=" << mean_diff
            << " at index " << max_index
            << " expected=" << expected[max_index]
            << " actual=" << actual[max_index];
        throw std::runtime_error(oss.str());
    }
}

std::vector<float> make_speech_like_waveform(
    int64_t sample_rate,
    double seconds,
    WaveformMode mode = WaveformMode::Fixed,
    int variant = 0) {
    const int64_t samples = static_cast<int64_t>(sample_rate * seconds);
    std::vector<float> waveform(static_cast<size_t>(samples), 0.0f);
    const WaveformProfile & profile =
        mode == WaveformMode::PerRun
            ? get_per_run_waveform_profile()
            : WaveformProfile{};
    const double variant_scale = 1.0 + 0.035 * static_cast<double>(variant);
    const double envelope_phase = profile.envelope_phase + 0.47 * static_cast<double>(variant);
    const double voiced_scale = profile.voiced_rate_scale * variant_scale;
    const double fricative_scale = profile.fricative_rate_scale * (1.0 + 0.02 * static_cast<double>(variant));
    const double pause_period = std::max(1.5, profile.pause_period + 0.06 * static_cast<double>(variant));
    const double pause_start = std::min(pause_period - 0.12, profile.pause_start + 0.03 * static_cast<double>(variant));
    const double pause_end = std::min(pause_period - 0.04, profile.pause_end + 0.03 * static_cast<double>(variant));
    for (int64_t i = 0; i < samples; ++i) {
        const double t = static_cast<double>(i) / static_cast<double>(sample_rate);
        const double envelope =
            0.55 + profile.envelope_bias +
            0.30 * std::sin(2.0 * kPi * 0.7 * profile.envelope_rate_scale * t + envelope_phase) +
            0.10 * std::sin(2.0 * kPi * 0.11 * profile.envelope_rate_scale * t + 0.3 + envelope_phase * 0.5);
        const double voiced =
            0.55 * std::sin(2.0 * kPi * 180.0 * voiced_scale * t) +
            0.25 * std::sin(2.0 * kPi * 360.0 * voiced_scale * t + 0.2 + envelope_phase * 0.25) +
            0.12 * std::sin(2.0 * kPi * 520.0 * voiced_scale * t + 0.5 + envelope_phase * 0.4);
        const double fricative =
            0.03 * std::sin(2.0 * kPi * 1900.0 * fricative_scale * t) +
            0.02 * std::sin(2.0 * kPi * 2600.0 * fricative_scale * t + 0.1 + envelope_phase * 0.15);
        const double pause_phase = std::fmod(t, pause_period);
        const double pause_gate = (pause_phase > pause_start && pause_phase < pause_end) ? 0.15 : 1.0;
        waveform[static_cast<size_t>(i)] = static_cast<float>(pause_gate * envelope * (voiced + fricative));
    }
    return waveform;
}

engine::audio::AudioTensor naive_istft_reference(
    const std::vector<float> & complex_spec,
    const std::vector<float> & window,
    int64_t batch,
    int64_t freq_bins,
    int64_t frames,
    int64_t samples,
    const engine::audio::STFTConfig & config) {
    const int64_t pad = config.n_fft / 2;
    const int64_t padded_samples = samples + 2 * pad;
    std::vector<float> accum(static_cast<size_t>(batch * padded_samples), 0.0f);
    std::vector<float> window_sums(static_cast<size_t>(batch * padded_samples), 0.0f);

    for (int64_t b = 0; b < batch; ++b) {
        for (int64_t frame_index = 0; frame_index < frames; ++frame_index) {
            std::vector<std::complex<float>> spectrum(static_cast<size_t>(config.n_fft), 0.0f);
            for (int64_t k = 0; k < freq_bins; ++k) {
                const size_t base = static_cast<size_t>((((b * freq_bins) + k) * frames + frame_index) * 2);
                spectrum[static_cast<size_t>(k)] = {complex_spec[base], complex_spec[base + 1]};
            }
            for (int64_t k = 1; k < freq_bins - 1; ++k) {
                spectrum[static_cast<size_t>(config.n_fft - k)] = std::conj(spectrum[static_cast<size_t>(k)]);
            }

            const int64_t start = frame_index * config.hop_length;
            for (int64_t n = 0; n < config.n_fft; ++n) {
                std::complex<float> sum = 0.0f;
                for (int64_t k = 0; k < config.n_fft; ++k) {
                    const float angle = 2.0f * kPi * static_cast<float>(k * n) / static_cast<float>(config.n_fft);
                    sum += spectrum[static_cast<size_t>(k)] * std::complex<float>(std::cos(angle), std::sin(angle));
                }
                float value = sum.real() / static_cast<float>(config.n_fft);
                if (n < config.win_length) {
                    value *= window[static_cast<size_t>(n)];
                } else {
                    value = 0.0f;
                }
                accum[static_cast<size_t>(b * padded_samples + start + n)] += value;
                if (n < config.win_length) {
                    const float w = window[static_cast<size_t>(n)];
                    window_sums[static_cast<size_t>(b * padded_samples + start + n)] += w * w;
                }
            }
        }
    }

    engine::audio::AudioTensor result;
    result.shape = {batch, samples};
    result.values.assign(static_cast<size_t>(batch * samples), 0.0f);
    for (int64_t b = 0; b < batch; ++b) {
        for (int64_t i = 0; i < samples; ++i) {
            const size_t padded_index = static_cast<size_t>(b * padded_samples + i + pad);
            const float denom = window_sums[padded_index] > 1e-8f ? window_sums[padded_index] : 1.0f;
            result.values[static_cast<size_t>(b * samples + i)] = accum[padded_index] / denom;
        }
    }
    return result;
}

void test_log_mel_matches_reference_pipeline() {
    constexpr int64_t kSampleRate = 16000;
    constexpr float kLogZeroGuard = 5.960464477539063e-8f;  // 2^-24
    const auto waveform = make_speech_like_waveform(kSampleRate, 10.0);
    const engine::audio::STFTConfig config{
        512,
        160,
        400,
        true,
        engine::audio::STFTPadMode::Reflect,
    };

    auto reference = engine::audio::MelSpectrogram().compute(
        waveform,
        1,
        static_cast<int64_t>(waveform.size()),
        kSampleRate,
        config,
        80);
    for (float & value : reference.values) {
        value = std::log(value + kLogZeroGuard);
    }
    const auto candidate = engine::audio::LogMelSpectrogram().compute(
        waveform,
        1,
        static_cast<int64_t>(waveform.size()),
        kSampleRate,
        config,
        80);

    require_shape_equal(candidate.shape, reference.shape, "log_mel");
    require_close(candidate.values, reference.values, 2.0e-5f, 2.0e-6, "log_mel");
}

void test_stft_istft_round_trip_matches_waveform() {
    constexpr int64_t kSampleRate = 24000;
    const auto waveform = make_speech_like_waveform(kSampleRate, 3.0);
    const engine::audio::STFTConfig config{
        20,
        5,
        20,
        true,
        engine::audio::STFTPadMode::Reflect,
        engine::audio::STFTFamily::Kokoro,
    };

    const auto & window = engine::audio::get_cached_stft_window(config);
    const auto complex = engine::audio::STFT().compute_complex(
        waveform,
        window,
        1,
        static_cast<int64_t>(waveform.size()),
        config);

    require_shape_equal(
        complex.shape,
        {1, (config.n_fft / 2) + 1, 1 + static_cast<int64_t>(waveform.size()) / config.hop_length, 2},
        "stft_complex");

    const auto reconstructed = engine::audio::ISTFT().compute(
        complex.values,
        window,
        1,
        complex.shape[1],
        complex.shape[2],
        static_cast<int64_t>(waveform.size()),
        config);

    require_shape_equal(
        reconstructed.shape,
        {1, static_cast<int64_t>(waveform.size())},
        "istft_waveform");
    require_close(
        reconstructed.values,
        waveform,
        2.0e-4f,
        2.0e-5,
        "stft_istft_round_trip");
}

void test_hann_window_modes_and_cache() {
    using engine::audio::STFTFamily;
    constexpr long double pi = 3.14159265358979323846264338327950288L;
    for (const int64_t length : {1, 20, 640, 1764}) {
        for (const auto family : {STFTFamily::Default, STFTFamily::Kokoro, STFTFamily::PeriodicF32}) {
            const engine::audio::STFTConfig config{length, 1, length, true,
                engine::audio::STFTPadMode::Reflect, family};
            const auto & actual = engine::audio::get_cached_stft_window(config);
            std::vector<float> expected(length, 1.0f);
            for (int64_t i = 0; length > 1 && i < length; ++i) {
                if (family == STFTFamily::PeriodicF32) {
                    const float step = static_cast<float>(2.0 * static_cast<double>(pi) / length);
                    expected[i] = 0.5f - 0.5f * std::cos(step * static_cast<float>(i));
                } else {
                    const float denominator = static_cast<float>(length - (family == STFTFamily::Default));
                    expected[i] = 0.5f - 0.5f * std::cos(2.0f * pi * static_cast<float>(i) / denominator);
                }
            }
            require_close(actual, expected, 0.0f, 0.0, "hann_window_mode");
            if (&actual != &engine::audio::get_cached_stft_window(config)) {
                throw std::runtime_error("Hann window cache did not reuse its entry");
            }
        }
    }
}

void test_resample_kernel_cache_eviction() {
    using engine::audio::resample_kernel_cache_stats;
    using engine::audio::set_resample_kernel_cache_budget_bytes;
    std::vector<float> input(2000);
    for (size_t i = 0; i < input.size(); ++i) {
        input[i] = std::sin(0.05f * static_cast<float>(i));
    }
    const auto options = engine::audio::torchaudio_sinc_hann_float32_options();
    const auto resample = [&](int source_rate, int target_rate) {
        return engine::audio::resample_mono_torchaudio_sinc_hann(input, source_rate, target_rate, options);
    };
    // 16 KiB holds a few of these 24 kernels (93 KiB in all), so the sweep evicts.
    constexpr size_t kBudget = 16 * 1024;
    std::vector<int> rates;
    for (int rate = 17000; rate <= 40000; rate += 1000) {
        rates.push_back(rate);
    }
    set_resample_kernel_cache_budget_bytes(kBudget);
    std::vector<std::vector<float>> expected;
    for (const int rate : rates) {
        expected.push_back(resample(rate, 16000));
        require(resample_kernel_cache_stats().bytes <= kBudget, "resample kernel cache exceeded its budget");
    }
    require(resample_kernel_cache_stats().entries < rates.size(), "resample kernel sweep did not evict");
    for (size_t i = 0; i < rates.size(); ++i) {
        require_close(resample(rates[i], 16000), expected[i], 0.0f, 0.0, "resample_after_eviction");
    }

    // Pairs that reduce to 3:1 build kernels of equal size under distinct keys.
    set_resample_kernel_cache_budget_bytes(0);
    set_resample_kernel_cache_budget_bytes(kBudget);
    resample(24000, 8000);
    const size_t kernel_bytes = resample_kernel_cache_stats().bytes;
    set_resample_kernel_cache_budget_bytes(3 * kernel_bytes);
    resample(48000, 16000);
    resample(96000, 32000);
    resample(24000, 8000);
    resample(12000, 4000);
    auto stats = resample_kernel_cache_stats();
    require(stats.entries == 3, "resample kernel cache did not hold three equal kernels");
    resample(24000, 8000);
    require(resample_kernel_cache_stats().hits == stats.hits + 1, "recently used resample kernel was evicted");
    resample(48000, 16000);
    require(resample_kernel_cache_stats().misses == stats.misses + 1, "least recently used resample kernel was kept");

    set_resample_kernel_cache_budget_bytes(kernel_bytes);
    stats = resample_kernel_cache_stats();
    require(stats.entries == 1 && stats.bytes == kernel_bytes, "shrinking the budget did not evict at once");
    resample(48000, 16000);
    require(resample_kernel_cache_stats().hits == stats.hits + 1, "shrinking evicted the most recently used kernel");

    set_resample_kernel_cache_budget_bytes(0);
    for (size_t i = 0; i < rates.size(); ++i) {
        require_close(resample(rates[i], 16000), expected[i], 0.0f, 0.0, "resample_uncached");
    }
    stats = resample_kernel_cache_stats();
    require(stats.entries == 0 && stats.bytes == 0, "budget 0 cached a resample kernel");

    // Concurrent misses evict entries other threads may still be using.
    set_resample_kernel_cache_budget_bytes(kBudget);
    std::vector<std::string> errors(4);
    std::vector<std::thread> threads;
    for (size_t t = 0; t < errors.size(); ++t) {
        threads.emplace_back([&, t] {
            try {
                for (int pass = 0; pass < 3; ++pass) {
                    for (size_t i = t; i < rates.size(); i += 2) {
                        require_close(resample(rates[i], 16000), expected[i], 0.0f, 0.0,
                                      "resample_concurrent_eviction");
                    }
                }
            } catch (const std::exception & ex) {
                errors[t] = ex.what();
            }
        });
    }
    for (auto & thread : threads) {
        thread.join();
    }
    for (const auto & error : errors) {
        require(error.empty(), error);
    }
    require(resample_kernel_cache_stats().bytes <= kBudget, "concurrent resampling exceeded the kernel budget");
    // The library default, for any later test.
    set_resample_kernel_cache_budget_bytes(32 * 1024 * 1024);
}

void test_istft_matches_reference_across_configs_and_variants() {
    const std::vector<engine::audio::STFTConfig> configs{
        engine::audio::STFTConfig{
            20,
            5,
            20,
            true,
            engine::audio::STFTPadMode::Reflect,
            engine::audio::STFTFamily::Kokoro,
        },
        engine::audio::STFTConfig{
            512,
            160,
            400,
            true,
            engine::audio::STFTPadMode::Reflect,
            engine::audio::STFTFamily::Default,
        },
    };

    for (const auto & config : configs) {
        for (int variant = 0; variant < 4; ++variant) {
            const auto waveform = make_speech_like_waveform(
                24000,
                2.0,
            WaveformMode::PerRun,
            variant);
            const auto & window = engine::audio::get_cached_stft_window(config);
            const auto complex = engine::audio::STFT().compute_complex(
                waveform,
                window,
                1,
                static_cast<int64_t>(waveform.size()),
                config);

            const auto reference = naive_istft_reference(
                complex.values,
                window,
                1,
                complex.shape[1],
                complex.shape[2],
                static_cast<int64_t>(waveform.size()),
                config);
            const auto reconstructed = engine::audio::ISTFT().compute(
                complex.values,
                window,
                1,
                complex.shape[1],
                complex.shape[2],
                static_cast<int64_t>(waveform.size()),
                config);

            require_shape_equal(reconstructed.shape, reference.shape, "istft_variant_shape");
            require_close(
                reconstructed.values,
                reference.values,
                3.0e-5f,
                5.0e-6,
                "istft_variant_parity");
        }
    }
}

void test_probability_segments() {
    using engine::audio::ProbabilitySegmenter;
    ProbabilitySegmenter frames({0.5f});
    require(!frames.push(0.1f, 0, 1), "leading silence");
    require(!frames.push(0.5f, 1, 2), "inclusive threshold");
    require(!frames.push(1.0f, 2, 3), "contiguous speech");
    const auto first = frames.push(0.0f, 3, 4);
    require(first && first->start == 1 && first->end == 3 && first->confidence == 0.75f,
            "frame segment and confidence");
    require(!frames.finish(), "finished segment must not repeat");
    require(!frames.push(1.0f, 4, 5), "next segment");
    const auto last = frames.finish();
    require(last && last->start == 4 && last->end == 5 && last->confidence == 1.0f,
            "finish resets confidence and boundaries");

    ProbabilitySegmenter windows({0.5f, 3200, 800});
    require(!windows.push(1.0f, 0, 3200), "positive window");
    require(!windows.push(0.0f, 1600, 4800), "overlap is not silence duration");
    require(!windows.push(0.0f, 3200, 6400), "silence shorter than minimum");
    const auto window = windows.push(0.0f, 4000, 7200);
    require(window && window->start == 0 && window->end == 3200,
            "end is last positive window end");

    ProbabilitySegmenter minimum({0.5f, 2, 0});
    require(!minimum.push(1.0f, 0, 1), "short segment begins");
    require(!minimum.push(0.0f, 1, 2), "short segment discarded");
    require(!minimum.push(0.5f, 2, 3), "reuse after discard");
    require(!minimum.push(0.5f, 3, 4), "minimum duration reached");
    const auto kept = minimum.finish();
    require(kept && kept->start == 2 && kept->end == 4 && kept->confidence == 0.5f,
            "minimum duration is inclusive");
}

}  // namespace

int main() {
    try {
        test_probability_segments();
        test_log_mel_matches_reference_pipeline();
        test_hann_window_modes_and_cache();
        test_resample_kernel_cache_eviction();
        test_stft_istft_round_trip_matches_waveform();
        test_istft_matches_reference_across_configs_and_variants();
        std::cout << "audio_dsp_test: ok\n";
        return 0;
    } catch (const std::exception & ex) {
        std::cerr << "audio_dsp_test: " << ex.what() << "\n";
        return 1;
    }
}