#include <core/scheduler.h>
#include <algorithm>

namespace extra2d {

namespace {

class IntervalTimer : public Timer {
public:
    IntervalTimer(Scheduler::Cb cb, float interval, uint32 repeat, float delay)
        : cb_(std::move(cb)) {
        interval_ = interval;
        repeat_ = repeat;
        delay_ = delay;
        useDelay_ = delay > 0.0f;
        runForever_ = repeat == 0;
        elapsed_ = useDelay_ ? 0.0f : -interval_;
    }

    void update(float dt) override {
        if (paused_ || done_) return;

        elapsed_ += dt;

        if (useDelay_) {
            if (elapsed_ < delay_) return;
            elapsed_ -= delay_;
            useDelay_ = false;
        }

        if (elapsed_ >= interval_) {
            trigger();
            elapsed_ -= interval_;
            if (!runForever_) {
                timesExecuted_++;
                if (timesExecuted_ >= repeat_) {
                    done_ = true;
                }
            }
        }
    }

    void trigger() override {
        if (cb_) cb_(elapsed_);
    }

private:
    Scheduler::Cb cb_;
};

class OnceTimer : public Timer {
public:
    OnceTimer(Scheduler::VoidCb cb, float delay)
        : cb_(std::move(cb)) {
        delay_ = delay;
        elapsed_ = 0.0f;
    }

    void update(float dt) override {
        if (paused_ || done_) return;

        elapsed_ += dt;
        if (elapsed_ >= delay_) {
            trigger();
            done_ = true;
        }
    }

    void trigger() override {
        if (cb_) cb_();
    }

private:
    Scheduler::VoidCb cb_;
};

}

Scheduler& Scheduler::inst() {
    static Scheduler instance;
    return instance;
}

TimerHdl Scheduler::scheduleUpdate(TimerTarget* target, int pri) {
    if (!target) return INVALID_HDL;

    UpdateEntry entry{target, pri, false, false};
    updates_.push_back(entry);
    updateIndex_.insert({target, updates_.size() - 1});

    return genHdl();
}

void Scheduler::unscheduleUpdate(TimerTarget* target) {
    if (!target) return;

    tbb::concurrent_hash_map<TimerTarget*, size_t>::accessor acc;
    if (updateIndex_.find(acc, target)) {
        size_t idx = acc->second;
        if (idx < updates_.size()) {
            updates_[idx].markedForDel = true;
        }
        updateIndex_.erase(acc);
    }
}

TimerHdl Scheduler::schedule(Cb cb, float interval, uint32 repeat, float delay) {
    if (!cb) return INVALID_HDL;

    auto timer = makePtr<IntervalTimer>(std::move(cb), interval, repeat, delay);
    TimerHdl hdl = genHdl();
    timer->hdl_ = hdl;

    timers_.insert({hdl, timer});
    return hdl;
}

TimerHdl Scheduler::scheduleOnce(VoidCb cb, float delay) {
    if (!cb) return INVALID_HDL;

    auto timer = makePtr<OnceTimer>(std::move(cb), delay);
    TimerHdl hdl = genHdl();
    timer->hdl_ = hdl;

    timers_.insert({hdl, timer});
    return hdl;
}

TimerHdl Scheduler::scheduleForever(Cb cb, float interval) {
    return schedule(std::move(cb), interval, 0, 0.0f);
}

void Scheduler::unschedule(TimerHdl hdl) {
    timers_.erase(hdl);
}

void Scheduler::unscheduleAll() {
    timers_.clear();
    updates_.clear();
    updateIndex_.clear();
}

void Scheduler::pause(TimerHdl hdl) {
    tbb::concurrent_hash_map<TimerHdl, Ptr<Timer>>::accessor acc;
    if (timers_.find(acc, hdl)) {
        acc->second->pause();
    }
}

void Scheduler::resume(TimerHdl hdl) {
    tbb::concurrent_hash_map<TimerHdl, Ptr<Timer>>::accessor acc;
    if (timers_.find(acc, hdl)) {
        acc->second->resume();
    }
}

void Scheduler::update(float dt) {
    float scaledDt = dt * timeScale_.load();

    locked_ = true;

    for (auto& entry : updates_) {
        if (!entry.markedForDel && !entry.paused && entry.target) {
            entry.target->update(scaledDt);
        }
    }

    updates_.erase(
        std::remove_if(updates_.begin(), updates_.end(),
            [](const UpdateEntry& e) { return e.markedForDel; }),
        updates_.end()
    );

    std::vector<TimerHdl> toRemove;
    for (auto it = timers_.begin(); it != timers_.end(); ++it) {
        auto& timer = it->second;
        timer->update(scaledDt);
        if (timer->isDone()) {
            toRemove.push_back(it->first);
        }
    }

    for (auto hdl : toRemove) {
        timers_.erase(hdl);
    }

    locked_ = false;
}

void Scheduler::updateParallel(float dt) {
    float scaledDt = dt * timeScale_.load();

    locked_ = true;

    tbb::parallel_for(tbb::blocked_range<size_t>(0, updates_.size()),
        [this, scaledDt](const tbb::blocked_range<size_t>& range) {
            for (size_t i = range.begin(); i != range.end(); ++i) {
                auto& entry = updates_[i];
                if (!entry.markedForDel && !entry.paused && entry.target) {
                    entry.target->update(scaledDt);
                }
            }
        });

    updates_.erase(
        std::remove_if(updates_.begin(), updates_.end(),
            [](const UpdateEntry& e) { return e.markedForDel; }),
        updates_.end()
    );

    std::vector<TimerHdl> toRemove;
    for (auto it = timers_.begin(); it != timers_.end(); ++it) {
        auto& timer = it->second;
        timer->update(scaledDt);
        if (timer->isDone()) {
            toRemove.push_back(it->first);
        }
    }

    for (auto hdl : toRemove) {
        timers_.erase(hdl);
    }

    locked_ = false;
}

bool Scheduler::isScheduled(TimerHdl hdl) const {
    tbb::concurrent_hash_map<TimerHdl, Ptr<Timer>>::const_accessor acc;
    return timers_.find(acc, hdl);
}

size_t Scheduler::count() const {
    return timers_.size();
}

TimerHdl Scheduler::genHdl() {
    return nextHdl_.fetch_add(1, std::memory_order_relaxed);
}

} // namespace extra2d