#include <mbgl/actor/scheduler.hpp>
#include <mbgl/gfx/uniform_buffer.hpp>
#include <mbgl/renderer/layer_group.hpp>
#include <mbgl/renderer/render_orchestrator.hpp>
#include <mbgl/renderer/renderer_observer.hpp>
#include <mbgl/renderer/retained_underlay_cache.hpp>
#include <mbgl/renderer/render_source.hpp>
#include <mbgl/renderer/render_source_observer.hpp>
#include <mbgl/renderer/tile_render_data.hpp>
#include <mbgl/style/sources/geojson_source_impl.hpp>
#include <mbgl/tile/tile.hpp>
#include <mbgl/tile/tile_observer.hpp>

#include <gtest/gtest.h>

#include <atomic>
#include <stdexcept>

using namespace mbgl;

namespace mbgl {
struct RenderOrchestratorTeardownTestAccess {
    static void add(RenderOrchestrator& orchestrator, std::unique_ptr<RenderSource> source) {
        source->setObserver(&orchestrator);
        orchestrator.renderSources.emplace(source->baseImpl->id, std::move(source));
    }
};
} // namespace mbgl

namespace {

class TestUniformBuffers final : public gfx::UniformBufferArray {
public:
    void bind(gfx::RenderPass&) override {}
protected:
    std::unique_ptr<gfx::UniformBuffer> copy(const gfx::UniformBuffer&) override { return nullptr; }
};

class TestLayerGroup final : public LayerGroup {
public:
    explicit TestLayerGroup(int32_t index)
        : LayerGroup(index, 0, "step097-error-group") {}

    const gfx::UniformBufferArray& getUniformBuffers() const override { return buffers; }
    gfx::UniformBufferArray& mutableUniformBuffers() override { return buffers; }

private:
    TestUniformBuffers buffers;
};

class ErrorObserver final : public RendererObserver {
public:
    void onRenderError(std::exception_ptr error) override {
        ++errors;
        try {
            std::rethrow_exception(error);
        } catch (const std::runtime_error& e) {
            message = e.what();
        }
    }

    int errors = 0;
    std::string message;
};

} // namespace

TEST(RenderOrchestratorError, RangeVisitorReportsOpaqueAndTranslucentFailures) {
    TaggedScheduler scheduler{Scheduler::GetBackground(), {}};
    RenderOrchestrator orchestrator(false, scheduler, std::nullopt);
    ErrorObserver observer;
    orchestrator.setObserver(&observer);
    ASSERT_TRUE(orchestrator.addLayerGroup(std::make_shared<TestLayerGroup>(0)));
    ASSERT_TRUE(orchestrator.addLayerGroup(std::make_shared<TestLayerGroup>(1)));

    bool failed = false;
    orchestrator.visitLayerGroupsInRangeReversed(0, 2, [](LayerGroupBase&, size_t) {
        throw std::runtime_error("opaque failure");
    }, &failed);
    EXPECT_TRUE(failed);
    EXPECT_EQ(observer.errors, 2);
    EXPECT_EQ(observer.message, "opaque failure");

    failed = false;
    observer.errors = 0;
    orchestrator.visitLayerGroupsInRange(0, 2, [](LayerGroupBase&, size_t) {
        throw std::runtime_error("translucent failure");
    }, &failed);
    EXPECT_TRUE(failed);
    EXPECT_EQ(observer.errors, 2);
    EXPECT_EQ(observer.message, "translucent failure");

    failed = false;
    int visits = 0;
    orchestrator.visitLayerGroupsInRange(0, 2, [&](LayerGroupBase&, size_t) { ++visits; }, &failed);
    EXPECT_FALSE(failed);
    EXPECT_EQ(visits, 2);
}

TEST(RenderOrchestratorError, FailedRangeCannotReuseAndCleanRetryCanRecover) {
    TaggedScheduler scheduler{Scheduler::GetBackground(), {}};
    RenderOrchestrator orchestrator(false, scheduler, std::nullopt);
    ErrorObserver observer;
    orchestrator.setObserver(&observer);
    ASSERT_TRUE(orchestrator.addLayerGroup(std::make_shared<TestLayerGroup>(0)));

    RetainedUnderlayCache cache;
    RetainedUnderlayKey key{{4, 4}, 0, 1, {}};
    key.cameraMatrix[0] = key.cameraMatrix[5] = key.cameraMatrix[10] = key.cameraMatrix[15] = 1.0;
    for (bool opaque : {true, false}) {
        EXPECT_EQ(cache.evaluate(key, 1, true), RetainedUnderlayDecision::Build);
        bool failed = false;
        const auto throwOnRender = [](LayerGroupBase&, size_t) { throw std::runtime_error("render failure"); };
        if (opaque) orchestrator.visitLayerGroupsInRangeReversed(0, 1, throwOnRender, &failed);
        else orchestrator.visitLayerGroupsInRange(0, 1, throwOnRender, &failed);
        ASSERT_TRUE(failed);
        cache.markBuildFailed("layer-render-error");
        EXPECT_FALSE(cache.isValid());
        EXPECT_EQ(cache.evaluate(key, 1, true), RetainedUnderlayDecision::Build);
    }
    bool failed = false;
    orchestrator.visitLayerGroupsInRange(0, 1, [](LayerGroupBase&, size_t) {}, &failed);
    ASSERT_FALSE(failed);
    cache.markBuildSucceeded();
    EXPECT_EQ(cache.evaluate(key, 1, true), RetainedUnderlayDecision::Reuse);
}

namespace {
class TeardownTile final : public Tile {
public:
    TeardownTile(TileObserver* observer_)
        : Tile(Kind::Geometry, OverscaledTileID{0, 0, {0, 0, 0}}, "step101-source", observer_) {}

    ~TeardownTile() override { observer->onTileAction(id, sourceID, TileOperation::Cancelled); }
    std::unique_ptr<TileRenderData> createRenderData() override { return nullptr; }
    void cancel() override {}
    bool layerPropertiesUpdated(const Immutable<style::LayerProperties>&) override { return true; }
    void emitChanged() { observer->onTileChanged(*this); }
    void emitError() { observer->onTileError(*this, std::make_exception_ptr(std::runtime_error("tile error"))); }
    void emitAction() { observer->onTileAction(id, sourceID, TileOperation::Cancelled); }
};

class TeardownSource final : public RenderSource {
public:
    TeardownSource(const TaggedScheduler& scheduler_, std::atomic<int>& deferred_, std::atomic<int>& renderJobs_,
                   std::atomic<int>& detached_)
        : RenderSource(makeMutable<style::GeoJSONSource::Impl>(
              "step101-source", style::GeoJSONOptions::defaultOptions())),
          scheduler(scheduler_), deferred(deferred_), renderJobs(renderJobs_), detached(detached_),
          tile(std::make_unique<TeardownTile>(static_cast<TileObserver*>(this))) {}

    ~TeardownSource() override {
        if (observer != oldObserver) ++detached;
        tile.reset(); // A pending tile emits Cancelled synchronously during source destruction.
        scheduler.schedule([&count = deferred] { ++count; });
        scheduler.runOnRenderThread([&count = renderJobs] { ++count; });
    }

    void expectDetachFrom(RenderSourceObserver* observer_) { oldObserver = observer_; }

    void emitAll() {
        tile->emitChanged();
        tile->emitError();
        tile->emitAction();
    }
    bool isLoaded() const override { return false; }
    void update(Immutable<style::Source::Impl>, const std::vector<Immutable<style::LayerProperties>>&,
                bool, bool, const TileParameters&) override {}
    std::unique_ptr<RenderItem> createRenderItem() override { return nullptr; }
    void prepare(const SourcePrepareParameters&) override {}
    void updateFadingTiles() override {}
    bool hasFadingTiles() const override { return false; }
    RenderSourceReadinessStats getReadinessDiagnostics() const override {
        return {.tileBacked = true,
                .tileCount = std::nullopt,
                .incompleteTileCount = std::nullopt,
                .loadedTileCount = std::nullopt,
                .hasFadingTiles = false};
    }
    std::unordered_map<std::string, std::vector<Feature>> queryRenderedFeatures(
        const ScreenLineString&, const TransformState&,
        const std::unordered_map<std::string, const RenderLayer*>&,
        const RenderedQueryOptions&, const mat4&) const override { return {}; }
    std::vector<Feature> querySourceFeatures(const SourceQueryOptions&) const override { return {}; }
    void reduceMemoryUse() override {}
    void dumpDebugLogs() const override {}

private:
    TaggedScheduler scheduler;
    std::atomic<int>& deferred;
    std::atomic<int>& renderJobs;
    std::atomic<int>& detached;
    RenderSourceObserver* oldObserver = nullptr;
    std::unique_ptr<TeardownTile> tile;
};

class CountingSourceObserver final : public RenderSourceObserver {
public:
    void onTileChanged(RenderSource&, const OverscaledTileID&) override { ++changed; }
    void onTileError(RenderSource&, const OverscaledTileID&, std::exception_ptr) override { ++errors; }
    void onTileAction(RenderSource&, TileOperation, const OverscaledTileID&, const std::string&) override { ++actions; }
    int changed = 0;
    int errors = 0;
    int actions = 0;
};
} // namespace

TEST(RenderOrchestratorTeardown, NullObserverAndDetachCoverEveryTileCallback) {
    TaggedScheduler scheduler{Scheduler::GetBackground(), {}};
    std::atomic<int> deferred{0};
    std::atomic<int> renderJobs{0};
    std::atomic<int> detached{0};
    auto source = std::make_unique<TeardownSource>(scheduler, deferred, renderJobs, detached);
    source->emitAll(); // Constructor default NullObserver.

    CountingSourceObserver observer;
    source->setObserver(&observer);
    source->expectDetachFrom(&observer);
    source->emitAll();
    EXPECT_EQ(observer.changed, 1);
    EXPECT_EQ(observer.errors, 1);
    EXPECT_EQ(observer.actions, 1);

    source->setObserver(nullptr);
    source->emitAll();
    source.reset(); // Cancelled must not reach the detached observer.
    EXPECT_EQ(detached.load(), 1);
    EXPECT_EQ(observer.changed, 1);
    EXPECT_EQ(observer.errors, 1);
    EXPECT_EQ(observer.actions, 1);
    scheduler.waitForEmpty();
    scheduler.runRenderJobs();
}

TEST(RenderOrchestratorTeardown, PendingTileCancellationAndSecondDrainRepeatSafely) {
    TaggedScheduler scheduler{Scheduler::GetBackground(), {}};
    std::atomic<int> deferred{0};
    std::atomic<int> renderJobs{0};
    std::atomic<int> detached{0};
    for (int iteration = 0; iteration < 20; ++iteration) {
        auto orchestrator = std::make_unique<RenderOrchestrator>(false, scheduler, std::nullopt);
        auto source = std::make_unique<TeardownSource>(scheduler, deferred, renderJobs, detached);
        source->expectDetachFrom(orchestrator.get());
        RenderOrchestratorTeardownTestAccess::add(*orchestrator, std::move(source));
        orchestrator.reset();
        EXPECT_EQ(detached.load(), iteration + 1); // Source was detached before its tile cancelled.
        EXPECT_EQ(deferred.load(), iteration + 1); // Work queued by source destruction was drained.
        scheduler.runRenderJobs(); // Render jobs have a separate backend-owned queue.
        EXPECT_EQ(renderJobs.load(), iteration + 1);
    }
}