#include "engine/community_models/parakeet_tdt/bark3.h"
#include "engine/community_models/parakeet_tdt/implementation.h"
#include "test_assert.h"
#include "engine/framework/core/backend.h"
#include "engine/framework/runtime/errors.h"
#include "ggml-alloc.h"

#include <algorithm>
#include <cmath>
#include <cstring>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <utility>

namespace rt = engine::runtime;
namespace bark = engine::community_models::parakeet_tdt;

namespace {
void streaming_frontend_test() {
    engine::test::require_eq(bark::kBarkImplementation.input_prefix_samples, int64_t(4000), "Bark prefix policy");
    engine::test::require_eq(bark::kParakeetImplementation.input_prefix_samples, int64_t(0), "Parakeet has no prefix");
    engine::test::require_eq(bark::kPhononImplementation.input_prefix_samples, int64_t(0), "Phonon has no prefix");
    auto assets = std::make_shared<bark::ParakeetTDTAssets>();
    auto & config = assets->config.frontend;
    config.sample_rate = 16000;
    config.feature_size = 128;
    config.n_fft = 512;
    config.hop_length = 160;
    config.win_length = 400;
    config.preemphasis = .97f;
    config.log_zero_guard = 0x1p-24f;
    bark::ParakeetFrontend reference(assets);
    bark::BarkStreamingFrontend cached(assets);
    std::vector<float> waveform(320000);
    for (size_t i = 0; i < waveform.size(); ++i)
        waveform[i] = .3f * std::sin(i * .017) + .2f * std::cos(i * .13);
    const auto check = [&](int64_t begin, int64_t end) {
        rt::AudioBuffer window{16000, 1, {waveform.begin() + begin, waveform.begin() + end}};
        const auto expected = reference.extract(window, true);
        const auto actual = cached.extract(window, begin);
        engine::test::require_eq(actual.frames, expected.frames, "cached frontend frame count");
        engine::test::require_eq(actual.valid_frames, expected.valid_frames, "cached frontend valid frames");
        engine::test::require(actual.values.size() == expected.values.size() &&
            std::memcmp(actual.values.data(), expected.values.data(), actual.values.size() * sizeof(float)) == 0,
            "cached frontend must preserve every normalized value, including boundary frames");
    };
    // Growing and sliding windows, unaligned tails/hop grids, short streams,
    // a gap with no overlap, and zero/one valid frame normalization.
    for (const auto & span : std::vector<std::pair<int64_t, int64_t>>{
            {0, 64000}, {0, 96000}, {0, 128000}, {32000, 160000},
            {64000, 192000}, {96000, 200013}, {128000, 200013},
            {160001, 210001}, {176001, 240006}, {300000, 310000}})
        check(span.first, span.second);
    cached.reset();
    std::fill(waveform.begin(), waveform.end(), 0.f);
    check(0, 64000); check(0, 96000); check(16000, 128000);
    cached.reset(); check(0, 1); check(0, 159); check(0, 160);
    cached.reset(); check(0, 64000);
    waveform[32000] = 1.f;
    cached.reset(); check(0, 64000); check(16000, 96000);
}

size_t allocation_probe(bool bark_policy) {
    auto context = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
        ggml_init({4 * 1024 * 1024, nullptr, true}), ggml_free);
    engine::test::require(context != nullptr, "probe context");
    auto * input = ggml_new_tensor_2d(context.get(), GGML_TYPE_F32, 128, 128);
    ggml_set_input(input);
    auto * result = input;
    for (int i = 0; i < 32; ++i) {
        auto * scaled = ggml_scale(context.get(), result, 2.f);
        auto * cropped = ggml_view_2d(context.get(), scaled, scaled->ne[0] - 1, 128, scaled->nb[1], sizeof(float));
        result = ggml_cont(context.get(), cropped);
    }
    ggml_set_output(result);
    auto * graph = ggml_new_graph_custom(context.get(), 1024, false);
    ggml_build_forward_expand(graph, result);
    const auto cpu = rt::GraphOptimizationBackend::Cpu;
    auto options = bark_policy ? bark::bark3_encoder_graph_options(cpu)
                               : rt::graph_optimization_options_for_backend(cpu);
    options.enabled = true;  // Exercise allocation planning even with a caller opt-out env var.
    rt::optimize_graph(*graph, options);
    auto backend = std::unique_ptr<ggml_backend, decltype(&ggml_backend_free)>(
        engine::core::init_backend({engine::core::BackendType::Cpu, 0, 1}), ggml_backend_free);
    auto allocator = std::unique_ptr<ggml_gallocr, decltype(&ggml_gallocr_free)>(
        ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend.get())), ggml_gallocr_free);
    size_t bytes = 0;
    ggml_gallocr_reserve_n_size(allocator.get(), graph, nullptr, nullptr, &bytes);
    if (bark_policy) {
        engine::test::require(ggml_gallocr_reserve(allocator.get(), graph), "reserve view-lifetime probe");
        engine::test::require(ggml_gallocr_alloc_graph(allocator.get(), graph), "allocate view-lifetime probe");
        std::vector<float> values(128 * 128);
        for (size_t i = 0; i < values.size(); ++i) values[i] = static_cast<float>(i + 1);
        ggml_backend_tensor_set(input, values.data(), 0, values.size() * sizeof(float));
        engine::test::require(ggml_backend_graph_compute(backend.get(), graph) == GGML_STATUS_SUCCESS, "compute probe");
        std::vector<float> actual(96 * 128);
        ggml_backend_tensor_get(result, actual.data(), 0, actual.size() * sizeof(float));
        for (int row = 0; row < 128; ++row)
            for (int col = 0; col < 96; ++col)
                engine::test::require_eq(actual[row * 96 + col], std::ldexp(values[row * 128 + col + 32], 32), "reused buffers retain correct values");
    }
    return bytes;
}
void allocation_test() {
    const auto old_bytes = allocation_probe(false);
    const auto fixed_bytes = allocation_probe(true);
    engine::test::require(fixed_bytes < old_bytes / 4, "view lifetimes bound temporary allocation across layers");
}
rt::WordTimestamp word(int64_t begin, int64_t end) {
    rt::WordTimestamp w; w.word = "hello"; w.span.start_sample = begin; w.span.end_sample = end; return w;
}
template<class F> void rejects(F operation) {
    bool rejected = false;
    try { operation(); } catch (const std::exception &) { rejected = true; }
    engine::test::require(rejected, "invalid operation must be rejected");
}
class Offline final : public rt::IOfflineVoiceTaskSession {
public:
    rt::SessionPreparationRequest preparation;
    rt::AudioBuffer input;
    bool fail = false;
    std::string family() const override { return "parakeet_tdt"; }
    rt::VoiceTaskKind task_kind() const override { return rt::VoiceTaskKind::Asr; }
    rt::RunMode run_mode() const override { return rt::RunMode::Offline; }
    void prepare(const rt::SessionPreparationRequest & r) override { preparation = r; }
    rt::TaskResult run(const rt::TaskRequest & r) override {
        input = *r.audio_input;
        if (std::exchange(fail, false)) throw std::runtime_error("injected failure");
        rt::TaskResult out; out.word_timestamps = {word(3000, 5000), word(6000, 999999)}; return out;
    }
};
class Streaming final : public rt::IStreamingVoiceTaskSession {
public:
    int64_t received = 0;
    bool fail = false;
    rt::StreamEventCallback sink;
    std::string family() const override { return "parakeet_tdt"; }
    rt::VoiceTaskKind task_kind() const override { return rt::VoiceTaskKind::Asr; }
    rt::RunMode run_mode() const override { return rt::RunMode::Streaming; }
    void prepare(const rt::SessionPreparationRequest &) override {}
    void reset() override { received = 0; }
    void set_stream_event_sink(rt::StreamEventCallback callback) override { sink = std::move(callback); }
    rt::StreamEvent process_audio_chunk(const rt::AudioChunk & chunk) override {
        engine::test::require_eq(chunk.start_sample, received, "padded chunks remain contiguous");
        if (received == 0) {
            engine::test::require(chunk.samples.size() >= 4000, "first chunk has prefix");
            engine::test::require(std::all_of(chunk.samples.begin(), chunk.samples.begin() + 4000, [](float v) { return v == 0; }), "prefix is silence");
        }
        received += static_cast<int64_t>(chunk.samples.size());
        if (std::exchange(fail, false)) throw std::runtime_error("injected failure");
        rt::StreamEvent event; event.word_timestamps = {word(4000, received)};
        if (sink) sink(event);
        return event;
    }
    rt::TaskResult finalize() override { rt::TaskResult out; out.word_timestamps = {word(4000, received)}; return out; }
};
void offline_test() {
    auto implementation = std::make_unique<Offline>(); auto * raw = implementation.get();
    auto session = bark::wrap_bark3_offline(std::move(implementation));
    rt::TaskRequest request; request.audio_input = rt::AudioBuffer{16000, 1, std::vector<float>(16000, .2f)};
    rejects([&] { session->run(request); });
    rt::SessionPreparationRequest preparation; preparation.audio = rt::AudioPreparationContract{16000, 1, 16000};
    session->prepare(preparation);
    engine::test::require_eq(raw->preparation.audio->max_input_samples, int64_t(20000), "reserve includes prefix");
    auto result = session->run(request);
    engine::test::require_eq(raw->input.samples.size(), size_t(20000), "pad once");
    engine::test::require_eq(request.audio_input->samples.size(), size_t(16000), "caller input remains intact");
    engine::test::require_eq(result.word_timestamps[0].span.start_sample, int64_t(0), "clamp negative time");
    engine::test::require_eq(result.word_timestamps[0].span.end_sample, int64_t(1000), "remove exact 250 ms");
    engine::test::require_eq(result.word_timestamps[1].span.end_sample, int64_t(16000), "clamp to original end");
    raw->fail = true; rejects([&] { session->run(request); }); session->run(request);
    engine::test::require_eq(raw->input.samples.size(), size_t(20000), "failure/reuse does not accumulate padding");
    auto vad = request; vad.options["audio_chunk_mode"] = "vad";
    session->run(vad);
    engine::test::require_eq(raw->input.samples.size(), size_t(20000), "offline VAD dispatch receives one recording prefix");
    session->run(request);
    auto bad = request; bad.audio_input->samples.clear(); rejects([&] { session->run(bad); });
    bad = request; bad.audio_input->channels = 3; rejects([&] { session->run(bad); });
    for (const auto * duration : {"0", "-1", "nan", "inf"}) {
        bad = request; bad.options["audio_chunk_duration_sec"] = duration;
        bool invalid = false;
        try { session->run(bad); } catch (const rt::InvalidRequestError &) { invalid = true; }
        engine::test::require(invalid, "invalid VAD duration is a client error");
        session->run(request);
    }
    auto stereo = request; stereo.audio_input = rt::AudioBuffer{8000, 2, std::vector<float>(4000, .2f)};
    session->run(stereo);
    engine::test::require_eq(raw->input.sample_rate, 16000, "resample before padding");
    engine::test::require_eq(raw->input.channels, 1, "downmix before padding");
    engine::test::require_eq(raw->input.samples.size(), size_t(8000), "prefix is 4000 target samples");
    preparation.audio->max_input_samples = std::numeric_limits<int64_t>::max();
    rejects([&] { session->prepare(preparation); }); rejects([&] { session->run(request); });
}
void streaming_test() {
    auto implementation = std::make_unique<Streaming>(); auto * raw = implementation.get();
    auto session = bark::wrap_bark3_streaming(std::move(implementation));
    session->prepare({});
    int64_t callback_end = -1;
    session->set_stream_event_sink([&](const rt::StreamEvent & event) { callback_end = event.word_timestamps[0].span.end_sample; });
    rt::TaskRequest vad; vad.options["audio_chunk_mode"] = "vad";
    rejects([&] { session->start_stream(vad); });
    session->start_stream({});
    rejects([&] { session->prepare({}); });
    auto event = session->process_audio_chunk({16000, 1, 0, std::vector<float>(100, .2f)});
    engine::test::require_eq(event.word_timestamps[0].span.end_sample, int64_t(100), "return timestamp shifted once");
    engine::test::require_eq(callback_end, int64_t(100), "callback timestamp shifted once");
    rejects([&] { session->process_audio_chunk({16000, 1, 99, {.2f}}); });
    session->process_audio_chunk({16000, 1, 100, std::vector<float>(50, .2f)});
    engine::test::require_eq(raw->received, int64_t(4150), "one prefix across chunks");
    engine::test::require_eq(session->finalize().word_timestamps[0].span.end_sample, int64_t(150), "final timestamps on original clock");
    rejects([&] { session->finalize(); });
    session->start_stream({}); rejects([&] { session->finalize(); });
    session->start_stream({}); raw->fail = true;
    rejects([&] { session->process_audio_chunk({16000, 1, 0, {.2f}}); });
    rejects([&] { session->process_audio_chunk({16000, 1, 1, {.2f}}); });
    session->start_stream({}); session->process_audio_chunk({16000, 1, 0, {.2f}}); session->finalize();
    engine::test::require_eq(raw->received, int64_t(4001), "failed stream reset recovers");
}
}  // namespace
int main() {
    try { streaming_frontend_test(); allocation_test(); offline_test(); streaming_test(); std::cout << "PASS Bark frontend reuse, allocation lifetimes, prefix, sample clocks, repeated requests, resampling, streaming and failure recovery\n"; }
    catch (const std::exception & e) { std::cerr << e.what() << '\n'; return 1; }
}