#include "NoiseReducer.h"
#include "AudioBuffer.h"
#include <rnnoise.h>
#include <cstring>
#include <algorithm>
#include <iostream>
#include <filesystem>

namespace secondvoice {

NoiseReducer::NoiseReducer() {
    state_ = rnnoise_create(nullptr);
    upsample_buffer_.resize(RNNOISE_FRAME_SIZE);
    downsample_buffer_.resize(INPUT_FRAME_SIZE);
}

NoiseReducer::~NoiseReducer() {
    if (state_) {
        rnnoise_destroy(state_);
    }
}

void NoiseReducer::upsample(const float* in, float* out, int in_samples) {
    // Simple linear interpolation 16kHz -> 48kHz (3x)
    for (int i = 0; i < in_samples; ++i) {
        int out_idx = i * 3;
        float curr = in[i];
        float next = (i + 1 < in_samples) ? in[i + 1] : curr;

        out[out_idx] = curr;
        out[out_idx + 1] = curr + (next - curr) * 0.333f;
        out[out_idx + 2] = curr + (next - curr) * 0.667f;
    }
}

void NoiseReducer::downsample(const float* in, float* out, int out_samples) {
    // Simple decimation 48kHz -> 16kHz (take every 3rd sample with averaging)
    for (int i = 0; i < out_samples; ++i) {
        int in_idx = i * 3;
        // Average of 3 samples for anti-aliasing
        out[i] = (in[in_idx] + in[in_idx + 1] + in[in_idx + 2]) / 3.0f;
    }
}

void NoiseReducer::suppressTransients(float* samples, size_t count) {
    // Transient suppressor: attenuates sudden spikes (claps, clicks, pops)
    // while preserving speech which has smoother amplitude changes

    int transients_detected = 0;

    for (size_t i = 0; i < count; ++i) {
        float abs_sample = std::abs(samples[i]);

        // IMPORTANT: Use envelope BEFORE updating for transient detection
        // Otherwise the envelope rises with the clap and we miss it!
        float safe_envelope = std::max(envelope_, 0.001f);
        float ratio = abs_sample / safe_envelope;

        bool is_transient = (ratio > transient_threshold_);

        if (is_transient) {
            // This is a transient - attenuate it heavily
            float target_amplitude = safe_envelope * transient_ratio_;
            float attenuation = target_amplitude / abs_sample;
            samples[i] *= attenuation;
            transients_detected++;

            // DON'T update envelope from transients - keep it stable
            // This prevents the envelope from "learning" the clap
        } else {
            // Only update envelope from non-transient samples
            if (abs_sample > envelope_) {
                envelope_ = envelope_ * (1.0f - envelope_attack_) + abs_sample * envelope_attack_;
            } else {
                envelope_ = envelope_ * (1.0f - envelope_release_) + abs_sample * envelope_release_;
            }
        }
    }

    if (transients_detected > 10) {
        std::cout << "[RNNoise] Suppressed " << transients_detected << " transient samples (clap/click)" << std::endl;
    }
}

void NoiseReducer::process(std::vector<float>& samples) {
    if (!enabled_ || !state_ || samples.empty()) {
        return;
    }

    // Process in chunks of INPUT_FRAME_SIZE (160 samples = 10ms at 16kHz)
    size_t num_frames = samples.size() / INPUT_FRAME_SIZE;
    float duration_ms = static_cast<float>(samples.size()) / 16.0f;  // 16 samples/ms at 16kHz
    std::cout << "[RNNoise] Denoising " << duration_ms << "ms of audio (" << num_frames << " frames)" << std::endl;

    for (size_t frame = 0; frame < num_frames; ++frame) {
        float* frame_ptr = samples.data() + frame * INPUT_FRAME_SIZE;

        // Upsample 16kHz -> 48kHz
        upsample(frame_ptr, upsample_buffer_.data(), INPUT_FRAME_SIZE);

        // RNNoise expects float input in range [-32768, 32768] (like int16)
        // Our input is [-1, 1], so scale up
        for (int i = 0; i < RNNOISE_FRAME_SIZE; ++i) {
            upsample_buffer_[i] *= 32768.0f;
        }

        // Apply RNNoise
        rnnoise_process_frame(state_, upsample_buffer_.data(), upsample_buffer_.data());

        // Scale back to [-1, 1]
        for (int i = 0; i < RNNOISE_FRAME_SIZE; ++i) {
            upsample_buffer_[i] /= 32768.0f;
            // Clamp to prevent any overflow
            upsample_buffer_[i] = std::clamp(upsample_buffer_[i], -1.0f, 1.0f);
        }

        // Downsample 48kHz -> 16kHz
        downsample(upsample_buffer_.data(), frame_ptr, INPUT_FRAME_SIZE);
    }

    // Handle remaining samples (less than a full frame) - just leave them as-is
    // They'll be processed in the next call
}

std::vector<float> NoiseReducer::processCopy(const std::vector<float>& samples) {
    std::vector<float> result = samples;
    process(result);
    return result;
}

std::vector<float> NoiseReducer::processRealtime(const float* samples, size_t count) {
    std::vector<float> output;

    if (!enabled_ || !state_ || count == 0) {
        // Pass through unchanged
        output.assign(samples, samples + count);
        return output;
    }

    // Add new samples to streaming buffer
    stream_input_buffer_.insert(stream_input_buffer_.end(), samples, samples + count);

    // Process complete frames
    while (stream_input_buffer_.size() >= INPUT_FRAME_SIZE) {
        // Get frame pointer
        float* frame_ptr = stream_input_buffer_.data();

        // Step 1: Suppress transients (claps, clicks, pops) BEFORE RNNoise
        suppressTransients(frame_ptr, INPUT_FRAME_SIZE);

        // Step 2: Upsample 16kHz -> 48kHz
        upsample(frame_ptr, upsample_buffer_.data(), INPUT_FRAME_SIZE);

        // Scale to int16 range for RNNoise
        for (int i = 0; i < RNNOISE_FRAME_SIZE; ++i) {
            upsample_buffer_[i] *= 32768.0f;
        }

        // Apply RNNoise
        rnnoise_process_frame(state_, upsample_buffer_.data(), upsample_buffer_.data());

        // Scale back to [-1, 1]
        for (int i = 0; i < RNNOISE_FRAME_SIZE; ++i) {
            upsample_buffer_[i] /= 32768.0f;
            upsample_buffer_[i] = std::clamp(upsample_buffer_[i], -1.0f, 1.0f);
        }

        // Downsample 48kHz -> 16kHz and add to output
        std::vector<float> denoised_frame(INPUT_FRAME_SIZE);
        downsample(upsample_buffer_.data(), denoised_frame.data(), INPUT_FRAME_SIZE);
        output.insert(output.end(), denoised_frame.begin(), denoised_frame.end());

        // Remove processed samples from buffer
        stream_input_buffer_.erase(stream_input_buffer_.begin(),
                                   stream_input_buffer_.begin() + INPUT_FRAME_SIZE);
    }

    return output;
}

std::vector<float> NoiseReducer::flushStream() {
    std::vector<float> output;

    if (!enabled_ || !state_ || stream_input_buffer_.empty()) {
        // Return any remaining samples as-is
        output = std::move(stream_input_buffer_);
        stream_input_buffer_.clear();
        return output;
    }

    // Process remaining samples (zero-pad to complete frame if needed)
    while (!stream_input_buffer_.empty()) {
        // Pad with zeros if we have less than a full frame
        while (stream_input_buffer_.size() < INPUT_FRAME_SIZE) {
            stream_input_buffer_.push_back(0.0f);
        }

        // Get frame pointer
        float* frame_ptr = stream_input_buffer_.data();

        // Step 1: Suppress transients
        suppressTransients(frame_ptr, INPUT_FRAME_SIZE);

        // Step 2: Upsample 16kHz -> 48kHz
        upsample(frame_ptr, upsample_buffer_.data(), INPUT_FRAME_SIZE);

        // Scale to int16 range for RNNoise
        for (int i = 0; i < RNNOISE_FRAME_SIZE; ++i) {
            upsample_buffer_[i] *= 32768.0f;
        }

        // Apply RNNoise
        rnnoise_process_frame(state_, upsample_buffer_.data(), upsample_buffer_.data());

        // Scale back to [-1, 1]
        for (int i = 0; i < RNNOISE_FRAME_SIZE; ++i) {
            upsample_buffer_[i] /= 32768.0f;
            upsample_buffer_[i] = std::clamp(upsample_buffer_[i], -1.0f, 1.0f);
        }

        // Downsample 48kHz -> 16kHz and add to output
        std::vector<float> denoised_frame(INPUT_FRAME_SIZE);
        downsample(upsample_buffer_.data(), denoised_frame.data(), INPUT_FRAME_SIZE);
        output.insert(output.end(), denoised_frame.begin(), denoised_frame.end());

        // Remove processed samples from buffer
        stream_input_buffer_.erase(stream_input_buffer_.begin(),
                                   stream_input_buffer_.begin() + INPUT_FRAME_SIZE);
    }

    return output;
}

void NoiseReducer::savePreview(const std::vector<float>& samples) {
    if (!save_preview_ || samples.empty()) {
        return;
    }

    std::filesystem::create_directories("./denoised");
    std::string filename = "./denoised/denoised_" + std::to_string(preview_count_++) + ".ogg";

    AudioBuffer buffer(16000, 1);  // 16kHz mono
    buffer.addSamples(samples);
    if (buffer.saveToOpus(filename)) {
        std::cout << "[RNNoise] Saved denoised preview: " << filename << std::endl;
    }
}

} // namespace secondvoice
