package nextflow.k8s

import groovy.util.logging.Slf4j
import nextflow.k8s.client.K8sClient
import nextflow.k8s.model.PodHostMount
import nextflow.k8s.model.PodSpecBuilder

@Slf4j
class K8sNodeInitDeployer {
    private K8sClient client
    private K8sConfig config
    private String runName

    K8sNodeInitDeployer(K8sClient client, K8sConfig config, String runName) {
        this.client = client
        this.config = config
        this.runName = runName
    }

    void deploy() {
        final init = config.nodeInit
        if ( !init?.enabled )
            return

        log.info("deploying init pods")

        final nodes = getNodes()
        for ( String nodeName : nodes ) {
            log.info(" ... deploying to " + nodeName)
            final spec = makePodSpec(init, nodeName)
            client.podCreate(spec)
        }

        log.info("waiting for init pods")
        waitForPods(nodes)
    }

    void cleanup() {
        final init = config.nodeInit
        if ( !init?.enabled || !init?.cleanup )
            return

        final nodes = getNodes()
        for ( String nodeName : nodes ) {
            final podName = buildPodName(nodeName)
            client.podDelete(podName)
        }
    }

    String buildCleanupScript() {
        if ( !config.nodeInit?.enabled || !config.nodeInit?.cleanup )
            return "#!/usr/bin/env/bash\nexit 0"

        final nodes = getNodes()
        final podNames = nodes.collect {buildPodName(it)}

        String script = '''
#!/usr/bin/env bash
set -u

PODS=('''
        for ( String podName : podNames ) {
            script += "\"${podName}\"\n"
        }
        script += '''
)

SERVICE_ACCOUNT_DIR="/var/run/secrets/kubernetes.io/serviceaccount"
TOKEN="$(cat "${SERVICE_ACCOUNT_DIR}/token")"
NAMESPACE="$(cat "${SERVICE_ACCOUNT_DIR}/namespace")"
CA_CERT="${SERVICE_ACCOUNT_DIR}/ca.crt"

KUBE_API="https://${KUBERNETES_SERVICE_HOST}:${KUBERNETES_SERVICE_PORT_HTTPS:-443}"

for POD in "${PODS[@]}"; do
  echo "Deleting pod: ${POD}"

  HTTP_CODE="$(
    curl \\
      --silent \\
      --show-error \\
      --output /tmp/delete-pod-response.json \\
      --write-out "%{http_code}" \\
      --request DELETE \\
      --cacert "${CA_CERT}" \\
      --header "Authorization: Bearer ${TOKEN}" \\
      --header "Accept: application/json" \\
      "${KUBE_API}/api/v1/namespaces/${NAMESPACE}/pods/${POD}"
  )"

  case "${HTTP_CODE}" in
    200|202)
      echo "Deleted pod: ${POD}"
      ;;
    404)
      echo "Pod already absent: ${POD}"
      ;;
    *)
      echo "Failed to delete pod: ${POD}; HTTP ${HTTP_CODE}" >&2
      cat /tmp/delete-pod-response.json >&2 || true
      ;;
  esac
done

exit 0
        '''
        return script
    }

    private List<String> getNodes() {
        final resp = client.nodeList()
        ArrayList<String> nodes = new ArrayList<String>()
        for ( Map item : resp.items ) {
            nodes.add(item.metadata.name as String)
        }
        return nodes
    }

    private String buildPodName(String nodeName) {
        // TODO: Remove anything that is not lowercase alpha-numeric or dash
        String sanitizedNodeName = nodeName.toLowerCase()
        String name = "nf-init-${runName}-${sanitizedNodeName}"
        if ( name.length() > 63 )
            name = name.substring(0, 63)
        return name
    }

    private Map makePodSpec(K8sNodeInitConfig config, String nodeName) {
        ArrayList<PodHostMount> mounts = new ArrayList<PodHostMount>()
        mounts.add(new PodHostMount("/sys", "/sys"))
        mounts.add(new PodHostMount("/dev", "/dev"))
        mounts.add(new PodHostMount("/lib/modules", "/lib/modules"))

        PodSpecBuilder builder = new PodSpecBuilder()
        return builder.withNodeName(nodeName)
            .withImageName(config.image)
            .withCommand(config.command)
            .withPrivileged(true)
            .withHostMounts(mounts)
            .withPodName(buildPodName(nodeName))
            .build()
    }

    private void waitForPods(List<String> nodes) {
        if ( config.nodeInit.wait == 'Running' ) {
            for (String nodeName : nodes) {
                String podName = buildPodName(nodeName)
                while (true) {
                    sleep 1000
                    final state = client.podState(podName)
                    if (state && !state.containsKey('waiting')) {
                        break
                    }
                }
            }
        } else if ( config.nodeInit.wait == 'Succeeded' ) {
            for ( String nodeName : nodes ) {
                final String podName = buildPodName(nodeName)
                final currentState = client.podState(podName)
                if ( currentState && currentState?.running instanceof Map ) {
                    try {
                        while (true) {
                            sleep 10000
                            final state = client.podState(podName)
                            if (state && !(state?.running instanceof Map)) {
                                println "$podName has changed from running state $state"
                                break
                            }
                        }
                    }
                    catch (Exception e) {
                        println "Caught exception while waiting for ${podName} to stop running"
                    }
                }
            }
        }
    }
}
