#include "pipipelinethread.h"
#include "pisystemtime.h"
#include "pithread.h"
#include "pitime.h"

#include "gtest/gtest.h"
#include <atomic>
#include <functional>
#include <thread>

namespace {

class Stage: public PIPipelineThread<int, int> {
	PIOBJECT_SUBCLASS(Stage, PIPipelineThread<int, int>)

public:
	std::function<int(int)> f;
	bool drop_odd = false;

	Stage(): f([](int v) { return v; }) {}

protected:
	int calc(int & v, bool & ok) override {
		ok = !drop_odd || !(v & 1);
		return f(v);
	}
};

// Stage with slow calculation, used to keep the consumer busy
class SlowStage: public PIPipelineThread<int, int> {
	PIOBJECT_SUBCLASS(SlowStage, PIPipelineThread<int, int>)

public:
	int calc_ms = 300;

protected:
	int calc(int & v, bool & ok) override {
		ok = true;
		piMSleep(calc_ms);
		return v;
	}
};

bool waitForCount(const PIPipelineThread<int, int> & stage, ullong count, double timeout_s = 5.) {
	PITimeMeasurer tm;
	while (stage.counter() < count && tm.elapsed_s() < timeout_s)
		piMSleep(1);
	return stage.counter() >= count;
}

} // namespace


TEST(PIPipelineThread_Basic, ProcessItems) {
	Stage stage;
	stage.f = [](int v) { return v * 2; };
	stage.start();

	const int N = 100;
	for (int i = 0; i < N; ++i)
		stage.enqueue(i);

	EXPECT_TRUE(waitForCount(stage, N));
	EXPECT_EQ(stage.counter(), N);
	EXPECT_EQ(stage.queSize(), 0);
	EXPECT_TRUE(stage.isEmpty());
	EXPECT_EQ(stage.getLast(), (N - 1) * 2);

	stage.stopCalc(2000);
	EXPECT_FALSE(stage.isRunning());
}


TEST(PIPipelineThread_Basic, DropByCalc) {
	Stage stage;
	stage.drop_odd = true;
	stage.start();

	for (int i = 0; i < 20; ++i)
		stage.enqueue(i);

	EXPECT_TRUE(waitForCount(stage, 10));
	EXPECT_EQ(stage.counter(), 10); // only even values were published
	EXPECT_EQ(stage.queSize(), 0);
	EXPECT_EQ(stage.getLast(), 18);

	stage.stopCalc(2000);
	EXPECT_FALSE(stage.isRunning());
}


TEST(PIPipelineThread_Chain, TwoStages) {
	Stage s1;
	Stage s2;
	s1.f = [](int v) { return v * 2; };
	s2.f = [](int v) { return v + 10; };
	s1.connectTo(&s2);

	s1.start();
	s2.start();

	const int N = 50;
	for (int i = 0; i < N; ++i)
		s1.enqueue(i);

	EXPECT_TRUE(waitForCount(s1, N));
	EXPECT_TRUE(waitForCount(s2, N));
	EXPECT_EQ(s2.queSize(), 0);
	EXPECT_EQ(s2.getLast(), (N - 1) * 2 + 10);

	s2.stopCalc(2000);
	s1.stopCalc(2000);
	EXPECT_FALSE(s1.isRunning());
	EXPECT_FALSE(s2.isRunning());
}


TEST(PIPipelineThread_Bounded, DropWhenFull) {
	Stage stage;
	stage.setMaxQueSize(2);
	EXPECT_EQ(stage.maxQueSize(), 2u);

	bool ov1 = false;
	bool ov2 = false;
	bool ov3 = true;
	stage.enqueue(1, false, &ov1);
	stage.enqueue(2, false, &ov2);
	stage.enqueue(3, false, &ov3);

	EXPECT_FALSE(ov1);
	EXPECT_FALSE(ov2);
	EXPECT_TRUE(ov3);
	EXPECT_EQ(stage.queSize(), 2);
}


TEST(PIPipelineThread_Bounded, ResizeTrimsQueue) {
	Stage stage;
	stage.setMaxQueSize(10);
	for (int i = 0; i < 5; ++i)
		stage.enqueue(i);
	EXPECT_EQ(stage.queSize(), 5);

	stage.setMaxQueSize(2);
	EXPECT_EQ(stage.queSize(), 2);
}


TEST(PIPipelineThread_Bounded, WaitUnblockedByClear) {
	Stage stage;
	stage.setMaxQueSize(1);
	stage.enqueue(1, false);
	ASSERT_EQ(stage.queSize(), 1);

	std::atomic_bool done(false);
	std::atomic_bool producer_overload(true);
	std::thread producer([&stage, &done, &producer_overload] {
		bool overload = true;
		stage.enqueue(2, true, &overload);
		producer_overload = overload;
		done              = true;
	});

	piMSleep(100);
	EXPECT_FALSE(done);
	EXPECT_EQ(stage.queSize(), 1); // 2 is still waiting for space

	stage.clear(); // frees space and wakes the producer

	producer.join();
	EXPECT_FALSE(producer_overload);
	EXPECT_EQ(stage.queSize(), 1); // only 2 remains in the queue
}


TEST(PIPipelineThread_Bounded, WaitUnblockedByStopCalc) {
	SlowStage stage;
	stage.setMaxQueSize(1);
	stage.start();

	// First item goes to the slow calc, then the second item fills the queue
	stage.enqueue(1, false);
	PITimeMeasurer tm;
	while (stage.queSize() > 0 && tm.elapsed_s() < 5.)
		piMSleep(1); // wait until the consumer takes item 1 into the slow calc
	stage.enqueue(2, false);
	ASSERT_EQ(stage.queSize(), 1);

	std::atomic_bool done(false);
	bool overload = false;
	std::thread producer([&stage, &done, &overload] {
		stage.enqueue(3, true, &overload);
		done = true;
	});

	piMSleep(100);
	EXPECT_FALSE(done); // blocked on the full queue

	stage.stopCalc(2000);

	// producer must be woken by stopCalc, not hang forever
	producer.join();
	EXPECT_TRUE(done);
	EXPECT_TRUE(overload); // item was dropped because the stage is stopping
	EXPECT_FALSE(stage.isRunning());
}


TEST(PIPipelineThread_Stop, FastStopOnEmptyQueue) {
	Stage stage;
	stage.start();
	piMSleep(50); // let the thread reach the wait

	PITimeMeasurer tm;
	stage.stopCalc(1000);
	EXPECT_LT(tm.elapsed_m(), 500.);
	EXPECT_FALSE(stage.isRunning());
}


TEST(PIPipelineThread_Stop, StopWithPendingItems) {
	Stage stage;
	stage.start();
	for (int i = 0; i < 10; ++i)
		stage.enqueue(i);
	piMSleep(50);

	PITimeMeasurer tm;
	stage.stopCalc(2000);
	EXPECT_LT(tm.elapsed_m(), 1500.);
	EXPECT_FALSE(stage.isRunning());
}