#!/usr/bin/env python2
#
# Copyright Pravega Authors.
#
# 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.
#
import argparse
import subprocess
import re

try:
    env = subprocess.check_output(
        "docker service inspect pravega_segmentstore --format='{{.Spec.TaskTemplate.ContainerSpec.Env}}'", shell=True)

    def get_config(name):
        return re.findall(name + "=(.+//)?(.+?)( |]|\r\n|\n)", env)[0][1]

    published_address = get_config("publishedIPAddress")
    hdfs_url = get_config("HDFS_URL")
    zk_url = get_config("ZK_URL")

    existing = subprocess.check_output("docker service ls --format '{{.Name}},{{.Image}}' -f 'name=pravega_segmentstore'", shell=True)
    instances = set([int(instance[1]) for instance in re.findall("(pravega_segmentstore_(\d+))", existing)])
    instances.add(1) # add the first service
    image=existing.split()[0].split(",")[1]
except Exception as e:
    print(e)
    print("pravega_segmentstore isn't running, can't be queried, or is improperly configured for scaling")
    exit(1)

def instance_count(value):
    ivalue = int(value)
    if ivalue < 1:
        raise argparse.ArgumentTypeError("%s must be greater than or equal to 1" % value)
    return ivalue

parser = argparse.ArgumentParser()
parser.add_argument("instance_count", help="the number of instance to scale up/down to", type=instance_count)
args = parser.parse_args()
requested = set(range(1, args.instance_count + 1))

to_destroy = sorted(list(instances - requested))
to_create = sorted(list(requested - instances))

for instance in to_destroy:
    name = "pravega_segmentstore_%i" % instance
    print("Destroying " + name)
    subprocess.call("docker service rm " + name, shell=True)

for instance in to_create:
    name = "pravega_segmentstore_%i" % instance
    port = 12345 + instance
    print("Creating " + name)
    template = """docker service create --name={name} \
          --network pravega_default \
          --label com.docker.stack.namespace='pravega' \
          -p '{port}:{port}' \
          -e HDFS_URL={hdfs_url} \
          -e TIER2_STORAGE=HDFS \
          -e HDFS_REPLICATION=1 \
          -e ZK_URL={zk_url} \
          -e CONTROLLER_URL=tcp://controller:9090 \
          -e JAVA_OPTS='-Dpravegaservice.service.published.host.nameOrIp={published_address} -Dpravegaservice.service.listener.port={port} -Dpravegaservice.service.listener.host.nameOrIp=0 -Dbookkeeper.ensemble.size=2 -Dbookkeeper.ack.quorum.size=2 -Dbookkeeper.write.quorum.size=2 -Xmx900m -XX:OnError="kill -9 p%" -XX:+ExitOnOutOfMemoryError -XX:+CrashOnOutOfMemoryError -XX:+HeapDumpOnOutOfMemoryError' \
          {image} segmentstore"""

    cmd = template.format(
        name=name,
        port=port,
        hdfs_url=hdfs_url,
        zk_url=zk_url,
        published_address=published_address,
        image=image
    )
    subprocess.call(cmd, shell=True)
