package recreationaltech.plugin.strategies

import groovy.transform.CompileStatic
import groovy.util.logging.Slf4j
import recreationaltech.plugin.K8sDVFSClient
import recreationaltech.plugin.K8sRuntimeEstimator
import recreationaltech.plugin.K8sSchedulingDecision
import recreationaltech.plugin.K8sSchedulingRequest
import recreationaltech.plugin.K8sSchedulingStrategy
import recreationaltech.plugin.K8sTaskHandler
import recreationaltech.plugin.K8sTaskScheduler
import recreationaltech.plugin.client.K8sClient
import nextflow.processor.TaskRun
import nextflow.util.Duration

/**
 * Implements a scheduling strategy utilizing dvfs to reduce the energy consumption
 * of workflow execution, while attempting to maintain the same makespan.
 */
@Slf4j
@CompileStatic
class K8sDVFSSchedulingStrategy implements K8sSchedulingStrategy {

    /** Used for passing K8sExecutor.getClient.
     * We cannot pass the client directly, becaue it can be refreshed
     * during workflow execution.
     */
    public interface K8sClientGetter {
        K8sClient getClient()
    }

    private static long getTaskMemoryRequirment(TaskRun task) {
        return task.config.getMemory() ? task.config.getMemory().bytes : 64 * 1024 * 1024
    }

    private static int getTaskCPURequirement(TaskRun task) {
        return task.config.hasCpus() ? task.config.getCpus() * 1000 : 1000
    }

    @Slf4j
    private class WorkerNode {
        private class AssignedTask {
            TaskRun task
            long frequency

            AssignedTask(TaskRun t, long f) {
                this.task = t
                this.frequency = f
            }
        }

        String name

        long maxFrequency
        long minFrequency
        long currentFrequency

        ArrayList<AssignedTask> tasks

        WorkerNode(String name, long maxF, long minF, long curF, K8sClientGetter clientGetter) {
            this.name = name
            this.maxFrequency = maxF
            this.minFrequency = minF
            this.currentFrequency = curF
            this.tasks = new ArrayList<>()
        }

        /* Return the number of available (unoccupied) bytes */
        long getAvailableMemory() {
            def k8sClient = clientGetter.getClient()
            Long available = k8sClient.getNodeMemoryAvailableBytes(this.name)
            if (available == null) {
                log.warn "[K8s] failed to retrieve available memory for node ${name}"
                // Fallback: try capacity - allocated
                Long capacity = k8sClient.getNodeMemoryCapacityBytes(this.name)
                Long allocated = k8sClient.getNodeMemoryAllocatedBytes(this.name)
                if (capacity != null && allocated != null) {
                    return capacity - allocated
                }
                return 0
            }
            return available
        }

        /* Return the number of available (unoccupied) cpu cores */
        long getAvailableCPUs() {
            def k8sClient = clientGetter.getClient()
            Long available = k8sClient.getNodeCpuAvailableMillis(this.name)
            if (available == null) {
                log.warn "[K8s] failed to retrieve available CPU for node ${name}"
                // Fallback: try capacity - allocated
                Long capacity = k8sClient.getNodeCpuCapacityMillis(this.name)
                Long allocated = k8sClient.getNodeCpuAllocatedMillis(this.name)
                if (capacity != null && allocated != null) {
                    return capacity - allocated
                }
                return 0
            }
            return available
        }

        /* Return the total amount of installed memory */
        long getMemoryAmount() {
            def k8sClient = clientGetter.getClient()
            Long capacity = k8sClient.getNodeMemoryCapacityBytes(this.name)
            return capacity != null ? capacity.longValue() : 0
        }

        /* Return the total number of installed cpu cores */
        long getCPUCount() {
            def k8sClient = clientGetter.getClient()
            Long capacity = k8sClient.getNodeCpuCapacityMillis(this.name)
            return capacity != null ? capacity.longValue(): 0
        }

        // Sets the frequency to the max. requested frequency of all currently running tasks.
        private void updateFrequency(K8sDVFSClient dvfsClient) {
            if (tasks.size() == 0)
                return
            long max = Long.MIN_VALUE
            for (AssignedTask t : tasks) {
                max = Long.max(t.frequency, max)
            }
            log.info "[K8s] node ${name} running at ${max} Hz / ${maxFrequency} Hz ${((double)max/(double)maxFrequency) * 100.0}%"
            dvfsClient.setNodeFrequency(name, (int)max)
            this.currentFrequency = max
        }

        void assignTask(TaskRun task, long frequency, K8sDVFSClient dvfsClient) {
            log.info "[K8s] node ${name}: task ${task.name} assigned with ${frequency}/${maxFrequency}"
            this.tasks.add(new AssignedTask(task, frequency))
            updateFrequency(dvfsClient)
        }

        // Returns the frequency that was assigned to the task
        long taskFinished(TaskRun task, K8sDVFSClient dvfsClient) {
            long f = Long.MAX_VALUE
            log.info "[K8s] node ${name}: task ${task.name} finished"
            AssignedTask t = this.tasks.find { it.task == task }
            if (t != null) {
                f = t.frequency
                this.tasks.remove(t)
                updateFrequency(dvfsClient)
            }
            return f
        }
    }

    class SchedulingRequestComparator implements Comparator<K8sSchedulingRequest> {
        K8sRuntimeEstimator runtimeEstimator
        long currentTime
        double epsilon

        @Override
        int compare(K8sSchedulingRequest o1, K8sSchedulingRequest o2) {
            // First, check if one of the tasks is (estimated to be) on the critical path
            double t1 = runtimeEstimator.estimate(o1.handler)
            double t2 = runtimeEstimator.estimate(o2.handler)

            if (t1 > t2 + epsilon)
                return -1
            else if (t2 > t1 + epsilon)
                return 1

            // Both are not on the critical path. Sort based on the time they spent in the queue
            long w1 = currentTime - o1.submitTimeMillis
            long w2 = currentTime - o2.submitTimeMillis
            if (w1 > w2)
                return -1
            else if (w2 > w1)
                return 1
            return 0
        }
    }

    private K8sRuntimeEstimator runtimeEstimator
    private K8sDVFSClient dvfsClient

    private ArrayList<WorkerNode> nodes
    private HashMap<String, WorkerNode> taskToNode
    private long globalMaxFrequency
    private long globalMinFrequency

    private K8sClientGetter clientGetter

    private double comparisonEpsilonMillis

    private double[] topRuntimes
    private double averageRuntime
    private double finishedTaskCount

    boolean fullSpeedMode

    K8sDVFSSchedulingStrategy(K8sRuntimeEstimator runtimeEstimator,
                              K8sDVFSClient dvfsClient,
                              K8sClientGetter clientGetter,
                              Duration runtimeComparisonEpsilon,
                              int topRuntimeCount) {
        this.runtimeEstimator = runtimeEstimator
        this.dvfsClient = dvfsClient
        this.nodes = new ArrayList<>()
        this.taskToNode = new HashMap<>();
        this.clientGetter = clientGetter
        this.comparisonEpsilonMillis = (double)runtimeComparisonEpsilon.toMillis()
        this.topRuntimes = new double[topRuntimeCount]
        for (int i = 0; i < topRuntimeCount; i++) {
            this.topRuntimes[i] = 0.0
        }
        this.averageRuntime = 0.0
        this.finishedTaskCount = 0.0
    }

    private boolean isInTopRuntimes(double rt) {
        for (int i = 0; i < topRuntimes.size(); i++) {
            if (rt >= topRuntimes[i])
                return true
        }
        return false
    }

    private void updateTopRuntimes(double rt) {
        for (int i = 0; i < topRuntimes.size(); i++) {
            if (rt > topRuntimes[i]) {
                /* Move all one down */
                for (int j = topRuntimes.size() - 1; j > i; j--) {
                    topRuntimes[j] = topRuntimes[j - 1];
                }
                topRuntimes[i] = rt
                break
            }
        }
    }

    @Override
    K8sSchedulingDecision schedule(K8sTaskScheduler scheduler, List<K8sSchedulingRequest> queue) {
        if (nodes.isEmpty()) {
            if (!initNodes(scheduler))
                return null
        }

        /* Step 1: Sort by task priority. We will attempt to schedule tasks "in order", so that the
         * highest priority tasks are assigned to nodes as soon as possible.
         *
         * Priority is based on a) the tasks estimated runtime and b) the wait time of the task.
         */
        SchedulingRequestComparator comparator = new SchedulingRequestComparator()
        comparator.runtimeEstimator = runtimeEstimator
        comparator.epsilon = comparisonEpsilonMillis
        comparator.currentTime = System.currentTimeMillis()
        queue.sort(comparator)

        /* Step 2: For each task attempt to schedule it onto a node */
        for (K8sSchedulingRequest req : queue) {
            /* Step 2.1: Determine if the task is on the critical path.
             * If yes, we just schedule it at max frequency on the node with the highest available
             * frequency. If not, we determine a frequency (see below).
             */
            final double taskEstimation = runtimeEstimator.estimate(req.handler)
            final boolean isCriticalPath = isInTopRuntimes(taskEstimation)
            long frequency = globalMaxFrequency
            if (!isCriticalPath && !fullSpeedMode) {
                /* Set frequency so that we expect the runtime to be close to the mean runtime. */
                frequency = (long)Math.floor((taskEstimation * globalMaxFrequency) / averageRuntime)
                frequency = Math.max(frequency, globalMinFrequency)
                frequency = Math.min(frequency, globalMaxFrequency)
            }

            /* Step 2.2: Filter nodes based on task requirements */
            ArrayList<WorkerNode> suitableNodes = filterNodes(req.task)
            if (suitableNodes.size() == 0) {
                if (!anyNode(req.task)) {
                    log.error "[K8s] unable to schedule task ${req.task} - no node satisfies resource requirements ${getTaskMemoryRequirment(req.task)} bytes ${getTaskCPURequirement(req.task)} cpus"
                    return null
                }
                /* No node can currently execute this task, but it should be possible in the future */
                log.info "[K8s] ${req.task} can not be scheduled: ${getTaskMemoryRequirment(req.task)} bytes ${getTaskCPURequirement(req.task)} CPUs"
                for (WorkerNode n : this.nodes) {
                    log.info "[K8s] node ${n.name} - ${n.availableMemory}, ${n.availableCPUs}"
                }
                continue
            }

            /* Step 2.3: Assign to node based on "best fit" - current node frequency is closest to determined frequency */
            long minDist = Math.abs(suitableNodes[0].currentFrequency - frequency)
            WorkerNode closest = suitableNodes[0]
            for (WorkerNode node : suitableNodes) {
                long dist = Math.abs(node.currentFrequency - frequency)
                if (dist < minDist) {
                    closest = node
                    minDist = dist
                }
            }

            closest.assignTask(req.task, frequency, dvfsClient)
            taskToNode.put(req.task.hash.toString(), closest)
            log.info "[K8s] DVFS: Assigned task ${req.task} to node ${closest.name} - ${taskToNode.size()} assigned tasks"
            return new K8sSchedulingDecision(req, closest.name)
        }

        log.info "[K8s] unable to schedule any task. The queue contains ${queue.size()} tasks."
        return null
    }

    @Override
    boolean scheduleImmediately(K8sTaskScheduler scheduler, List<K8sSchedulingRequest> queue) {
        if (nodes.isEmpty()) {
            if (!initNodes(scheduler))
                return false
        }
        /* We want to schedule immediately if there are unoccupied nodes */
        boolean doIt = queue != null && queue.size() > 0 && taskToNode.size() < nodes.size()
        log.info "[K8s] scheduleImmediately: ${queue.size()} tasks in queue, ${taskToNode.size()} tasks running on ${nodes.size()} nodes: ${doIt}"
        return doIt
    }

    @Override
    synchronized void taskFinished(K8sTaskHandler task) {
        long freq = globalMaxFrequency
        /* Free resources allocated by this task */
        WorkerNode node = taskToNode.get(task.task.hash.toString())
        if (node != null) {
            freq = node.taskFinished(task.task, dvfsClient)
            taskToNode.remove(task.task.hash.toString())
        } else {
            log.warn "[K8s] no node recorded for task ${task.toString()}"
        }

        log.info "[K8s] task ${task.toString()} finished - ${taskToNode.size()} tasks running"
        if (node != null) {
            log.info "[K8s] task ran on node ${node.name} - ${node.availableMemory} bytes ${node.availableCPUs} at ${freq}/${globalMaxFrequency} Hz (${(double)freq / (double)globalMaxFrequency}%)"
        }

        /* We scale the runtime by the tasks relative frequency to avoid skewing the average runtime towards
         * longer runtimes. This is obviously only a rough approximation.
         */
        double runtime = (double)(task.getCompleteTimeMillis() - task.getStartTimeMillis())
        runtime *= (double)freq / (double)globalMaxFrequency
        averageRuntime = (runtime + finishedTaskCount * averageRuntime) / (finishedTaskCount + 1.0)
        finishedTaskCount += 1.0
        updateTopRuntimes(runtime)
    }

    private synchronized boolean initNodes(K8sTaskScheduler scheduler) {
        this.globalMaxFrequency = Long.MAX_VALUE
        final nodes = scheduler.getNodes()
        for (String node : nodes) {
            final cur = dvfsClient.getNodeCurrentFrequency(node)
            final min = dvfsClient.getNodeMinFrequency(node)
            final max = dvfsClient.getNodeMaxFrequency(node)
            final cpus = dvfsClient.getCPUCount(node)
            final mem = dvfsClient.getMemoryAmount(node)

            if (cur.empty || min.empty || max.empty || cpus.empty || mem.empty) {
                log.error "[K8s] failed to query node $node information"
                continue
            }
            globalMaxFrequency = Long.min(globalMaxFrequency, max.asLong)
            globalMinFrequency = Long.max(globalMinFrequency, min.asLong)

            log.info "[K8s] node ${node}: ${cpus.asLong} CPUs, ${mem.asLong} bytes RAM ${min.asLong} Hz - ${max.asLong} Hz current ${cur.asLong}"
            this.nodes.add(new WorkerNode(node, max.asLong, min.asLong, cur.asLong, clientGetter))
        }
        return !this.nodes.isEmpty()
    }

    /* Returns a list of nodes that fulfill the tasks resource requirements
     */
    private ArrayList<WorkerNode> filterNodes(TaskRun task) {
        final long reqBytes = getTaskMemoryRequirment(task)
        final int reqCPUs = getTaskCPURequirement(task)
        ArrayList<WorkerNode> suitableNodes = new ArrayList<>()
        for (WorkerNode n : nodes) {
            if (n.availableMemory >= reqBytes && n.availableCPUs >= reqCPUs) {
                log.info "[K8s] task ${task.name}: ${reqBytes} bytes ${reqCPUs} cpus: node ${n.name} has ${n.availableMemory} bytes, ${n.availableCPUs} cpus"
                suitableNodes.add(n)
            }
        }
        return suitableNodes
    }

    private boolean anyNode(TaskRun task) {
        final long reqBytes = getTaskMemoryRequirment(task)
        final int reqCPUs = getTaskCPURequirement(task)
        for (WorkerNode n : nodes) {
            if (n.memoryAmount >= reqBytes && n.CPUCount >= reqCPUs)
                return true
        }
        return false
    }
}