package com.hubspot.singularity.scheduler;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;

import org.apache.mesos.Protos.FrameworkID;
import org.apache.mesos.Protos.Offer;
import org.apache.mesos.Protos.OfferID;
import org.apache.mesos.Protos.Resource;
import org.apache.mesos.Protos.SlaveID;
import org.apache.mesos.Protos.Value.Range;
import org.apache.mesos.Protos.Value.Ranges;
import org.apache.mesos.Protos.Value.Scalar;
import org.apache.mesos.Protos.Value.Type;
import org.junit.Assert;
import org.junit.Test;

import com.google.common.base.Optional;
import com.google.common.collect.Lists;
import com.hubspot.mesos.MesosUtils;
import com.hubspot.singularity.ExtendedTaskState;
import com.hubspot.singularity.SingularityTaskHistoryUpdate;
import com.hubspot.singularity.SingularityTaskId;

public class MesosUtilsTest {

  private void assertFound(int numPorts, Resource resource) {
    int portsFound = 0;
    for (Range r : resource.getRanges().getRangeList()) {
      portsFound += (r.getEnd() - r.getBegin()) + 1;
    }
    Assert.assertEquals(numPorts, portsFound);
  }

  private void test(int numPorts, String... ranges) {
    Resource resource = MesosUtils.getPortsResource(numPorts, buildOffer(ranges));

    assertFound(numPorts, resource);
  }

  @Test
  public void testTaskOrdering() {
    final SingularityTaskId taskId = new SingularityTaskId("r", "d", System.currentTimeMillis(), 1, "h", "r");
    final Optional<String> msg = Optional.absent();

    SingularityTaskHistoryUpdate update1 = new SingularityTaskHistoryUpdate(taskId, 1L, ExtendedTaskState.TASK_LAUNCHED, msg);
    SingularityTaskHistoryUpdate update2 = new SingularityTaskHistoryUpdate(taskId, 2L, ExtendedTaskState.TASK_RUNNING, msg);
    SingularityTaskHistoryUpdate update3 = new SingularityTaskHistoryUpdate(taskId, 2L, ExtendedTaskState.TASK_FAILED, msg);

    List<SingularityTaskHistoryUpdate> list = Arrays.asList(update2, update1, update3);

    Collections.sort(list);

    Assert.assertTrue(list.get(0).getTaskState() == ExtendedTaskState.TASK_LAUNCHED);
    Assert.assertTrue(list.get(1).getTaskState() == ExtendedTaskState.TASK_RUNNING);
    Assert.assertTrue(list.get(2).getTaskState() == ExtendedTaskState.TASK_FAILED);


  }

  @Test
  public void testSubtractResources() {
    Assert.assertEquals(createResources(3, 60, "23:23", "100:175", "771:1000"),
        MesosUtils.subtractResources(createResources(5, 100, "23:23", "100:1000"), createResources(2, 40, "176:770")));

    List<Resource> subtracted = createResources(100, 1000, "1:100", "101:1000");

    subtracted = MesosUtils.subtractResources(subtracted, createResources(5, 100, "23:74", "101:120", "125:130", "750:756"));

    Assert.assertEquals(createResources(95, 900, "1:22", "75:100", "121:124", "131:749", "757:1000"), subtracted);

    subtracted = MesosUtils.subtractResources(subtracted, createResources(20, 20, "75:90", "121:121", "757:1000"));

    Assert.assertEquals(createResources(75, 880, "1:22", "91:100", "122:124", "131:749"), subtracted);
  }


  private List<Resource> createResources(int cpus, int memory, String... ranges) {
    List<Resource> resources = Lists.newArrayList();

    if (cpus > 0) {
      resources.add(Resource.newBuilder().setType(Type.SCALAR).setName(MesosUtils.CPUS).setScalar(Scalar.newBuilder().setValue(cpus).build()).build());
    }

    if (memory > 0) {
      resources.add(Resource.newBuilder().setType(Type.SCALAR).setName(MesosUtils.MEMORY).setScalar(Scalar.newBuilder().setValue(memory).build()).build());
    }

    if (ranges.length > 0) {
      resources.add(buildRanges(ranges));
    }

    return resources;
  }

  @Test
  public void testRangeSelection() {
    test(4, "23:24", "26:26", "28:28", "29:29", "31:32");
    test(2, "22:23");
    test(3, "22:22", "23:23", "24:24", "25:25");
    test(10, "100:10000");
    test(23, "90:100", "9100:9100", "185:1000");

  }

  private Resource buildRanges(String... ranges) {
    Resource.Builder resources = Resource.newBuilder()
        .setType(Type.RANGES)
        .setName(MesosUtils.PORTS);

    Ranges.Builder rangesBuilder = Ranges.newBuilder();

    for (String range : ranges) {
      String[] split = range.split("\\:");

      rangesBuilder.addRange(
          Range.newBuilder()
          .setBegin(Long.parseLong(split[0]))
          .setEnd(Long.parseLong(split[1])));
    }

    resources.setRanges(rangesBuilder);

    return resources.build();
  }

  private Offer buildOffer(String... ranges) {
    Offer.Builder offer = Offer.newBuilder()
        .setId(OfferID.newBuilder().setValue("offerid").build())
        .setFrameworkId(FrameworkID.newBuilder().setValue("frameworkid").build())
        .setHostname("hostname")
        .setSlaveId(SlaveID.newBuilder().setValue("slaveid").build());

    offer.addResources(buildRanges(ranges));

    return offer.build();
  }


}

@@ -859,263 +859,8 @@
  %7B%0A%0A
-  private void assertFound(int numPorts, Resource resource) %7B%0A    int portsFound = 0;%0A    for (Range r : resource.getRanges().getRangeList()) %7B%0A      portsFound += (r.getEnd() - r.getBegin()) + 1;%0A    %7D%0A    Assert.assertEquals(numPorts, portsFound);%0A  %7D%0A%0A
   pr
@@ -1001,30 +1001,87 @@
     
-assertFound(numPorts, 
+Assert.assertEquals(numPorts, MesosUtils.getNumPorts(Collections.singletonList(
 reso
@@ -1085,16 +1085,18 @@
 esource)
+))
 ;%0A  %7D%0A%0A 

use MesosUtils.getNumPorts() instead of redundant MesosUtilsTest.assertFound()