package net.engin33r.luaspigot.lua.lib;

import net.engin33r.luaspigot.lua.Library;
import net.engin33r.luaspigot.lua.annotation.LibFunctionDef;
import org.bukkit.Bukkit;
import org.bukkit.plugin.java.JavaPlugin;
import org.bukkit.scheduler.BukkitScheduler;
import org.bukkit.scheduler.BukkitTask;
import org.luaj.vm2.LuaFunction;
import org.luaj.vm2.LuaNumber;
import org.luaj.vm2.Varargs;

import static org.luaj.vm2.LuaValue.NIL;

/**
 * Library for creating and managing Spigot async tasks.
 */
public class TaskLibrary extends Library {
    private final JavaPlugin plugin;
    private final BukkitScheduler scheduler;

    public TaskLibrary(JavaPlugin plugin) {
        this.plugin = plugin;
        this.scheduler = Bukkit.getScheduler();
    }

    @Override
    public String getName() {
        return "tasks";
    }

    @LibFunctionDef(name = "sched", module = "async")
    public Varargs schedAsync(Varargs args) {
        LuaFunction run = args.checkfunction(1);
        BukkitTask task = scheduler.runTaskAsynchronously(plugin, run::call);
        return LuaNumber.valueOf(task.getTaskId());
    }

    @LibFunctionDef(name = "repeating", module = "async")
    public Varargs repeatingAsync(Varargs args) {
        LuaFunction run = args.checkfunction(1);
        BukkitTask task = scheduler.runTaskTimerAsynchronously(plugin,
                run::call, args.optlong(3, 0), args.checklong(2));
        return LuaNumber.valueOf(task.getTaskId());
    }

    @LibFunctionDef(name = "delayed", module = "async")
    public Varargs delayedAsync(Varargs args) {
        LuaFunction run = args.checkfunction(1);
        BukkitTask task = scheduler.runTaskLaterAsynchronously(plugin,
                run::call, args.checklong(2));
        return LuaNumber.valueOf(task.getTaskId());
    }

    @LibFunctionDef(name = "sched", module = "sync")
    public Varargs schedSync(Varargs args) {
        LuaFunction run = args.checkfunction(1);
        BukkitTask task = scheduler.runTask(plugin, run::call);
        return LuaNumber.valueOf(task.getTaskId());
    }

    @LibFunctionDef(name = "repeating", module = "sync")
    public Varargs repeatingSync(Varargs args) {
        LuaFunction run = args.checkfunction(1);
        BukkitTask task = scheduler.runTaskTimer(plugin, run::call,
                args.optlong(3, 0), args.checklong(2));
        return LuaNumber.valueOf(task.getTaskId());
    }

    @LibFunctionDef(name = "delayed", module = "sync")
    public Varargs delayedSync(Varargs args) {
        LuaFunction run = args.checkfunction(1);
        BukkitTask task = scheduler.runTaskLater(plugin, run::call,
                args.checklong(2));
        return LuaNumber.valueOf(task.getTaskId());
    }

    @LibFunctionDef(name = "cancel")
    public Varargs cancel(Varargs args) {
        scheduler.cancelTask(args.checkint(1));
        return NIL;
    }

    @LibFunctionDef(name = "cancelAll")
    public Varargs cancelAll(Varargs args) {
        scheduler.cancelAllTasks();
        return NIL;
    }
}

@@ -385,16 +385,65 @@
 rargs;%0A%0A
+import java.util.HashSet;%0Aimport java.util.Set;%0A%0A
 import s
@@ -665,16 +665,82 @@
 heduler;
+%0A    private final Set%3CBukkitTask%3E spawnedTasks = new HashSet%3C%3E();
 %0A%0A    pu
@@ -1156,32 +1156,64 @@
 in, run::call);%0A
+        spawnedTasks.add(task);%0A
         return L
@@ -1542,32 +1542,64 @@
 .checklong(2));%0A
+        spawnedTasks.add(task);%0A
         return L
@@ -1904,32 +1904,64 @@
 .checklong(2));%0A
+        spawnedTasks.add(task);%0A
         return L
@@ -2206,32 +2206,64 @@
 in, run::call);%0A
+        spawnedTasks.add(task);%0A
         return L
@@ -2576,32 +2576,64 @@
 .checklong(2));%0A
+        spawnedTasks.add(task);%0A
         return L
@@ -2922,32 +2922,64 @@
 .checklong(2));%0A
+        spawnedTasks.add(task);%0A
         return L
@@ -3112,214 +3112,552 @@
     
-scheduler.cancelTask(args.checkint(1));%0A        return NIL;%0A    %7D%0A%0A    @LibFunctionDef(name = %22cancelAll%22)%0A    public Varargs cancelAll(Varargs args) %7B%0A        scheduler.cancelAllTasks();%0A        return NIL
+int id = args.checkint(1);%0A        spawnedTasks.stream().filter(t -%3E t.getTaskId() == id)%0A                .forEach(BukkitTask::cancel);%0A        spawnedTasks.removeIf(t -%3E t.getTaskId() == id);%0A        scheduler.cancelTask(id);%0A        return NIL;%0A    %7D%0A%0A    @LibFunctionDef(name = %22cancelAll%22)%0A    public Varargs cancelAll(Varargs args) %7B%0A        scheduler.cancelAllTasks();%0A        spawnedTasks.clear();%0A        return NIL;%0A    %7D%0A%0A    public void stop() %7B%0A        spawnedTasks.stream().forEach(BukkitTask::cancel);%0A        spawnedTasks.clear()
 ;%0A  

Add a registry of all spawned tasks and manage them