/*
 * Copyright (C) 2019 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 com.android.server.wifi.util;

import android.util.ArrayMap;

import java.lang.reflect.Array;
import java.util.Iterator;
import java.util.Map;

/**
 * Counts occurrences of keys of type K, where K is a custom key class. This is most commonly used
 * for counting occurrences of composite keys.
 * @param <K> Custom key type K must override {@link Object#hashCode()},
 * {@link Object#equals(Object)}, and {@link Object#toString()}. These methods can be auto-generated
 *           by IntelliJ by right clicking inside your class and clicking "Generate...".
 */
public class ObjectCounter<K> implements Iterable<Map.Entry<K, Integer>> {
    private ArrayMap<K, Integer> mCounter;

    public ObjectCounter() {
        mCounter = new ArrayMap<>();
    }

    /**
     * Removes all keys and counts from this counter.
     */
    public void clear() {
        mCounter.clear();
    }

    /**
     * Returns the number of keys in this counter.
     */
    public int size() {
        return mCounter.size();
    }

    /**
     * Gets the number of occurrences of a key
     */
    public int getCount(K key) {
        return mCounter.getOrDefault(key, 0);
    }

    /**
     * Increments the count of a key by 1.
     */
    public void increment(K key) {
        add(key, 1);
    }

    /**
     * Increments the count of a key by <code>count</code>.
     */
    public void add(K key, int count) {
        int curCount = getCount(key);
        mCounter.put(key, curCount + count);
    }

    /**
     * Returns a string representation of this counter object suitable for dump(). Note that the
     * type K must have an overridden implementation of {@link Object#toString()}.
     */
    @Override
    public String toString() {
        return mCounter.toString();
    }

    /**
     * Iterates over all (key, count) pairs.
     */
    @Override
    public Iterator<Map.Entry<K, Integer>> iterator() {
        return mCounter.entrySet().iterator();
    }

    /**
     * Converter function that converts a single (key, count) pair to a Protobuf object.
     * @param <I> The type of the key.
     * @param <O> The type of the Protobuf output.
     */
    public interface ProtobufConverter<I, O> {
        /**
         * Converter function that converts a single (key, count) pair to a Protobuf object.
         * @param key the key that we are counting occurrences for
         * @param count the number of occurrences for this key
         * @return the Protobuf output
         */
        O convert(I key, int count);
    }

    /**
     * Converts this object to a custom Protobuf representation.
     * @param protoClass the class object for the Protobuf type.
     * @param converter a conversion function.
     * @param <T> the type of the Protobuf output.
     * @return an array of Protobuf representation of buckets generated by the converter function.
     */
    public <T> T[] toProto(Class<T> protoClass, ProtobufConverter<K, T> converter) {
        @SuppressWarnings("unchecked")
        T[] output = (T[]) Array.newInstance(protoClass, size());
        int i = 0;
        for (Map.Entry<K, Integer> entry : this) {
            output[i] = converter.convert(entry.getKey(), entry.getValue());
            i++;
        }
        return output;
    }
}