/*
 * Bach - Java Shell Builder
 * Copyright (C) 2020 Christian Stein
 *
 * 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
 *
 *     https://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 test.base;

import static java.lang.annotation.ElementType.METHOD;
import static java.lang.annotation.RetentionPolicy.RUNTIME;

import java.io.ByteArrayOutputStream;
import java.io.PrintStream;
import java.lang.annotation.Retention;
import java.lang.annotation.Target;
import java.util.List;
import java.util.stream.Collectors;
import org.junit.jupiter.api.extension.BeforeEachCallback;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.api.extension.ExtensionContext;
import org.junit.jupiter.api.extension.ExtensionContext.Namespace;
import org.junit.jupiter.api.extension.ParameterContext;
import org.junit.jupiter.api.extension.ParameterResolver;
import org.junit.jupiter.api.parallel.ResourceLock;
import org.junit.jupiter.api.parallel.Resources;

/** Swallow and capture texts printed to system streams. */
@Target(METHOD)
@Retention(RUNTIME)
@ResourceLock(Resources.SYSTEM_OUT)
@ResourceLock(Resources.SYSTEM_ERR)
@ExtendWith(SwallowSystem.Extension.class)
public @interface SwallowSystem {

  class Extension implements BeforeEachCallback, ParameterResolver {

    @Override
    public void beforeEach(ExtensionContext context) {
      getOrComputeStreamsIfAbsent(context);
    }

    @Override
    public boolean supportsParameter(ParameterContext parameterContext, ExtensionContext context) {
      var isTestMethod = context.getTestMethod().isPresent();
      var isStreamsParameter = parameterContext.getParameter().getType() == Streams.class;
      return isTestMethod && isStreamsParameter;
    }

    @Override
    public Streams resolveParameter(ParameterContext parameterContext, ExtensionContext context) {
      return getOrComputeStreamsIfAbsent(context);
    }

    private Streams getOrComputeStreamsIfAbsent(ExtensionContext context) {
      var store = context.getStore(Namespace.create(getClass(), context.getRequiredTestMethod()));
      return store.getOrComputeIfAbsent(Streams.class);
    }
  }

  class Streams implements ExtensionContext.Store.CloseableResource {

    private final PrintStream standardOut, standardErr;
    private final ByteArrayOutputStream out, err;

    Streams() {
      this.standardOut = System.out;
      this.standardErr = System.err;
      this.out = new ByteArrayOutputStream();
      this.err = new ByteArrayOutputStream();
      System.setOut(new PrintStream(out));
      System.setErr(new PrintStream(err));
    }

    @Override
    public void close() {
      System.setOut(standardOut);
      System.setErr(standardErr);
    }

    public List<String> lines() {
      return out.toString().lines().collect(Collectors.toList());
    }

    public List<String> errors() {
      return err.toString().lines().collect(Collectors.toList());
    }

    @Override
    public String toString() {
      return "out=```" + out.toString() + "```, err=```" + err.toString() + "```";
    }
  }
}

@@ -904,16 +904,44 @@
 Target;%0A
+import java.util.ArrayList;%0A
 import j
@@ -2790,16 +2790,64 @@
 ut, err;
+%0A    private final List%3CRunnable%3E shutdownHooks;
 %0A%0A    St
@@ -3018,24 +3018,70 @@
 utStream();%0A
+      this.shutdownHooks = new ArrayList%3C%3E();%0A
       System
@@ -3153,32 +3153,127 @@
 m(err));%0A    %7D%0A%0A
+    public void addShutdownHook(Runnable runnable) %7B%0A      shutdownHooks.add(runnable);%0A    %7D%0A%0A
     @Override%0A  
@@ -3364,16 +3364,60 @@
 rdErr);%0A
+      shutdownHooks.forEach(Runnable::run);%0A
     %7D%0A%0A 

Support runnable shutdown hooks