package com.ebay.myriad.scheduler;

import com.ebay.myriad.configuration.MyriadConfiguration;
import com.ebay.myriad.configuration.NodeManagerConfiguration;
import com.ebay.myriad.executor.NMTaskConfig;
import com.ebay.myriad.state.NodeTask;
import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
import com.google.gson.Gson;
import com.google.protobuf.ByteString;
import org.apache.mesos.Protos;
import org.apache.mesos.Protos.*;
import org.apache.mesos.Protos.CommandInfo.URI;
import org.apache.mesos.Protos.Value.Scalar;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.inject.Inject;
import java.util.Objects;

public interface TaskFactory {
    TaskInfo createTask(Offer offer, TaskID taskId, NodeTask nodeTask);

    class NMTaskFactoryImpl implements TaskFactory {
        public static final String EXECUTOR_NAME = "myriad_task";
        public static final String EXECUTOR_PREFIX = "myriad_executor";

        private static final Logger LOGGER = LoggerFactory
                .getLogger(NMTaskFactoryImpl.class);
        private MyriadConfiguration cfg;
        private TaskUtils taskUtils;

        @Inject
        public NMTaskFactoryImpl(MyriadConfiguration cfg, TaskUtils taskUtils) {
            this.cfg = cfg;
            this.taskUtils = taskUtils;
        }

        private static String getFileName(String uri) {
            int lastSlash = uri.lastIndexOf('/');
            if (lastSlash == -1) {
                return uri;
            } else {
                String fileName = uri.substring(lastSlash + 1);
                Preconditions.checkArgument(!Strings.isNullOrEmpty(fileName),
                        "URI should not have a slash at the end");
                return fileName;
            }
        }

        @Override
        public TaskInfo createTask(Offer offer, TaskID taskId, NodeTask nodeTask) {
            Objects.requireNonNull(offer, "Offer should be non-null");
            Objects.requireNonNull(nodeTask, "NodeTask should be non-null");

            NMProfile profile = nodeTask.getProfile();
            NMTaskConfig nmTaskConfig = new NMTaskConfig();
            nmTaskConfig.setAdvertisableCpus(profile.getCpus());
            nmTaskConfig.setAdvertisableMem(profile.getMemory());
            NodeManagerConfiguration nodeManagerConfiguration = this.cfg.getNodeManagerConfiguration();
            nmTaskConfig.setUser(nodeManagerConfiguration
                    .getUser().orNull());
            nmTaskConfig.setJvmOpts(nodeManagerConfiguration
                    .getJvmOpts().orNull());
            nmTaskConfig.setCgroups(nodeManagerConfiguration.getCgroups().or(Boolean.FALSE));
            nmTaskConfig.setYarnEnvironment(cfg.getYarnEnvironment());

            String taskConfigJSON = new Gson().toJson(nmTaskConfig);

            Scalar taskMemory = Value.Scalar.newBuilder()
                    .setValue(taskUtils.getTaskMemory(profile)).build();
            Scalar taskCpus = Value.Scalar.newBuilder()
                    .setValue(taskUtils.getTaskCpus(profile)).build();
            Scalar executorMemory = Value.Scalar.newBuilder()
                    .setValue(taskUtils.getExecutorMemory()).build();
            Scalar executorCpus = Value.Scalar.newBuilder()
                    .setValue(taskUtils.getExecutorCpus()).build();

            String executorPath = cfg.getMyriadExecutorConfiguration()
                    .getPath();
            URI executorURI = URI.newBuilder().setValue(executorPath)
                    .setExecutable(true).build();

            CommandInfo commandInfo = CommandInfo.newBuilder()
                    .addUris(executorURI).setUser("root")
                    .setValue("export CAPSULE_CACHE_DIR=`pwd`;echo $CAPSULE_CACHE_DIR; java -Dcapsule.log=verbose -jar " + getFileName(executorPath))
                    .build();

            ExecutorID executorId = Protos.ExecutorID.newBuilder()
                    .setValue(EXECUTOR_PREFIX + offer.getSlaveId().getValue())
                    .build();
            ExecutorInfo executorInfo = ExecutorInfo
                    .newBuilder()
                    .setCommand(commandInfo)
                    .setName(EXECUTOR_NAME)
                    .addResources(
                            Resource.newBuilder().setName("cpus")
                                    .setType(Value.Type.SCALAR)
                                    .setScalar(executorCpus).build())
                    .addResources(
                            Resource.newBuilder().setName("mem")
                                    .setType(Value.Type.SCALAR)
                                    .setScalar(executorMemory).build())
                    .setExecutorId(executorId).setCommand(commandInfo).build();

            TaskInfo.Builder taskBuilder = TaskInfo.newBuilder()
                    .setName("task-" + taskId.getValue()).setTaskId(taskId)
                    .setSlaveId(offer.getSlaveId());

            // TODO (mohit): Configure ports for multi-tenancy
            ByteString data = ByteString.copyFrom(taskConfigJSON.getBytes());
            return taskBuilder
                    .addResources(
                            Resource.newBuilder().setName("cpus")
                                    .setType(Value.Type.SCALAR)
                                    .setScalar(taskCpus).build())
                    .addResources(
                            Resource.newBuilder().setName("mem")
                                    .setType(Value.Type.SCALAR)
                                    .setScalar(taskMemory).build())
                    .setExecutor(executorInfo).setData(data).build();
        }
    }
}

@@ -953,16 +953,205 @@
 ecutor%22;
+%0A        public static final String YARN_NODEMANAGER_OPTS_KEY = %22YARN_NODEMANAGER_OPTS%22;%0A        private static final String YARN_RESOURCEMANAGER_HOSTNAME = %22yarn.resourcemanager.hostname%22;
 %0A%0A      
@@ -2937,24 +2937,1152 @@
 onment());%0A%0A
+            // if RM's hostname is passed in as a system property, pass it along%0A            // to Node Managers launched via Myriad%0A            String rmHostName = System.getProperty(YARN_RESOURCEMANAGER_HOSTNAME);%0A            if (rmHostName != null && !rmHostName.isEmpty()) %7B%0A%0A                String nmOpts = nmTaskConfig.getYarnEnvironment().get(YARN_NODEMANAGER_OPTS_KEY);%0A                if (nmOpts == null) %7B%0A                    nmOpts = %22%22;%0A                %7D%0A                nmOpts += %22 %22 + %22-D%22 + YARN_RESOURCEMANAGER_HOSTNAME + %22=%22 + rmHostName;%0A                nmTaskConfig.getYarnEnvironment().put(YARN_NODEMANAGER_OPTS_KEY, nmOpts);%0A                LOGGER.info(YARN_RESOURCEMANAGER_HOSTNAME + %22 is set to %22 + rmHostName +%0A                    %22 via YARN_RESOURCEMANAGER_OPTS. Passing it into YARN_NODEMANAGER_OPTS.%22);%0A            %7D%0A//            else %7B%0A                // TODO(Santosh): Handle this case. Couple of options:%0A                // 1. Lookup a hostname here and use it as %22RM's hostname%22%0A                // 2. Abort here.. RM cannot start unless a hostname is passed in as it requires it to pass to NMs.%0A%0A
             

Allow ResourceManager to be launched via Marathon and discoverable via Mesos DNS.