#include "scheduler.h"
#include "data_source_sqlite.h"
#include "fake_probe.h"

#include <atomic>
#include <future>
#include <stdexcept>

#include <gtest/gtest.h>

namespace
{

using namespace std::chrono_literals;

TEST(Scheduler, NextDueTimePreservesCadence)
{
    using Time = Scheduler::Clock::time_point;
    auto next = Scheduler::nextDueTime(Time{10s}, 5s, Time{10s});
    ASSERT_TRUE(next);
    EXPECT_EQ(*next, Time{15s});
    auto late = Scheduler::nextDueTime(Time{10s}, 5s, Time{27s});
    ASSERT_TRUE(late);
    EXPECT_EQ(*late, Time{30s});
    auto exact = Scheduler::nextDueTime(Time{10s}, 5s, Time{30s});
    ASSERT_TRUE(exact);
    EXPECT_EQ(*exact, Time{35s});
}

TEST(Scheduler, NextDueTimeRejectsInvalidSchedule)
{
    using Time = Scheduler::Clock::time_point;
    EXPECT_FALSE(Scheduler::nextDueTime(Time{}, 0s, Time{}));
    EXPECT_FALSE(Scheduler::nextDueTime(Time{}, -1s, Time{}));
    EXPECT_FALSE(Scheduler::nextDueTime(Time{-1s}, 1s, Time{}));
    EXPECT_FALSE(Scheduler::nextDueTime(Time{1s}, 1s, Time{}));
}

TEST(Scheduler, NextDueTimeDetectsOverflow)
{
    using Time = Scheduler::Clock::time_point;
    EXPECT_FALSE(Scheduler::nextDueTime(Time::max(), 1s, Time::max()));
}

// Keep SQLite persistence real while observing worker completion in tests.
class WatchingSource final : public DataSourceInterface
{
public:
    enum class SaveMode { SUCCESS, ERROR, EXCEPTION };

    const Scheduler* scheduler = nullptr;
    SaveMode mode = SaveMode::SUCCESS;
    std::atomic<bool> flight_during_save{true};

    mw::E<void> save(const StatusRecord& record) override
    {
        if(scheduler)
        {
            const auto* service = scheduler->findService(record.service_id);
            if(!service || !service->in_flight.load())
            {
                flight_during_save = false;
            }
        }
        if(mode == SaveMode::ERROR)
        {
            return std::unexpected(mw::runtimeError("Save failed"));
        }
        if(mode == SaveMode::EXCEPTION)
        {
            throw std::runtime_error("Unexpected save exception");
        }
        auto result = database->save(record);
        if(result)
        {
            {
                std::lock_guard lock(mutex);
                ++counts[record.service_id];
            }
            changed.notify_all();
        }
        return result;
    }

    mw::E<std::optional<StatusRecord>> latest(
        const std::string& service_id) override
    {
        return database->latest(service_id);
    }

    mw::E<std::vector<StatusRecord>> history(
        const std::string& service_id, std::int64_t start,
        std::int64_t end) override
    {
        return database->history(service_id, start, end);
    }

    bool waitForRecords(const std::string& service_id, std::size_t count,
                        std::chrono::milliseconds timeout = 1s)
    {
        std::unique_lock lock(mutex);
        const auto deadline = std::chrono::steady_clock::now() + timeout;
        while(counts[service_id] < count)
        {
            if(changed.wait_until(lock, deadline) == std::cv_status::timeout)
            {
                return counts[service_id] >= count;
            }
        }
        return true;
    }

private:
    std::unique_ptr<DataSourceSqlite> database =
        DataSourceSqlite::open(":memory:").value();
    std::mutex mutex;
    std::condition_variable changed;
    std::unordered_map<std::string, std::size_t> counts;
};

ServiceConfig serviceConfig(const std::string& id,
                            Scheduler::Clock::duration interval = 1h)
{
    return {id, "Display " + id, "Description", TcpEndpoint{"localhost", 80},
            1s, interval, "https://example.invalid"};
}

Configuration oneService()
{
    return {1, ":memory:", {{"Group", {serviceConfig("service")}}}};
}

mw::E<std::unique_ptr<ProbeInterface>> makeFake(
    [[maybe_unused]] const Service& service)
{
    return std::make_unique<FakeProbe>(
        ProbeResult{ProbeStatus::BAD, 123, 456}, 0ms);
}

// Cancel before the async run's future is destroyed after a failed assertion.
struct StopOnExit
{
    std::stop_source& source;

    ~StopOnExit()
    {
        source.request_stop();
    }
};

TEST(Scheduler, BuildsServicesAcrossGroupsAndPersistsResults)
{
    WatchingSource source;
    Configuration configuration{2, ":memory:", {
        {"First", {serviceConfig("first")}},
        {"Second", {serviceConfig("second")}}
    }};
    auto created = Scheduler::create(configuration, source, makeFake);
    ASSERT_TRUE(created) << created.error().msg();
    auto scheduler = std::move(*created);
    source.scheduler = scheduler.get();
    ASSERT_NE(scheduler->findService("first"), nullptr);
    ASSERT_NE(scheduler->findService("second"), nullptr);
    EXPECT_EQ(scheduler->findService("missing"), nullptr);
    EXPECT_EQ(scheduler->findService("first")->config.name, "Display first");
    configuration.groups.clear();
    EXPECT_EQ(scheduler->findService("first")->config.id, "first");

    std::stop_source stop;
    auto running = std::async(std::launch::async, &Scheduler::run,
                              scheduler.get(), stop.get_token());
    StopOnExit cleanup{stop};
    ASSERT_TRUE(source.waitForRecords("first", 1));
    ASSERT_TRUE(source.waitForRecords("second", 1));
    stop.request_stop();
    ASSERT_EQ(running.wait_for(1s), std::future_status::ready);
    EXPECT_TRUE(running.get());
    for(const auto& id : {"first", "second"})
    {
        auto record = source.latest(id);
        ASSERT_TRUE(record);
        ASSERT_TRUE(*record);
        EXPECT_EQ((**record), (StatusRecord{id, 123, 456, ProbeStatus::BAD}));
        EXPECT_FALSE(scheduler->findService(id)->in_flight.load());
    }
    EXPECT_TRUE(source.flight_during_save.load());
    EXPECT_FALSE(scheduler->run(stop.get_token()));
}

TEST(Scheduler, CancellationWakesEmptyQueue)
{
    WatchingSource source;
    auto scheduler = Scheduler::create(
        Configuration{1, ":memory:", {}}, source);
    ASSERT_TRUE(scheduler);
    std::stop_source stop;
    auto running = std::async(std::launch::async, &Scheduler::run,
                              scheduler->get(), stop.get_token());
    StopOnExit cleanup{stop};
    EXPECT_EQ(running.wait_for(20ms), std::future_status::timeout);
    stop.request_stop();
    ASSERT_EQ(running.wait_for(1s), std::future_status::ready);
    EXPECT_TRUE(running.get());
}

mw::E<std::unique_ptr<ProbeInterface>> factoryError(
    [[maybe_unused]] const Service& service)
{
    return std::unexpected(mw::runtimeError("Factory failed"));
}

TEST(Scheduler, CancellationBeforeStart)
{
    WatchingSource source;
    auto scheduler = Scheduler::create(oneService(), source, factoryError);
    ASSERT_TRUE(scheduler);
    std::stop_source stop;
    stop.request_stop();
    EXPECT_TRUE((*scheduler)->run(stop.get_token()));
    EXPECT_FALSE((*scheduler)->findService("service")->in_flight.load());
    EXPECT_FALSE(source.latest("service").value());
}

TEST(Scheduler, CancellationWakesFutureDueTime)
{
    WatchingSource source;
    auto scheduler = Scheduler::create(oneService(), source, makeFake);
    ASSERT_TRUE(scheduler);
    std::stop_source stop;
    auto running = std::async(std::launch::async, &Scheduler::run,
                              scheduler->get(), stop.get_token());
    StopOnExit cleanup{stop};
    ASSERT_TRUE(source.waitForRecords("service", 1));
    EXPECT_EQ(running.wait_for(20ms), std::future_status::timeout);
    stop.request_stop();
    ASSERT_EQ(running.wait_for(1s), std::future_status::ready);
    EXPECT_TRUE(running.get());
}

class FailureProbe final : public ProbeInterface
{
public:
    explicit FailureProbe(bool throw_exception)
        : throw_exception(throw_exception)
    {}

    mw::E<ProbeResult> probe() override
    {
        if(throw_exception)
        {
            throw std::runtime_error("Unexpected probe exception");
        }
        return std::unexpected(mw::runtimeError("Probe failed"));
    }

private:
    bool throw_exception;
};

mw::E<std::unique_ptr<ProbeInterface>> failedProbe(
    [[maybe_unused]] const Service& service)
{
    return std::make_unique<FailureProbe>(false);
}

mw::E<std::unique_ptr<ProbeInterface>> throwingProbe(
    [[maybe_unused]] const Service& service)
{
    return std::make_unique<FailureProbe>(true);
}

mw::E<std::unique_ptr<ProbeInterface>> nullProbe(
    [[maybe_unused]] const Service& service)
{
    return std::unique_ptr<ProbeInterface>{};
}

mw::E<std::unique_ptr<ProbeInterface>> throwingFactory(
    [[maybe_unused]] const Service& service)
{
    throw std::runtime_error("Unexpected factory exception");
}

TEST(Scheduler, FailuresClearServiceFlightState)
{
    for(auto factory : {factoryError, failedProbe, nullProbe})
    {
        WatchingSource source;
        auto scheduler = Scheduler::create(oneService(), source, factory);
        ASSERT_TRUE(scheduler);
        std::stop_source stop;
        auto running = std::async(std::launch::async, &Scheduler::run,
                                  scheduler->get(), stop.get_token());
        StopOnExit cleanup{stop};
        ASSERT_EQ(running.wait_for(1s), std::future_status::ready);
        auto result = running.get();
        ASSERT_FALSE(result);
        EXPECT_NE(result.error().msg().find("service:"), std::string::npos);
        EXPECT_FALSE((*scheduler)->findService("service")->in_flight.load());
        EXPECT_FALSE(source.latest("service").value());
    }
}

TEST(Scheduler, PersistenceFailuresClearServiceFlightState)
{
    WatchingSource source;
    source.mode = WatchingSource::SaveMode::ERROR;
    auto scheduler = Scheduler::create(oneService(), source, makeFake);
    ASSERT_TRUE(scheduler);
    source.scheduler = scheduler->get();
    std::stop_source stop;
    auto running = std::async(std::launch::async, &Scheduler::run,
                              scheduler->get(), stop.get_token());
    StopOnExit cleanup{stop};
    ASSERT_EQ(running.wait_for(1s), std::future_status::ready);
    auto result = running.get();
    ASSERT_FALSE(result);
    EXPECT_NE(result.error().msg().find("service:"), std::string::npos);
    EXPECT_FALSE((*scheduler)->findService("service")->in_flight.load());
    EXPECT_TRUE(source.flight_during_save.load());
}

class BlockingProbe final : public ProbeInterface
{
public:
    BlockingProbe(std::shared_future<void> release, std::promise<void>& started)
        : release(std::move(release)), started(started)
    {}

    mw::E<ProbeResult> probe() override
    {
        started.set_value();
        release.wait();
        return ProbeResult{ProbeStatus::GOOD, 123, 456};
    }

private:
    std::shared_future<void> release;
    std::promise<void>& started;
};

struct BlockingFactory
{
    std::shared_future<void> release;
    std::promise<void>& started;
    std::atomic<int>& calls;

    mw::E<std::unique_ptr<ProbeInterface>> operator()(
        const Service& service) const
    {
        if(service.config.id == "a_slow")
        {
            ++calls;
            return std::make_unique<BlockingProbe>(release, started);
        }
        return makeFake(service);
    }
};

TEST(Scheduler, SkipsInFlightAndFinishesAcceptedWorkOnStop)
{
    WatchingSource source;
    std::atomic<int> calls{0};
    std::promise<void> started;
    auto ready = started.get_future();
    std::stop_source stop;
    auto release = std::make_unique<std::promise<void>>();
    Configuration configuration{2, ":memory:", {{"Group", {
        serviceConfig("a_slow", 1ms), serviceConfig("ticker", 1ms)
    }}}};
    auto scheduler = Scheduler::create(configuration, source, BlockingFactory{
        release->get_future().share(), started, calls});
    ASSERT_TRUE(scheduler);
    source.scheduler = scheduler->get();
    auto running = std::async(std::launch::async, &Scheduler::run,
                              scheduler->get(), stop.get_token());
    // Unblock the worker before destroying running on assertion failure.
    auto release_on_exit = std::move(release);
    StopOnExit cleanup{stop};
    ASSERT_EQ(ready.wait_for(1s), std::future_status::ready);
    ASSERT_TRUE(source.waitForRecords("ticker", 4));
    EXPECT_EQ(calls.load(), 1);
    EXPECT_TRUE((*scheduler)->findService("a_slow")->in_flight.load());
    stop.request_stop();
    EXPECT_EQ(running.wait_for(20ms), std::future_status::timeout);
    release_on_exit->set_value();
    ASSERT_EQ(running.wait_for(1s), std::future_status::ready);
    EXPECT_TRUE(running.get());
    EXPECT_FALSE((*scheduler)->findService("a_slow")->in_flight.load());
    EXPECT_FALSE((*scheduler)->findService("ticker")->in_flight.load());
    EXPECT_TRUE(source.flight_during_save.load());
    EXPECT_TRUE(source.latest("a_slow").value());
}

TEST(Scheduler, SkipsWhenFullAndRetriesAtNextInterval)
{
    WatchingSource source;
    std::atomic<int> calls{0};
    std::promise<void> started;
    auto ready = started.get_future();
    std::stop_source stop;
    auto release = std::make_unique<std::promise<void>>();
    Configuration configuration{1, ":memory:", {{"Group", {
        serviceConfig("a_slow"), serviceConfig("retry", 1ms)
    }}}};
    auto scheduler = Scheduler::create(configuration, source, BlockingFactory{
        release->get_future().share(), started, calls});
    ASSERT_TRUE(scheduler);
    auto running = std::async(std::launch::async, &Scheduler::run,
                              scheduler->get(), stop.get_token());
    auto release_on_exit = std::move(release);
    StopOnExit cleanup{stop};
    ASSERT_EQ(ready.wait_for(1s), std::future_status::ready);
    EXPECT_FALSE(source.waitForRecords("retry", 1, 20ms));
    release_on_exit->set_value();
    ASSERT_TRUE(source.waitForRecords("retry", 1));
    stop.request_stop();
    ASSERT_EQ(running.wait_for(1s), std::future_status::ready);
    EXPECT_TRUE(running.get());
    EXPECT_FALSE((*scheduler)->findService("retry")->in_flight.load());
    EXPECT_EQ(calls.load(), 1);
}

TEST(Scheduler, InvalidConfiguration)
{
    WatchingSource source;
    auto configuration = oneService();
    configuration.worker_count = 0;
    EXPECT_FALSE(Scheduler::create(configuration, source));
    configuration = oneService();
    EXPECT_FALSE(Scheduler::create(configuration, source, {}));
    configuration.groups.front().services.front().id.clear();
    EXPECT_FALSE(Scheduler::create(configuration, source));
    configuration = oneService();
    configuration.groups.front().services.front().interval = 0s;
    EXPECT_FALSE(Scheduler::create(configuration, source));
    configuration.groups.front().services.front().interval = -1s;
    EXPECT_FALSE(Scheduler::create(configuration, source));
    configuration.groups.front().services.front().interval =
        Scheduler::Clock::duration::max();
    EXPECT_FALSE(Scheduler::create(configuration, source));
    configuration = oneService();
    configuration.groups.front().services.front().timeout = 0s;
    EXPECT_FALSE(Scheduler::create(configuration, source));
    configuration = oneService();
    configuration.groups.front().services.front().endpoint =
        TcpEndpoint{"localhost", 0};
    EXPECT_FALSE(Scheduler::create(configuration, source));
    configuration = oneService();
    configuration.groups.push_back(
        {"Another group", {serviceConfig("service")}});
    auto duplicate = Scheduler::create(configuration, source);
    ASSERT_FALSE(duplicate);
    EXPECT_NE(duplicate.error().msg().find("Duplicate service ID"),
              std::string::npos);
}

}

TEST(SchedulerDeathTest, UnexpectedProbeExceptionsTerminate)
{
    for(auto factory : {throwingProbe, throwingFactory})
    {
        ASSERT_DEATH({
            WatchingSource source;
            auto scheduler = Scheduler::create(oneService(), source, factory);
            (*scheduler)->run(std::stop_token{});
        }, "");
    }
}

TEST(SchedulerDeathTest, UnexpectedSaveExceptionTerminates)
{
    ASSERT_DEATH({
        WatchingSource source;
        source.mode = WatchingSource::SaveMode::EXCEPTION;
        auto scheduler = Scheduler::create(oneService(), source, makeFake);
        source.scheduler = scheduler->get();
        (*scheduler)->run(std::stop_token{});
    }, "");
}