johnzon-core/src/main/java/org/apache/johnzon/core/Snippet.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 org.apache.johnzon.core;

import javax.json.Json;
import javax.json.JsonArray;
import javax.json.JsonObject;
import javax.json.JsonValue;
import javax.json.stream.JsonGenerator;
import javax.json.stream.JsonGeneratorFactory;
import java.io.ByteArrayOutputStream;
import java.io.Closeable;
import java.io.IOException;
import java.io.OutputStream;
import java.io.UncheckedIOException;
import java.util.Collections;
import java.util.Map;

public class Snippet {

    private final int max;
    private final JsonGeneratorFactory generatorFactory;

    public Snippet(final int max) {
        this(max, Json.createGeneratorFactory(Collections.EMPTY_MAP));
    }

    public Snippet(final int max, final JsonGeneratorFactory generatorFactory) {
        this.max = max;
        this.generatorFactory = generatorFactory;
    }

    public String of(final JsonValue value) {
        switch (value.getValueType()) {
            case TRUE: return "true";
            case FALSE: return "false";
            case NULL: return "null";
            default: {
                try (final Buffer buffer = new Buffer()) {
                    buffer.write(value);
                    return buffer.get();
                }
            }
        }
    }

    public static String of(final JsonValue value, final int max) {
        return new Snippet(max).of(value);
    }

    class Buffer implements Closeable {
        private final JsonGenerator generator;
        private final SnippetOutputStream snippet;
        private Runnable close;

        private Buffer() {
            this.snippet = new SnippetOutputStream(max);
            this.generator = generatorFactory.createGenerator(snippet);
            this.close = () -> {
                try {
                    generator.close();
                } finally {
                    this.close = () -> {
                    };
                }
            };
        }

        private void write(final JsonValue value) {
            if (snippet.isComplete()) {
                return;
            }

            switch (value.getValueType()) {
                case ARRAY: {
                    write(value.asJsonArray());
                    break;
                }
                case OBJECT: {
                    write(value.asJsonObject());
                    break;
                }
                default: {
                    generator.write(value);
                }
            }
        }

        private void write(final JsonArray array) {
            if (snippet.isComplete()) {
                return;
            }

            if (array.isEmpty()) {
                generator.write(array);
                return;
            }

            generator.writeStartArray();
            for (final JsonValue jsonValue : array) {
                if (snippet.isComplete()) {
                    break;
                }
                write(jsonValue);
            }
            generator.writeEnd();
        }

        private void write(final JsonObject object) {
            if (snippet.isComplete()) {
                return;
            }

            if (object.isEmpty()) {
                generator.write(object);
                return;
            }

            generator.writeStartObject();
            for (final Map.Entry<String, JsonValue> entry : object.entrySet()) {
                if (snippet.isComplete()) {
                    break;
                }
                write(entry.getKey(), entry.getValue());
            }
            generator.writeEnd();
        }

        private void write(final String name, final JsonValue value) {
            if (snippet.isComplete()) {
                return;
            }

            switch (value.getValueType()) {
                case ARRAY:
                    generator.writeStartArray(name);
                    final JsonArray array = value.asJsonArray();
                    for (final JsonValue jsonValue : array) {
                        if (snippet.isComplete()) {
                            break;
                        }
                        write(jsonValue);
                    }
                    generator.writeEnd();

                    break;
                case OBJECT:
                    generator.writeStartObject(name);
                    final JsonObject object = value.asJsonObject();
                    for (final Map.Entry<String, JsonValue> keyval : object.entrySet()) {
                        if (snippet.isComplete()) {
                            break;
                        }
                        write(keyval.getKey(), keyval.getValue());
                    }
                    generator.writeEnd();

                    break;
                default: generator.write(name, value);
            }
        }

        private String get() {
            // If close is not called the string may be empty
            close();
            return snippet.get();
        }

        @Override
        public void close() {
            close.run();
        }
    }

    static class SnippetOutputStream extends OutputStream {

        private final ByteArrayOutputStream buffer;
        private OutputStream mode;

        public SnippetOutputStream(final int max) {
            this.buffer = new ByteArrayOutputStream();
            this.mode = new Writing(max, buffer);
        }

        public String get() {
            if (isComplete()) {
                return buffer.toString() + "...";
            } else {
                return buffer.toString();
            }
        }

        public boolean isComplete() {
            return mode instanceof Ignoring;
        }

        @Override
        public void write(final int b) throws IOException {
            mode.write(b);
        }

        @Override
        public void write(final byte[] b) throws IOException {
            mode.write(b);
        }

        @Override
        public void write(final byte[] b, final int off, final int len) throws IOException {
            mode.write(b, off, len);
        }

        @Override
        public void flush() throws IOException {
            mode.flush();
        }

        @Override
        public void close() throws IOException {
            mode.close();
        }

        public void print(final String string) {
            try {
                mode.write(string.getBytes());
            } catch (IOException e) {
                throw new UncheckedIOException(e);
            }
        }

        class Writing extends OutputStream {
            private final int max;
            private int count;
            private final OutputStream out;

            public Writing(final int max, final OutputStream out) {
                this.max = max;
                this.out = out;
            }

            @Override
            public void write(final int b) throws IOException {
                if (++count < max) {
                    out.write(b);
                } else {
                    endReached();
                }
            }

            @Override
            public void write(final byte[] b) throws IOException {
                write(b, 0, b.length);
            }

            @Override
            public void write(final byte[] b, final int off, final int len) throws IOException {
                final int remaining = max - count;

                if (remaining <= 0) {

                    endReached();

                } else if (len > remaining) {

                    out.write(b, off, remaining);
                    endReached();

                } else {
                    out.write(b, off, len);
                }
            }

            private void endReached() throws IOException {
                mode = new Ignoring();
                flush();
                close();
            }
        }

        static class Ignoring extends OutputStream {
            @Override
            public void write(final int b) throws IOException {
            }
        }

    }
}