/**
 * 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.storm.grouping;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Holds a list of the current loads.
 */
public class LoadMapping {
    private static final Load NOT_CONNECTED = new Load(false, 1.0, 1.0);
    private final AtomicReference<Map<Integer, Load>> local = new AtomicReference<Map<Integer, Load>>(new HashMap<Integer, Load>());
    private final AtomicReference<Map<Integer, Load>> remote = new AtomicReference<Map<Integer, Load>>(new HashMap<Integer, Load>());

    public void setLocal(Map<Integer, Double> local) {
        Map<Integer, Load> newLocal = new HashMap<Integer, Load>();
        if (local != null) {
            for (Map.Entry<Integer, Double> entry : local.entrySet()) {
                newLocal.put(entry.getKey(), new Load(true, entry.getValue(), 0.0));
            }
        }
        this.local.set(newLocal);
    }

    public void setRemote(Map<Integer, Load> remote) {
        if (remote != null) {
            this.remote.set(new HashMap<Integer, Load>(remote));
        } else {
            this.remote.set(new HashMap<Integer, Load>());
        }
    }

    public Load getLoad(int task) {
        Load ret = local.get().get(task);
        if (ret == null) {
            ret = remote.get().get(task);
        }
        if (ret == null) {
            ret = NOT_CONNECTED;
        }
        return ret;
    }

    public double get(int task) {
        return getLoad(task).getLoad();
    }
}