package org.apache.mesos.elasticsearch.scheduler;

import com.google.protobuf.ByteString;
import org.apache.mesos.Protos;
import org.apache.mesos.elasticsearch.common.Discovery;
import org.apache.mesos.elasticsearch.scheduler.state.ClusterState;
import org.apache.mesos.elasticsearch.scheduler.state.FrameworkState;
import org.apache.mesos.elasticsearch.scheduler.util.Clock;
import org.joda.time.DateTime;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.runners.MockitoJUnitRunner;

import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.time.ZoneId;
import java.time.ZoneOffset;
import java.time.ZonedDateTime;
import java.util.*;
import java.util.stream.Collectors;

import static java.util.Arrays.asList;
import static org.junit.Assert.*;
import static org.mockito.Mockito.*;

/**
 * Tests TaskInfoFactory
 */
@RunWith(MockitoJUnitRunner.class)
@SuppressWarnings({"PMD.TooManyMethods"})
public class TaskInfoFactoryTest {

    private static final double EPSILON = 0.0001;
    public static final String SLAVEID = "SLAVEID";

    @Mock
    private FrameworkState frameworkState;

    @Mock
    private ClusterState clusterState;

    @Mock
    private Configuration configuration;

    @Mock
    private Clock clock;

    @Before
    public void before() {
        Protos.FrameworkID frameworkId = Protos.FrameworkID.newBuilder().setValue(UUID.randomUUID().toString()).build();
        when(frameworkState.getFrameworkID()).thenReturn(frameworkId);
        when(configuration.getTaskName()).thenReturn(Configuration.EXECUTOR_NAME);
        when(configuration.getMesosZKURL()).thenReturn("zk://zookeeper:2181/mesos");
        when(configuration.getExecutorImage()).thenReturn(Configuration.DEFAULT_EXECUTOR_IMAGE);
        when(configuration.getElasticsearchSettingsLocation()).thenReturn(Configuration.HOST_PATH_CONF);
        when(configuration.getElasticsearchNodes()).thenReturn(3);
        when(configuration.getElasticsearchClusterName()).thenReturn("cluster-name");
        when(configuration.getDataDir()).thenReturn(Configuration.DEFAULT_HOST_DATA_DIR);
        when(configuration.getFrameworkRole()).thenReturn("some-framework-role");
        when(configuration.isFrameworkUseDocker()).thenReturn(true);
        when(configuration.getElasticsearchPorts()).thenReturn(Collections.emptyList());
        when(configuration.taskSpecificHostDir(any())).thenReturn("/var/lib/mesos/slave/elasticsearch/cluster-name/" + SLAVEID);
        when(configuration.getElasticsearchBinary()).thenReturn("");
    }

    @Test
    public void testCreateTaskInfo() {
        TaskInfoFactory factory = new TaskInfoFactory(clusterState);

        Date now = new DateTime().withDayOfMonth(1).withDayOfYear(1).withYear(1970).withHourOfDay(1).withMinuteOfHour(2).withSecondOfMinute(3).withMillisOfSecond(400).toDate();
        when(clock.now()).thenReturn(now);
        when(clock.nowUTC()).thenReturn(ZonedDateTime.now(ZoneOffset.UTC));


        Protos.Offer offer = getOffer(frameworkState.getFrameworkID());

        Protos.TaskInfo taskInfo = factory.createTask(configuration, frameworkState, offer, clock);

        assertEquals(configuration.getTaskName(), taskInfo.getName());
        assertEquals(offer.getSlaveId(), taskInfo.getSlaveId());
        assertEquals("elasticsearch_localhost_19700101T010203.400Z", taskInfo.getTaskId().getValue());

        List<Protos.Resource> resourceList = taskInfo.getResourcesList();
        assertEquals(configuration.getCpus(), getResourceByName(resourceList, "cpus").getScalar().getValue(), EPSILON);

        assertEquals(configuration.getDisk(), getResourceByName(resourceList, "disk").getScalar().getValue(), EPSILON);

        assertEquals(configuration.getMem(), getResourceByName(resourceList, "mem").getScalar().getValue(), EPSILON);

        List<Protos.Resource> portsList = getResourceListByName(resourceList, "ports");
        assertEquals(9200, portsList.get(0).getRanges().getRange(0).getBegin());
        assertEquals(9200, portsList.get(0).getRanges().getRange(0).getEnd());

        assertEquals(9300, portsList.get(1).getRanges().getRange(0).getBegin());
        assertEquals(9300, portsList.get(1).getRanges().getRange(0).getEnd());

        assertEquals(9200, taskInfo.getDiscovery().getPorts().getPorts(0).getNumber());
        assertEquals(9300, taskInfo.getDiscovery().getPorts().getPorts(1).getNumber());
        assertEquals(Protos.DiscoveryInfo.Visibility.EXTERNAL, taskInfo.getDiscovery().getVisibility());

        assertEquals(2, taskInfo.getContainer().getVolumesCount());
        assertEquals(Configuration.CONTAINER_PATH_DATA, taskInfo.getContainer().getVolumes(0).getContainerPath());
        assertEquals(Configuration.DEFAULT_HOST_DATA_DIR + "/" + configuration.getElasticsearchClusterName() + "/" + offer.getSlaveId().getValue(), taskInfo.getContainer().getVolumes(0).getHostPath());
        assertEquals(Protos.Volume.Mode.RW, taskInfo.getContainer().getVolumes(0).getMode());
        assertEquals(Configuration.CONTAINER_PATH_CONF_YML, taskInfo.getContainer().getVolumes(1).getContainerPath());
        assertEquals(Configuration.HOST_PATH_CONF, taskInfo.getContainer().getVolumes(1).getHostPath());
        assertEquals(Protos.Volume.Mode.RO, taskInfo.getContainer().getVolumes(1).getMode());
    }

    private Protos.Resource getResourceByName(List<Protos.Resource> resourceList, String name) {
        return resourceList.stream().filter(resource -> resource.getName().equals(name)).findFirst().get();
    }

    private List<Protos.Resource> getResourceListByName(List<Protos.Resource> resourceList, String name) {
        return resourceList.stream().filter(resource -> resource.getName().equals(name)).collect(Collectors.toList());
    }

    private Protos.Offer getOffer(Protos.FrameworkID frameworkId) {
        return Protos.Offer.newBuilder()
                                                .setId(Protos.OfferID.newBuilder().setValue(UUID.randomUUID().toString()))
                .setSlaveId(Protos.SlaveID.newBuilder().setValue(SLAVEID))
                                                .setFrameworkId(frameworkId)
                                                .setHostname("localhost")
                                                .addAllResources(asList(
                                                        Resources.singlePortRange(9200, "some-framework-role"),
                                                        Resources.singlePortRange(9300, "some-framework-role"),
                                                        Resources.cpus(1.0, "some-framework-role"),
                                                        Resources.disk(2.0, "some-framework-role"),
                                                        Resources.mem(3.0, "some-framework-role")))
                                            .build();
    }

    @Test
    public void shouldAddJarInfoAndRemoveContainerInfo() {
        when(configuration.isFrameworkUseDocker()).thenReturn(false);
        String address = "http://localhost:1234";
        when(configuration.getFrameworkFileServerAddress()).thenReturn(address);
        when(configuration.nativeCommand(any())).thenReturn("ls");
        TaskInfoFactory factory = new TaskInfoFactory(clusterState);

        Date now = new DateTime().withDayOfMonth(1).withDayOfYear(1).withYear(1970).withHourOfDay(1).withMinuteOfHour(2).withSecondOfMinute(3).withMillisOfSecond(400).toDate();
        when(clock.now()).thenReturn(now);
        when(clock.nowUTC()).thenReturn(ZonedDateTime.now(ZoneOffset.UTC));

        Protos.TaskInfo taskInfo = factory.createTask(configuration, frameworkState, getOffer(frameworkState.getFrameworkID()), clock);
        assertFalse(taskInfo.getContainer().isInitialized());
        assertTrue(taskInfo.getExecutor().getCommand().isInitialized());
        assertEquals(2, taskInfo.getCommand().getUrisCount());
        assertTrue(taskInfo.getCommand().getUris(0).getValue().contains(address));
    }

    @Test
    public void canParseTask() throws Exception {
        final ZonedDateTime nowUTC = ZonedDateTime.now(ZoneOffset.UTC);
        when(clock.nowUTC()).thenReturn(nowUTC.minusYears(100));

        final Protos.TaskID taskId = Protos.TaskID.newBuilder().setValue("TaskID").build();

        final Protos.TaskInfo taskInfo = createTaskInfo(taskId, createData(Optional.of(nowUTC)));

        final Protos.TaskStatus taskStatus = Protos.TaskStatus.newBuilder()
                .setTaskId(taskId)
                .setState(Protos.TaskState.TASK_STAGING)
                .build();

        final Task task = TaskInfoFactory.parse(taskInfo, taskStatus, clock);
        assertEquals(nowUTC, task.getStartedAt());
        verify(clock, never()).nowUTC();
    }

    @Test
    public void canParseTaskWithoutTimestamp() throws Exception {
        final ZonedDateTime nowUTC = ZonedDateTime.now(ZoneOffset.UTC);
        when(clock.nowUTC()).thenReturn(nowUTC);

        final Protos.TaskID taskId = Protos.TaskID.newBuilder().setValue("TaskID").build();

        final Protos.TaskInfo taskInfo = createTaskInfo(taskId, createData(Optional.<ZonedDateTime>empty()));

        final Protos.TaskStatus taskStatus = Protos.TaskStatus.newBuilder()
                .setTaskId(taskId)
                .setState(Protos.TaskState.TASK_STAGING)
                .build();

        final Task task = TaskInfoFactory.parse(taskInfo, taskStatus, clock);
        assertEquals(nowUTC, task.getStartedAt());
    }

    @Test
    public void canParseTaskCappedTimestamp() throws Exception {
        final ZonedDateTime nowUTC = ZonedDateTime.now(ZoneOffset.UTC);
        when(clock.nowUTC()).thenReturn(nowUTC);

        final Protos.TaskID taskId = Protos.TaskID.newBuilder().setValue("TaskID").build();

        Properties data = new Properties();
        data.put("hostname", "hostname");
        data.put("ipAddress", "ip address");
        data.put("startedAt", nowUTC.withZoneSameInstant(ZoneId.of("Europe/Paris")).toString());

        StringWriter writer = new StringWriter();
        data.list(new PrintWriter(writer));

        final Protos.TaskInfo taskInfo = createTaskInfo(taskId, ByteString.copyFromUtf8(writer.getBuffer().toString()));

        final Protos.TaskStatus taskStatus = Protos.TaskStatus.newBuilder()
                .setTaskId(taskId)
                .setState(Protos.TaskState.TASK_STAGING)
                .build();

        final Task task = TaskInfoFactory.parse(taskInfo, taskStatus, clock);
        assertEquals(nowUTC, task.getStartedAt());
    }

    @Test
    public void shouldAllowUserSpecifiedPorts() {
        when(configuration.getElasticsearchPorts()).thenReturn(Arrays.asList(123, 456));
        TaskInfoFactory factory = new TaskInfoFactory(clusterState);
        Protos.TaskInfo taskInfo = factory.createTask(configuration, frameworkState, getOffer(frameworkState.getFrameworkID()), new Clock());
        assertTrue(taskInfo.isInitialized());
        assertTrue(taskInfo.toString().contains("123"));
        assertTrue(taskInfo.toString().contains("456"));
    }

    @Test
    public void shouldAllowUserRequestForRandomHttpAndTransportPorts() {
        when(configuration.getElasticsearchPorts()).thenReturn(Arrays.asList(0, 0));
        TaskInfoFactory factory = new TaskInfoFactory(clusterState);
        Protos.TaskInfo taskInfo = factory.createTask(configuration, frameworkState, getOffer(frameworkState.getFrameworkID()), new Clock());
        assertTrue(taskInfo.isInitialized());
        assertTrue(taskInfo.toString().contains("9200"));
        assertTrue(taskInfo.toString().contains("9300"));
    }

    @Test
    public void shouldAllowUserRequestForFixedHttpAndRandomTransportPort() {
        when(configuration.getElasticsearchPorts()).thenReturn(Arrays.asList(123, 0));
        TaskInfoFactory factory = new TaskInfoFactory(clusterState);
        Protos.TaskInfo taskInfo = factory.createTask(configuration, frameworkState, getOffer(frameworkState.getFrameworkID()), new Clock());
        assertTrue(taskInfo.isInitialized());
        assertTrue(taskInfo.toString().contains("123"));
        assertTrue(taskInfo.toString().contains("9200"));
    }

    @Test
    public void shouldAllowUserRequestForRandomHttpAndFixedTransportPort() {
        when(configuration.getElasticsearchPorts()).thenReturn(Arrays.asList(0, 456));
        TaskInfoFactory factory = new TaskInfoFactory(clusterState);
        Protos.TaskInfo taskInfo = factory.createTask(configuration, frameworkState, getOffer(frameworkState.getFrameworkID()), new Clock());
        assertTrue(taskInfo.isInitialized());
        assertTrue(taskInfo.toString().contains("9200"));
        assertTrue(taskInfo.toString().contains("456"));
    }


    @Test
    public void shouldUseMesosProvidedPorts() {
        TaskInfoFactory factory = new TaskInfoFactory(clusterState);
        Protos.TaskInfo taskInfo = factory.createTask(configuration, frameworkState, getOffer(frameworkState.getFrameworkID()), new Clock());
        assertTrue(taskInfo.getContainer().isInitialized());
        assertTrue(taskInfo.isInitialized());
        assertTrue(taskInfo.toString().contains("9200"));
        assertTrue(taskInfo.toString().contains("9300"));
    }

    private Protos.TaskInfo createTaskInfo(Protos.TaskID taskId, ByteString data) {
        Protos.DiscoveryInfo.Builder discovery = Protos.DiscoveryInfo.newBuilder()
                .setPorts(Protos.Ports.newBuilder()
                        .addPorts(Discovery.CLIENT_PORT_INDEX, Protos.Port.newBuilder().setNumber(9001).setName(Discovery.CLIENT_PORT_NAME))
                        .addPorts(Discovery.TRANSPORT_PORT_INDEX, Protos.Port.newBuilder().setNumber(9002).setName(Discovery.TRANSPORT_PORT_NAME)))
                .setVisibility(Protos.DiscoveryInfo.Visibility.EXTERNAL);

        return Protos.TaskInfo.newBuilder()
                .setName("Name")
                .setTaskId(taskId)
                .setSlaveId(Protos.SlaveID.newBuilder().setValue("SlaveID").build())
                .setData(data)
                .setDiscovery(discovery)
                .build();
    }

    private ByteString createData(Optional<ZonedDateTime> startedAt) {
        Properties data = new Properties();
        data.put("hostname", "hostname");
        data.put("ipAddress", "ip address");
        if (startedAt.isPresent()) {
            data.put("startedAt", startedAt.get().toString());
        }

        StringWriter writer = new StringWriter();
        try {
            data.store(new PrintWriter(writer), "Data");
        } catch (IOException e) {
            e.printStackTrace();
        }

        return ByteString.copyFromUtf8(writer.getBuffer().toString());
    }

}
