#!/usr/bin/env python3
"""
DMLC submission script by mesos

One need to make sure all slaves machines are ssh-able.
"""
from __future__ import absolute_import

import json
import logging
import os
import uuid
from threading import Thread

from . import tracker

try:
    import pymesos.subprocess

    logging.getLogger("pymesos").setLevel(logging.WARNING)

    def _run(prog, env, resources):
        cwd = os.getcwd()
        pymesos.subprocess.check_call(
            prog,
            shell=True,
            env=env,
            cwd=cwd,
            cpus=resources["cpus"],
            mem=resources["mem"],
        )

    _USE_PYMESOS = True

except ImportError:
    import subprocess

    DEVNULL = open(os.devnull, "w")

    def _run(prog, env, resources):
        master = os.environ["MESOS_MASTER"]
        if ":" not in master:
            master += ":5050"

        name = str(uuid.uuid4())
        cwd = os.getcwd()
        prog = "cd %s && %s" % (cwd, prog)

        resources = ";".join("%s:%s" % (k, v) for k, v in resources.items())
        prog = prog.replace("'", "\\'")
        env = json.dumps(env).replace("'", "\\'")
        resources = resources.replace("'", "\\'")
        cmd = (
            "mesos-execute --master=%s --name='%s'"
            " --command='%s' --env='%s' --resources='%s'"
            % (master, name, prog, env, resources)
        )

        subprocess.check_call(cmd, shell=True, stdout=DEVNULL, stderr=subprocess.STDOUT)

    _USE_PYMESOS = False


def get_env():
    # get system envs
    keys = set(["OMP_NUM_THREADS", "KMP_AFFINITY", "LD_LIBRARY_PATH"])
    return {k: v for k, v in os.environ.items() if k in keys}


def submit(args):
    def mesos_submit(nworker, nserver, pass_envs):
        """
        customized submit script
        """
        # launch jobs
        for i in range(nworker + nserver):
            resources = {}
            pass_envs["DMLC_ROLE"] = "server" if i < nserver else "worker"
            if i < nserver:
                pass_envs["DMLC_SERVER_ID"] = i
                resources["cpus"] = args.server_cores
                resources["mem"] = args.server_memory_mb
            else:
                pass_envs["DMLC_WORKER_ID"] = i - nserver
                resources["cpus"] = args.worker_cores
                resources["mem"] = args.worker_memory_mb

            env = {str(k): str(v) for k, v in pass_envs.items()}
            env.update(get_env())
            prog = " ".join(args.command)
            thread = Thread(target=_run, args=(prog, env, resources))
            thread.setDaemon(True)
            thread.start()

        return mesos_submit

    if not _USE_PYMESOS:
        logging.warning(
            "No PyMesos found, use mesos-execute instead," " no task output available"
        )

    if args.mesos_master:
        os.environ["MESOS_MASTER"] = args.mesos_master

    assert "MESOS_MASTER" in os.environ, "No mesos master configured!"

    tracker.submit(
        args.num_workers,
        args.num_servers,
        fun_submit=mesos_submit,
        pscmd=(" ".join(args.command)),
    )
