package nextflow.k8s

import groovy.transform.CompileDynamic
import groovy.transform.CompileStatic
import groovy.util.logging.Slf4j
import nextflow.k8s.client.K8sClient
import java.util.concurrent.LinkedBlockingQueue

@Slf4j
@CompileStatic
class K8sTaskScheduler implements Runnable {
    private final K8sExecutor executor
    private final K8sSchedulingStrategy strategy
    private final LinkedBlockingQueue<K8sSchedulingRequest> queue = new LinkedBlockingQueue<>()
    private String[] nodes

    private synchronized boolean shouldStop

    private HashMap<String, ArrayList<K8sTaskHandler>> occupancy;
    private HashMap<String, String> taskToNode

    K8sTaskScheduler(String[] nodes, K8sSchedulingStrategy strategy) {
        this.executor = executor
        this.strategy = strategy
        this.nodes = nodes
        this.occupancy = new HashMap<>()
        for (String node : nodes) {
            this.occupancy.put(node, new ArrayList<K8sTaskHandler>())
        }
        this.taskToNode = new HashMap<>()
    }

    /**
     * Adds a task to the queue of outstanding tasks
     * @param handler
     */
    void submit(K8sTaskHandler handler) {
        log.info "[K8s] received queued task ${handler.task.name}"
        queue.add(new K8sSchedulingRequest(handler))
        final pending = new ArrayList<K8sSchedulingRequest>(queue)
        if (strategy.scheduleImmediately(this, pending))
            drain()
    }

    /**
     * Notify the scheduler that a task has finished execution
     * @param handler
     */
    void taskFinished(K8sTaskHandler handler) {
        /* Remove from occupancy list */
        String nodeName = taskToNode.get(handler.task.name)
        if (nodeName == null) {
            log.error "[K8s] no node saved for task ${handler.task.name}"
            return
        }
        ArrayList<K8sTaskHandler> nodeTasks = occupancy.get(nodeName)
        if (nodeTasks == null) {
            log.error "[K8s] tried to remove task ${handler.task.name} from node ${nodeName} but no tasks are saved for that node"
            return
        }
        nodeTasks.remove(handler)
        log.info "[K8s] removed task ${handler.task.name} from node ${nodeName}"
    }

    protected synchronized void drain() {
        while( true ) {
            final pending = new ArrayList<K8sSchedulingRequest>(queue)
            final decision = strategy.schedule(this, pending)

            if ( !decision )
                return

            if ( !queue.remove(decision.request) )
                continue

            log.info "[K8s] launching queued task ${decision.request.task.name} on node: ${decision.nodeName}"
            decision.request.handler.submitNow(decision.nodeName)
            ArrayList nodeTasks = occupancy.get(decision.nodeName)
            assert nodeTasks != null
            nodeTasks.add(decision.request.handler)
            taskToNode.put(decision.request.task.name, decision.nodeName)
        }
    }

    /* Scheduling Strategy Interface */
    List<String> getFreeNodes() {
        def unoccupied = occupancy.findAll {it.value == null || it.value.size() == 0 }
        unoccupied.keySet().toList()
    }

    List<String> getNodes() {
        return nodes.toList()
    }

    /**
     * Run is the scheduler threads main function
     * */
    void run() {
        this.shouldStop = false
        def interval = this.executor.getK8sConfig().schedulerInterval
        while (!shouldStop) {
            sleep(interval.toMillis())
            drain()
        }
        log.info("[K8s] terminated scheduler loop")
    }

    /**
     * Stop terminates the scheduler thread
     */
    void stop() {
        log.info("[K8s] stopping scheduler loop")
        this.shouldStop = true
    }
}
