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

import kafka.common.TopicAndPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.storm.kafka.trident.GlobalPartitionInformation;

import java.util.*;

import static org.apache.storm.kafka.KafkaUtils.taskPrefix;

public class ZkCoordinator implements PartitionCoordinator {
    private static final Logger LOG = LoggerFactory.getLogger(ZkCoordinator.class);

    SpoutConfig _spoutConfig;
    int _taskIndex;
    int _totalTasks;
    int _taskId;
    String _topologyInstanceId;
    Map<Partition, PartitionManager> _managers = new HashMap<>();
    List<PartitionManager> _cachedList = new ArrayList<>();
    Long _lastRefreshTime = null;
    int _refreshFreqMs;
    DynamicPartitionConnections _connections;
    DynamicBrokersReader _reader;
    ZkState _state;
    Map _stormConf;

    public ZkCoordinator(DynamicPartitionConnections connections, Map stormConf, SpoutConfig spoutConfig, ZkState state, int taskIndex, int totalTasks, int taskId, String topologyInstanceId) {
        this(connections, stormConf, spoutConfig, state, taskIndex, totalTasks, taskId, topologyInstanceId, buildReader(stormConf, spoutConfig));
    }

    public ZkCoordinator(DynamicPartitionConnections connections, Map stormConf, SpoutConfig spoutConfig, ZkState state, int taskIndex, int totalTasks, int taskId, String topologyInstanceId, DynamicBrokersReader reader) {
        _spoutConfig = spoutConfig;
        _connections = connections;
        _taskIndex = taskIndex;
        _totalTasks = totalTasks;
        _taskId = taskId;
        _topologyInstanceId = topologyInstanceId;
        _stormConf = stormConf;
        _state = state;
        ZkHosts brokerConf = (ZkHosts) spoutConfig.hosts;
        _refreshFreqMs = brokerConf.refreshFreqSecs * 1000;
        _reader = reader;
    }

    private static DynamicBrokersReader buildReader(Map stormConf, SpoutConfig spoutConfig) {
        ZkHosts hosts = (ZkHosts) spoutConfig.hosts;
        return new DynamicBrokersReader(stormConf, hosts.brokerZkStr, hosts.brokerZkPath, spoutConfig.topic);
    }

    @Override
    public List<PartitionManager> getMyManagedPartitions() {
        if (_lastRefreshTime == null || (System.currentTimeMillis() - _lastRefreshTime) > _refreshFreqMs) {
            refresh();
            _lastRefreshTime = System.currentTimeMillis();
        }
        return _cachedList;
    }

    @Override
    public void refresh() {
        try {
            LOG.info(taskPrefix(_taskIndex, _totalTasks, _taskId) + " Refreshing partition manager connections");
            List<GlobalPartitionInformation> brokerInfo = _reader.getBrokerInfo();
            List<Partition> mine = KafkaUtils.calculatePartitionsForTask(brokerInfo, _totalTasks, _taskIndex, _taskId);

            Set<Partition> curr = _managers.keySet();
            Set<Partition> newPartitions = new HashSet<>(mine);
            newPartitions.removeAll(curr);

            Set<Partition> deletedPartitions = new HashSet<>(curr);
            deletedPartitions.removeAll(mine);

            LOG.info("{} Deleted partition managers: {}",
                taskPrefix(_taskIndex, _totalTasks, _taskId), deletedPartitions);

            Map<TopicAndPartition, PartitionManager> deletedManagers = new HashMap<>();
            for (Partition id : deletedPartitions) {
                PartitionManager manager = _managers.remove(id);
                manager.close();
                deletedManagers.put(new TopicAndPartition(id.topic, id.partition), manager);
            }

            LOG.info("{} New partition managers: {}",
                taskPrefix(_taskIndex, _totalTasks, _taskId), newPartitions);

            for (Partition id : newPartitions) {
                PartitionManager man = new PartitionManager(
                        _connections,
                        _topologyInstanceId,
                        _state,
                        _stormConf,
                        _spoutConfig,
                        id,
                        deletedManagers.get(new TopicAndPartition(id.topic, id.partition)));
                _managers.put(id, man);
            }

        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        _cachedList = new ArrayList<>(_managers.values());
        LOG.info(taskPrefix(_taskIndex, _totalTasks, _taskId) + " Finished refreshing");
    }

    @Override
    public PartitionManager getManager(Partition partition) {
        return _managers.get(partition);
    }
}
