#include "scheduler.h"

#include <utility>

bool Scheduler::Later::operator()(const ScheduledTask& left,
                                  const ScheduledTask& right) const
{
    if(left.due == right.due)
    {
        return left.service->config.id > right.service->config.id;
    }
    return left.due > right.due;
}

mw::E<Scheduler::Clock::time_point> Scheduler::nextDueTime(
    Clock::time_point previous, Clock::duration interval, Clock::time_point now)
{
    if(interval <= Clock::duration::zero() ||
       previous < Clock::time_point{} || now < previous)
    {
        return std::unexpected(mw::runtimeError("Invalid task schedule"));
    }
    const auto delay = interval - (now - previous) % interval;
    if(delay > Clock::time_point::max() - now)
    {
        return std::unexpected(mw::runtimeError("Task due time overflow"));
    }
    return now + delay;
}

Scheduler::Scheduler(DataSourceInterface& data_source,
                     ProbeFactory probe_factory)
    : data_source(data_source), probe_factory(std::move(probe_factory))
{}

mw::E<std::unique_ptr<Scheduler>> Scheduler::create(
    const Configuration& configuration, DataSourceInterface& data_source,
    ProbeFactory probe_factory)
{
    if(configuration.worker_count == 0 || !probe_factory)
    {
        return std::unexpected(mw::runtimeError(
            "Scheduler requires workers and a probe factory"));
    }
    auto scheduler = std::unique_ptr<Scheduler>(
        new Scheduler(data_source, std::move(probe_factory)));
    for(const auto& group : configuration.groups)
    {
        for(const auto& config : group.services)
        {
            if(config.id.empty() ||
               config.id.find('\0') != std::string::npos)
            {
                return std::unexpected(mw::runtimeError(
                    "Invalid service ID"));
            }
            if(scheduler->services.contains(config.id))
            {
                return std::unexpected(mw::runtimeError(
                    "Duplicate service ID: " + config.id));
            }
            const auto now = Clock::now();
            auto next = nextDueTime(now, config.interval, now);
            if(!next)
            {
                return std::unexpected(mw::runtimeError(
                    config.id + ": " + next.error().msg()));
            }
            auto service = std::make_unique<Service>();
            service->config = config;
            // Validate without running the probe or the injected factory.
            auto probe = service->createProbe();
            if(!probe)
            {
                return std::unexpected(mw::runtimeError(
                    config.id + ": " + probe.error().msg()));
            }
            scheduler->services.emplace(config.id, std::move(service));
        }
    }
    auto pool = ThreadPool::create(configuration.worker_count);
    if(!pool)
    {
        return std::unexpected(pool.error());
    }
    scheduler->pool = std::move(*pool);
    return scheduler;
}

const Service* Scheduler::findService(const std::string& service_id) const
{
    const auto found = services.find(service_id);
    return found == services.end() ? nullptr : found->second.get();
}

mw::E<void> Scheduler::run(std::stop_token stop_token)
{
    {
        std::lock_guard lock(mutex);
        if(started)
        {
            return std::unexpected(mw::runtimeError(
                "Scheduler has already started"));
        }
        started = true;
    }
    auto result = runLoop(stop_token);
    pool->waitIdle();
    std::lock_guard lock(mutex);
    if(result && error)
    {
        return std::unexpected(*error);
    }
    return result;
}

mw::E<void> Scheduler::runLoop(std::stop_token stop_token)
{
    const auto initial_due = Clock::now();
    for(const auto& entry : services)
    {
        queue.push({entry.second.get(), initial_due});
    }
    while(!stop_token.stop_requested())
    {
        std::unique_lock lock(mutex);
        if(error)
        {
            return std::unexpected(*error);
        }
        const auto now = Clock::now();
        if(!queue.empty() && queue.top().due <= now)
        {
            auto scheduled = queue.top();
            auto& service = *scheduled.service;
            auto next = nextDueTime(
                scheduled.due, service.config.interval, now);
            if(!next)
            {
                return std::unexpected(next.error());
            }
            queue.pop();
            scheduled.due = *next;
            queue.push(scheduled);
            lock.unlock();
            auto dispatched = dispatch(service);
            if(!dispatched)
            {
                return std::unexpected(dispatched.error());
            }
            continue;
        }
        if(!queue.empty())
        {
            const auto next = queue.top().due;
            wake.wait_until(lock, stop_token, next,
                            std::bind_front(&Scheduler::hasError, this));
        }
        else
        {
            wake.wait(lock, stop_token,
                      std::bind_front(&Scheduler::hasError, this));
        }
    }
    return {};
}

namespace
{

struct FlightGuard
{
    std::atomic<bool>& in_flight;
    bool active = true;

    ~FlightGuard()
    {
        if(active)
        {
            in_flight = false;
        }
    }
};

}

mw::E<void> Scheduler::dispatch(Service& service)
{
    if(service.in_flight.exchange(true))
    {
        return {};
    }
    FlightGuard guard{service.in_flight};
    auto submitted = pool->trySubmit(Task{*this, service, nullptr});
    if(!submitted)
    {
        return std::unexpected(submitted.error());
    }
    if(*submitted)
    {
        // The accepted task now clears the flag after probing and persistence.
        guard.active = false;
    }
    return {};
}

void Scheduler::Task::operator()()
{
    scheduler.execute(*this);
}

void Scheduler::execute(Task& task)
{
    FlightGuard guard{task.service.in_flight};
    const auto& service_id = task.service.config.id;
    auto probe = probe_factory(task.service);
    if(!probe)
    {
        reportError(service_id, std::move(probe.error()));
        return;
    }
    if(!*probe)
    {
        reportError(service_id, mw::runtimeError(
            "Probe factory returned a null probe"));
        return;
    }
    task.probe = std::move(*probe);
    auto result = task.probe->probe();
    if(!result)
    {
        reportError(service_id, std::move(result.error()));
        return;
    }
    auto saved = data_source.save({service_id, result->timestamp,
        result->duration_microsecond, result->status});
    if(!saved)
    {
        reportError(service_id, std::move(saved.error()));
    }
}

void Scheduler::reportError(const std::string& service_id, mw::Error failure)
{
    {
        std::lock_guard lock(mutex);
        if(!error)
        {
            error = mw::runtimeError(service_id + ": " + failure.msg());
        }
    }
    wake.notify_one();
}

bool Scheduler::hasError() const
{
    return error.has_value();
}