/*
 * Copyright (C) 2015 The Android Open Source Project
 *
 * 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 vogar.testing;

import com.google.common.base.Joiner;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.io.UnsupportedEncodingException;
import java.util.EnumMap;
import java.util.EnumSet;
import org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;

/**
 * A {@link TestRule} that will intercept content written to {@link System#out} and/or
 * {@link System#err} and collate it for use by the test.
 */
public class InterceptOutputStreams implements TestRule {

    /**
     * The streams that can be intercepted.
     */
    public enum Stream {
        OUT {
            @Override
            PrintStream get() {
                return System.out;
            }

            @Override
            void set(PrintStream stream) {
                System.setOut(stream);
            }
        },
        ERR {
            @Override
            PrintStream get() {
                return System.err;
            }

            @Override
            void set(PrintStream stream) {
                System.setErr(stream);
            }
        };

        abstract PrintStream get();

        abstract void set(PrintStream stream);
    }

    /**
     * The streams to intercept.
     */
    private final EnumSet<Stream> streams;
    private final EnumMap<Stream, State> streams2State;

    /**
     * The streams to intercept.
     */
    public InterceptOutputStreams(Stream... streams) {
        this.streams = EnumSet.of(streams[0], streams);
        streams2State = new EnumMap<>(Stream.class);
    }

    /**
     * Get the intercepted contents for the stream.
     * @param stream the stream whose contents are required.
     * @return the intercepted contents.
     * @throws IllegalStateException if the stream contents are not being intercepted (in which
     *     case the developer needs to add {@code stream} to the constructor parameters), or if the
     *     test is not actually running at the moment.
     */
    public String contents(Stream stream) {
        if (!streams.contains(stream)) {
            EnumSet<Stream> extra = streams.clone();
            extra.add(stream);
            String message = "Not intercepting " + stream + " output, try:\n"
                    + "    new " + InterceptOutputStreams.class.getSimpleName() + "("
                    + Joiner.on(", ").join(extra)
                    + ")";
            throw new IllegalStateException(message);
        }

        State state = streams2State.get(stream);
        if (state == null) {
            throw new IllegalStateException(
                    "Attempting to access stream contents outside the test");
        }

        return state.contents();
    }

    @Override
    public Statement apply(final Statement base, Description description) {
        return new Statement() {
            @Override
            public void evaluate() throws Throwable {
                for (Stream stream : streams) {
                    State state = new State(stream);
                    streams2State.put(stream, state);
                }

                try {
                    base.evaluate();
                } finally {
                    for (State state : streams2State.values()) {
                        state.reset();
                    }
                    streams2State.clear();
                }
            }
        };
    }

    private static class State {
        private final PrintStream original;
        private final ByteArrayOutputStream baos;
        private final Stream stream;

        State(Stream stream) throws IOException {
            this.stream = stream;
            original = stream.get();
            baos = new ByteArrayOutputStream();
            stream.set(new PrintStream(baos, true, "UTF-8"));
        }

        String contents() {
            try {
                return baos.toString("UTF-8");
            } catch (UnsupportedEncodingException e) {
                throw new RuntimeException(e);
            }
        }

        void reset() {
            stream.set(original);
        }
    }
}