#!/usr/bin/env python3
# -*- encoding: utf-8 -*-

#  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.

'''integration test topology builder'''
import copy
import heronpy.api.api_constants as api_constants
from heronpy.api.stream import Stream, Grouping
from heronpy.api.topology import TopologyBuilder, Topology, TopologyType
from ..core import constants as integ_const
from .aggregator_bolt import AggregatorBolt
from .integration_test_spout import IntegrationTestSpout
from .integration_test_bolt import IntegrationTestBolt

class TestTopologyBuilder(TopologyBuilder):
  """Topology Builder for integration tests

  Given spouts and bolts will be delegated by IntegrationTestSpout and IntegrationTestBolt
  classes respectively.
  """
  TERMINAL_BOLT_NAME = '__integration_test_aggregator_bolt'
  TERMINAL_BOLT_CLASS = AggregatorBolt
  DEFAULT_CONFIG = {api_constants.TOPOLOGY_DEBUG: True,
                    api_constants.TOPOLOGY_RELIABILITY_MODE:
                    api_constants.TopologyReliabilityMode.ATLEAST_ONCE,
                    api_constants.TOPOLOGY_PROJECT_NAME: "heron-integration-test"}
  def __init__(self, name, http_server_url):
    super(TestTopologyBuilder, self).__init__(name)
    self.output_location = "%s/%s" % (http_server_url, self.topology_name)
    self.set_config(self.DEFAULT_CONFIG)

    # map <name -> spout's component spec>
    self.spouts = {}
    # map <name -> bolt's component spec>
    self.bolts = {}
    # map <name -> set of parents>
    self.prev = {}

  def add_spout(self, name, spout_cls, par, config=None,
                optional_outputs=None, max_executions=None):
    """Add an integration_test spout"""
    user_spec = spout_cls.spec(name)
    spout_classpath = user_spec.python_class_path

    if hasattr(spout_cls, 'outputs'):
      user_outputs = spout_cls.outputs
    else:
      user_outputs = []

    if optional_outputs is not None:
      user_outputs.extend(optional_outputs)

    if config is None:
      _config = {}
    else:
      _config = copy.copy(config)

    if max_executions is not None:
      _config[integ_const.USER_MAX_EXECUTIONS] = max_executions

    test_spec = IntegrationTestSpout.spec(name, par, _config,
                                          user_spout_classpath=spout_classpath,
                                          user_output_fields=user_outputs)
    self.add_spec(test_spec)
    self.spouts[name] = test_spec
    return test_spec

  def add_bolt(self, name, bolt_cls, par, inputs, config=None, optional_outputs=None):
    """Add an integration_test bolt

    Only dict based inputs is supported
    """
    assert isinstance(inputs, dict)
    user_spec = bolt_cls.spec(name)
    bolt_classpath = user_spec.python_class_path

    if hasattr(bolt_cls, 'outputs'):
      user_outputs = bolt_cls.outputs
    else:
      user_outputs = []

    if optional_outputs is not None:
      user_outputs.extend(optional_outputs)

    if config is None:
      _config = {}
    else:
      _config = config

    test_spec = IntegrationTestBolt.spec(name, par, inputs, _config,
                                         user_bolt_classpath=bolt_classpath,
                                         user_output_fields=user_outputs)
    self.add_spec(test_spec)
    self.bolts[name] = test_spec
    return test_spec

  # pylint: disable=too-many-branches
  def create_topology(self):
    """Creates an integration_test topology class"""

    # first add the aggregation_bolt
    # inputs will be updated later
    aggregator_config = {integ_const.HTTP_POST_URL_KEY: self.output_location}
    self.add_bolt(self.TERMINAL_BOLT_NAME, self.TERMINAL_BOLT_CLASS, 1,
                  inputs={}, config=aggregator_config)

    # building a graph directed from children to parents, by looking only on bolts
    # since spouts don't have parents
    for name, bolt_spec in self.bolts.items():
      if name == self.TERMINAL_BOLT_NAME:
        continue

      bolt_protobuf = bolt_spec.get_protobuf()
      for istream in bolt_protobuf.inputs:
        parent = istream.stream.component_name
        if name in self.prev:
          self.prev[name].add(parent)
        else:
          parents = set()
          parents.add(parent)
          self.prev[name] = parents

    # Find the terminal bolts defined by users and link them with "AggregatorBolt".

    # set of terminal component names
    terminals = set()
    # set of non-terminal component names
    non_terminals = set()
    # 1. terminal bolts need upstream components, because we don't want isolated bolts
    # 2. terminal bolts should not exist in the prev.values(), meaning that no downstream
    for parent_set in list(self.prev.values()):
      non_terminals.update(parent_set)

    for bolt_name in list(self.prev.keys()):
      if bolt_name not in non_terminals:
        terminals.add(bolt_name)

    # will also consider the cases with spouts without children
    for spout_name in list(self.spouts.keys()):
      if spout_name not in non_terminals:
        terminals.add(spout_name)

    # add all grouping to components
    for child in list(self.prev.keys()):
      for parent in self.prev[child]:
        self._add_all_grouping(child, parent, integ_const.INTEGRATION_TEST_CONTROL_STREAM_ID)

    # then connect aggregator bolt with user's terminal components
    # terminal_outputs are output fields for terminals, list of either str or Stream
    for terminal in terminals:
      if terminal in self.bolts:
        terminal_outputs = self.bolts[terminal].outputs
      else:
        terminal_outputs = self.spouts[terminal].outputs

      # now get a set of stream ids
      stream_ids = ["default" if isinstance(out, str) else out.stream_id
                    for out in terminal_outputs]
      for stream_id in set(stream_ids):
        self._add_all_grouping(self.TERMINAL_BOLT_NAME, terminal, stream_id)

    # create topology class
    class_dict = self._construct_topo_class_dict()
    return TopologyType(self.topology_name, (Topology,), class_dict)

  def _add_all_grouping(self, child, parent, stream_id):
    """Adds all grouping between child component and parent component with a given stream id

    :type child: str
    :param child: child's component name
    :type parent: str
    :param parent: parent's component name
    :type stream_id: str
    :param stream_id: stream id
    """
    # child has to be a bolt
    child_component_spec = self.bolts[child]
    # child_inputs is dict mapping from <HeronComponentSpec|GlobalStreamId -> grouping>
    child_inputs = child_component_spec.inputs

    if parent in self.bolts:
      parent_component_spec = self.bolts[parent]
    else:
      parent_component_spec = self.spouts[parent]

    if stream_id == Stream.DEFAULT_STREAM_ID:
      child_inputs[parent_component_spec] = Grouping.ALL
    else:
      child_inputs[parent_component_spec[stream_id]] = Grouping.ALL