/**
 * Licensed to Apereo under one or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information regarding copyright ownership. Apereo
 * 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 the
 * following location:
 *
 * <p>http://www.apache.org/licenses/LICENSE-2.0
 *
 * <p>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.apereo.portal.portlet.container.cache;

import com.google.common.base.Function;
import java.io.IOException;
import java.io.Writer;
import org.apache.commons.io.output.NullWriter;
import org.apereo.portal.utils.TeeWriter;

/**
 * Subclass of {@link TeeWriter} that stops writing to the branch once the limit is hit by calling
 * {@link #setBranch(Writer)} with {@link NullWriter}.
 *
 * <p>A callback {@link Function} to be executed when the limit is hit can be provided as well.
 *
 */
public class LimitingTeeWriter extends TeeWriter {
    private final long maximumCharacters;
    private final Writer branch;
    private final Function<LimitingTeeWriter, ?> limitReachedCallback;
    private long characterCount = 0;
    private boolean limitReached = false;

    public LimitingTeeWriter(long maximumCharacters, Writer out, Writer branch) {
        this(maximumCharacters, out, branch, null);
    }

    public LimitingTeeWriter(
            long maximumCharacters,
            Writer out,
            Writer branch,
            Function<LimitingTeeWriter, ?> limitReachedCallback) {
        super(out, branch);
        this.maximumCharacters = maximumCharacters;
        this.branch = branch;
        this.limitReachedCallback = limitReachedCallback;
    }

    /**
     * Sets the character count back to 0. If {@link #isLimitReached()} is true it is set back to
     * false and the original branch {@link Writer} is set as the current branch again.
     */
    public void resetByteCount() {
        this.characterCount = 0;

        if (limitReached) {
            limitReached = false;
            this.setBranch(branch);
        }
    }

    /** @return Number of characters seen so far */
    public long getCharacterCount() {
        return characterCount;
    }

    /** @return true if the limit has been reached */
    public boolean isLimitReached() {
        return limitReached;
    }

    @Override
    protected void beforeWrite(int n) throws IOException {
        this.characterCount += n;

        if (this.maximumCharacters > 0
                && !this.limitReached
                && this.characterCount > this.maximumCharacters) {
            //Hit limit, replace tee'd writer with a null writer
            this.limitReached = true;
            this.setBranch(NullWriter.NULL_WRITER);

            if (this.limitReachedCallback != null) {
                this.limitReachedCallback.apply(this);
            }
        }
    }
}