/*
 * Copyright (c) 2016 Spotify AB.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package com.spotify.helios.testing;

import static com.google.common.base.Strings.isNullOrEmpty;
import static com.spotify.docker.client.DockerClient.LogsParam.follow;
import static com.spotify.docker.client.DockerClient.LogsParam.stderr;
import static com.spotify.docker.client.DockerClient.LogsParam.stdout;

import com.spotify.docker.client.DockerClient;
import com.spotify.docker.client.LogStream;
import com.spotify.helios.client.HeliosClient;
import com.spotify.helios.common.descriptors.HostStatus;
import com.spotify.helios.common.descriptors.JobId;
import com.spotify.helios.common.descriptors.TaskStatus;

import com.google.common.collect.Maps;
import com.google.common.util.concurrent.AbstractScheduledService;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

class HeliosSoloLogService extends AbstractScheduledService {

  private static final Logger log = LoggerFactory.getLogger(HeliosSoloLogService.class);

  private final HeliosClient heliosClient;
  private final DockerClient dockerClient;
  private final LogStreamProvider logStreamProvider;

  private Map<String, Future> logFutures = Maps.newHashMap();

  HeliosSoloLogService(@NotNull final HeliosClient heliosClient,
                       @NotNull final DockerClient dockerClient,
                       @NotNull final LogStreamProvider logStreamProvider) {
    super();
    this.heliosClient = heliosClient;
    this.dockerClient = dockerClient;
    this.logStreamProvider = logStreamProvider;
  }

  @Override
  protected void runOneIteration() throws Exception {
    try {
      // fetch all the jobs running on the solo deployment
      for (final String host : heliosClient.listHosts().get()) {
        final HostStatus hostStatus = heliosClient.hostStatus(host).get();
        final Map<JobId, TaskStatus> statuses = hostStatus.getStatuses();

        for (final TaskStatus status : statuses.values()) {
          final JobId jobId = status.getJob().getId();
          final String containerId = status.getContainerId();
          if (isNullOrEmpty(containerId)) {
            continue;
          }

          if (!logFutures.containsKey(containerId)) {
            // for any containers we're not already tracking, attach to their stdout/stderr
            final Future<?> future = this.executor().submit(new Runnable() {
              @Override
              public void run() {
                try {
                  final LogStream logStream = dockerClient.logs(containerId, stdout(), stderr(),
                                                                follow());

                  log.info("attaching stdout/stderr for job={}, container={}", jobId, containerId);
                  logStream.attach(logStreamProvider.getStdoutStream(jobId, containerId),
                                   logStreamProvider.getStderrStream(jobId, containerId));
                } catch (final InterruptedException e) {
                  // we're shutting down, no need to log anything
                } catch (final Exception e) {
                  log.warn("error streaming log for job={}, container={} - {}", jobId, containerId,
                           e);
                }
              }
            });

            logFutures.put(containerId, future);
          }
        }
      }
    } catch (Exception e) {
      // do nothing, we'll retry in a bit
    }
  }

  @Override
  protected void shutDown() throws Exception {
    for (final Future future : logFutures.values()) {
      future.cancel(true);
    }

    super.shutDown();
  }

  @Override
  protected Scheduler scheduler() {
    return Scheduler.newFixedDelaySchedule(0, 100, TimeUnit.MILLISECONDS);
  }

}

@@ -1317,16 +1317,79 @@
 ervice;%0A
+import com.google.common.util.concurrent.ThreadFactoryBuilder;%0A
 import o
@@ -1534,52 +1534,282 @@
 ent.
-Future;%0Aimport java.util.concurrent.TimeUnit
+ExecutionException;%0Aimport java.util.concurrent.Executors;%0Aimport java.util.concurrent.Future;%0Aimport java.util.concurrent.ScheduledExecutorService;%0Aimport java.util.concurrent.ThreadFactory;%0Aimport java.util.concurrent.TimeUnit;%0Aimport java.util.concurrent.TimeoutException
 ;%0A%0Ac
@@ -2508,24 +2508,477 @@
 vider;%0A  %7D%0A%0A
+  private static %3CT%3E T get(Future%3CT%3E future)%0A      throws InterruptedException, ExecutionException, TimeoutException %7B%0A    return future.get(1, TimeUnit.SECONDS);%0A  %7D%0A%0A  @Override%0A  protected ScheduledExecutorService executor() %7B%0A    final ThreadFactory threadFactory = new ThreadFactoryBuilder()%0A        .setNameFormat(serviceName())%0A        .setDaemon(true)%0A        .build();%0A    return Executors.newSingleThreadScheduledExecutor(threadFactory);%0A  %7D%0A%0A
   @Override%0A
@@ -3131,16 +3131,20 @@
  host : 
+get(
 heliosCl
@@ -3159,21 +3159,16 @@
 tHosts()
-.get(
 )) %7B%0A   
@@ -3202,16 +3202,20 @@
 tatus = 
+get(
 heliosCl
@@ -3235,21 +3235,16 @@
 us(host)
-.get(
 );%0A     
@@ -4794,43 +4794,53 @@
     
-// do nothing, we'll retry in a bit
+log.warn(%22Caught exception, will ignore%22, e);
 %0A   

use timeouts with Futures