# 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.
''' mock_proto.py '''
from heronpy.api import api_constants
import heron.proto.execution_state_pb2 as protoEState
import heron.proto.physical_plan_pb2 as protoPPlan
import heron.proto.packing_plan_pb2 as protoPackingPlan
import heron.proto.tmaster_pb2 as protoTmaster
import heron.proto.topology_pb2 as protoTopology

# pylint: disable=no-self-use, missing-docstring
class MockProto(object):
  ''' Mocking Proto'''
  topology_name = "mock_topology_name"
  topology_id = "mock_topology_id"
  cluster = "mock_topology_cluster"
  environ = "mock_topology_environ"

  def create_mock_spout(self,
                        spout_name,
                        output_streams,
                        spout_parallelism):
    spout = protoTopology.Spout()
    spout.comp.name = spout_name
    kv = spout.comp.config.kvs.add()
    kv.key = api_constants.TOPOLOGY_COMPONENT_PARALLELISM
    kv.type = protoTopology.ConfigValueType.Value('STRING_VALUE')
    kv.value = str(spout_parallelism)
    for stream in output_streams:
      spout.outputs.add().stream.CopyFrom(stream)
    return spout

  def create_mock_resource(self):
    resource = protoPackingPlan.Resource()
    resource.cpu = 1.0
    resource.ram = 1024
    resource.disk = 1024 * 2
    return resource

  def create_mock_instance_plan(self):
    instancePlan = protoPackingPlan.InstancePlan()
    instancePlan.component_name  = "word"
    instancePlan.task_id = 1
    instancePlan.component_index = 1
    instancePlan.resource.CopyFrom(self.create_mock_resource())
    return instancePlan

  def create_mock_simple_container_plan(self):
    containerPlan = protoPackingPlan.ContainerPlan()
    containerPlan.id = 1
    containerPlan.instance_plans.extend([self.create_mock_instance_plan()])
    containerPlan.requiredResource.CopyFrom(self.create_mock_resource())

    return containerPlan

  def create_mock_simple_container_plan2(self):
    containerPlan = protoPackingPlan.ContainerPlan()
    containerPlan.id = 1
    containerPlan.instance_plans.extend([self.create_mock_instance_plan()])
    containerPlan.requiredResource.CopyFrom(self.create_mock_resource())
    containerPlan.scheduledResource.CopyFrom(self.create_mock_resource())
    return containerPlan

  def create_mock_bolt(self,
                       bolt_name,
                       input_streams,
                       output_streams,
                       bolt_parallelism):
    bolt = protoTopology.Bolt()
    bolt.comp.name = bolt_name
    kv = bolt.comp.config.kvs.add()
    kv.key = api_constants.TOPOLOGY_COMPONENT_PARALLELISM
    kv.type = protoTopology.ConfigValueType.Value('STRING_VALUE')
    kv.value = str(bolt_parallelism)
    for stream in input_streams:
      bolt.inputs.add().stream.CopyFrom(stream)
    for stream in output_streams:
      bolt.outputs.add().stream.CopyFrom(stream)
    return bolt

  def create_mock_simple_topology(
      self,
      spout_parallelism=1,
      bolt_parallelism=1):
    """
    Simple topology contains one spout and one bolt.
    """
    topology = protoTopology.Topology()
    topology.id = MockProto.topology_id
    topology.name = MockProto.topology_name

    # Stream1
    stream1 = protoTopology.StreamId()
    stream1.id = "mock_stream1"
    stream1.component_name = "mock_spout"

    # Spout1
    spout = self.create_mock_spout("mock_spout", [stream1], spout_parallelism)
    topology.spouts.extend([spout])

    # Bolt1
    bolt = self.create_mock_bolt("mock_bolt", [stream1], [], bolt_parallelism)
    topology.bolts.extend([bolt])

    return topology

  def create_mock_medium_topology(
      self,
      spout_parallelism=1,
      bolt1_parallelism=1,
      bolt2_parallelism=1,
      bolt3_parallelism=1):
    """
    Medium topology is a three stage topology
    with one spout, two mid stage bolts, and one
    last stage bolt.
    S -str1-> B1 -str3-> B3
    S -str2-> B2 -str4-> B3
    """
    topology = protoTopology.Topology()
    topology.id = "mock_topology_id"
    topology.name = "mock_topology_name"

    # Streams
    stream1 = protoTopology.StreamId()
    stream1.id = "mock_stream1"
    stream1.component_name = "mock_spout1"

    stream2 = protoTopology.StreamId()
    stream2.id = "mock_stream2"
    stream2.component_name = "mock_spout1"

    stream3 = protoTopology.StreamId()
    stream3.id = "mock_stream3"
    stream3.component_name = "mock_bolt1"

    stream4 = protoTopology.StreamId()
    stream4.id = "mock_stream4"
    stream4.component_name = "mock_bolt2"

    # Spouts
    spout1 = self.create_mock_spout("mock_spout1",
                                    [stream1, stream2],
                                    spout_parallelism)
    topology.spouts.extend([spout1])

    # Bolts
    bolt1 = self.create_mock_bolt("mock_bolt1",
                                  [stream1],
                                  [stream3],
                                  bolt1_parallelism)
    bolt2 = self.create_mock_bolt("mock_bolt2",
                                  [stream2],
                                  [stream4],
                                  bolt2_parallelism)
    bolt3 = self.create_mock_bolt("mock_bolt3",
                                  [stream3, stream4],
                                  [],
                                  bolt3_parallelism)
    topology.bolts.extend([bolt1, bolt2, bolt3])


    return topology

  def create_mock_simple_physical_plan(
      self,
      spout_parallelism=1,
      bolt_parallelism=1):
    pplan = protoPPlan.PhysicalPlan()
    pplan.topology.CopyFrom(self.create_mock_simple_topology(
        spout_parallelism,
        bolt_parallelism))
    return pplan

  def create_mock_simple_packing_plan(
    self):
    packingPlan = protoPackingPlan.PackingPlan()
    packingPlan.id = "ExclamationTopology"
    packingPlan.container_plans.extend([self.create_mock_simple_container_plan()])
    return packingPlan

  def create_mock_simple_packing_plan2(
    self):
    packingPlan = protoPackingPlan.PackingPlan()
    packingPlan.id = "ExclamationTopology"
    packingPlan.container_plans.extend([self.create_mock_simple_container_plan2()])
    packingPlan.container_plans.extend([self.create_mock_simple_container_plan()])
    return packingPlan

  def create_mock_medium_physical_plan(
      self,
      spout_parallelism=1,
      bolt1_parallelism=1,
      bolt2_parallelism=1,
      bolt3_parallelism=1):
    pplan = protoPPlan.PhysicalPlan()
    pplan.topology.CopyFrom(self.create_mock_medium_topology(
        spout_parallelism,
        bolt1_parallelism,
        bolt2_parallelism,
        bolt3_parallelism))
    return pplan

  def create_mock_execution_state(self):
    estate = protoEState.ExecutionState()
    estate.topology_name = MockProto.topology_name
    estate.topology_id = MockProto.topology_id
    estate.cluster = MockProto.cluster
    estate.environ = MockProto.environ
    return estate

  def create_mock_tmaster(self):
    tmaster = protoTmaster.TMasterLocation()
    return tmaster

  def add_topology_config(self, topology, key, value):
    kv = topology.topology_config.kvs.add()
    kv.key = key
    kv.type = protoTopology.ConfigValueType.Value('STRING_VALUE')
    kv.value = str(value)