#include "livepipeline.h"

#include <chrono>

namespace lc {

namespace {
constexpr int kFrameSamples = 320;   // 20ms @16k
} // namespace

LivePipeline::LivePipeline(IAsr & asr, LivePipelineOptions opt)
    : asr_(asr), opt_(std::move(opt)), stream_(asr, opt_.stream) {}

LivePipeline::~LivePipeline()
{
    stop();
}

std::string LivePipeline::deviceName() const
{
    return capture_.deviceName();
}

bool LivePipeline::start(std::string * error)
{
    const bool ok = capture_.start(
        asr_.sampleRate(),
        [this](const float * s, size_t n) {
            {
                std::lock_guard<std::mutex> lk(mutex_);
                pending_.insert(pending_.end(), s, s + n);
            }
            cv_.notify_one();
        },
        error,
        opt_.captureDevice);
    if (!ok) return false;

    worker_ = std::thread([this] { workerLoop(); });
    return true;
}

void LivePipeline::stop()
{
    {
        std::lock_guard<std::mutex> lk(mutex_);
        quit_ = true;
    }
    cv_.notify_all();
    if (worker_.joinable()) worker_.join();
    capture_.stop();
}

void LivePipeline::workerLoop()
{
    for (;;) {
        std::vector<float> chunk;
        {
            std::unique_lock<std::mutex> lk(mutex_);
            cv_.wait_for(lk, std::chrono::milliseconds(50),
                         [this] { return quit_ || !pending_.empty(); });
            if (quit_ && pending_.empty()) break;
            chunk.swap(pending_);
        }
        if (!chunk.empty()) process(chunk.data(), chunk.size());
    }

    if (inSegment_) finishSegment();
}

void LivePipeline::emitStep(const StreamingAsr::Step & s)
{
    if (!s.ok || !onCaption_) return;
    onCaption_(s.committed, s.tail, s.language);
}

void LivePipeline::process(const float * pcm, size_t n)
{
    buffer_.insert(buffer_.end(), pcm, pcm + n);

    const size_t   prerollSamples = (size_t) (opt_.prerollSec * asr_.sampleRate());
    const float    frameSec       = (float) kFrameSamples / asr_.sampleRate();

    std::vector<float> segAudio;   // 本段要送进引擎的音频
    bool flushNow = false;

    size_t off = 0;
    for (; off + kFrameSamples <= buffer_.size(); off += kFrameSamples) {
        const float * frame = buffer_.data() + off;
        const bool speech = vad_.feed(frame, kFrameSamples);

        if (!inSegment_) {
            // 静音期间什么都不送给模型：否则模型会对静音产生幻觉
            if (speech) {
                inSegment_  = true;
                silenceAcc_ = 0.0f;
                segAudio.insert(segAudio.end(), preroll_.begin(), preroll_.end());
                preroll_.clear();
                segAudio.insert(segAudio.end(), frame, frame + kFrameSamples);
            } else {
                preroll_.insert(preroll_.end(), frame, frame + kFrameSamples);
                if (preroll_.size() > prerollSamples) {
                    preroll_.erase(preroll_.begin(),
                                   preroll_.begin() + (long) (preroll_.size() - prerollSamples));
                }
            }
        } else {
            segAudio.insert(segAudio.end(), frame, frame + kFrameSamples);
            if (speech) silenceAcc_ = 0.0f;
            else        silenceAcc_ += frameSec;

            if (silenceAcc_ >= (float) opt_.silenceTimeoutSec) {
                flushNow = true;
                off += kFrameSamples;   // 这一帧已经并入 segAudio
                break;
            }
        }
    }

    if (!segAudio.empty()) {
        const auto cb = [this](const StreamingAsr::Step & s) { emitStep(s); };
        stream_.feed(segAudio.data(), segAudio.size(), cb);
    }

    if (flushNow) {
        const auto cb = [this](const StreamingAsr::Step & s) { emitStep(s); };
        stream_.flush(cb);
        stream_.reset();
        inSegment_  = false;
        silenceAcc_ = 0.0f;
        preroll_.clear();
    }

    if (off > 0) buffer_.erase(buffer_.begin(), buffer_.begin() + (long) off);
}

void LivePipeline::finishSegment()
{
    if (!inSegment_) return;
    const auto cb = [this](const StreamingAsr::Step & s) { emitStep(s); };
    stream_.flush(cb);
    stream_.reset();
    inSegment_  = false;
    silenceAcc_ = 0.0f;
    preroll_.clear();
}

} // namespace lc