#!/bin/bash

set -eux

# By default we install from the main Ubuntu archive.
source=`config-get source`
case $source in
    archive)
        juju-log "Configuring zookeeper using the Ubuntu archive packages..."
        zk_package="zookeeperd"
        zk_service="zookeeper"
        zk_conf="/etc/zookeeper/conf/zoo.cfg"
        zk_myid="/etc/zookeeper/conf/myid"
        zk_install="/etc/init/zookeeper.conf"
        ;;
    dev|testing|stable)
        juju-log "Configuring zookeeper using the Hadoop Ubuntu Team PPA..."
        zk_package="hadoop-zookeeper-server"
        zk_service="hadoop-zookeeper-server"
        zk_conf="/etc/hadoop-zookeeper/conf/zoo.cfg"
        zk_myid="/etc/hadoop-zookeeper/conf/myid"
        zk_install="/etc/init/hadoop-zookeeper-server.conf"
        ;;
    *)
        juju-log "Unsupported source..."
        exit 1
        ;;
esac

base_packages="ntp"
zk_port="`config-get zk-port`"

configure_sources () {
    source=`config-get source`
    case $source in
        archive)
            juju-log "Configuring zookeeper using the Ubuntu archive packages..."
            ;;
        dev|testing|stable)
            juju-log "Configuring zookeeper using the Hadoop Ubuntu Team PPA..."
            add-apt-repository ppa:hadoop-ubuntu/$source
            apt-get update
            ;;
        *)
            juju-log "Unsupported source..."
            exit 1
            ;;
    esac
}

install_zookeeper () {
    juju-log "Installing ZooKeeper..."
    apt-get install -y $zk_package $base_packages
}

configure_zookeeper () {
    juju-log "Purging any standalone configuration..."
    # Purge
    sed -i "/^server./d" $zk_conf
    sed -i "/^weight./d" $zk_conf
    sed -i "/^group./d" $zk_conf
    juju-log "Generating unique ID for this instance..."
    unit_no=`echo $JUJU_UNIT_NAME | cut -d / -f 2`
    default_weight=`config-get default_weight`
    default_group=`config-get default_group`
    echo $unit_no > $zk_myid
    juju-log "Adding this unit to the quorum..."
    hostname=`unit-get private-address`
    # Add the config
    echo "server.${unit_no}=${hostname}:2888:3888" >> $zk_conf
    echo "group.${default_group}=${unit_no}" >> $zk_conf
    echo "weight.${unit_no}=${default_weight}" >> $zk_conf
    # Expose port as required
    open-port $zk_port
}

# Service Control Commands
restart_zookeeper () {
    juju-log "Restarting ZooKeeper"
    service $zk_service status && service $zk_service restart || :
}

start_zookeeper() {
    juju-log "Starting ZooKeeper"
    service $zk_service status || service $zk_service start
}

stop_zookeeper() {
    juju-log "Stopping ZooKeeper"
    service $zk_service stop || :
}

update_group() {
    member_id=$1
    [ -z ${member_id} ] && return 0
    default_group=`config-get default_group`
    group_arr=( $(grep "group.${default_group}" ${zk_conf} | awk -F'=' '{ print $2 }' | tr ":" "\n") )
    member_found="no"
    for member in "${group_arr[@]}"
    do
        [ "${member_id}" == "${member}" ] && member_found="yes"
    done
    [ "${member_found}" == "no" ] && group_arr=("${group_arr[@]}" "${member_id}")
    retVal="group.${default_group}=${group_arr[@]}"
    retVal=`echo ${retVal} | tr " " ":"`
    sed -i "s/^group.${default_group}.*/${retVal}/" ${zk_conf}
    sed -i "s/^group.${default_group}=:/group.${default_group}=/" ${zk_conf}
}

update_external_server() {
    # Add extra server ( if configured )
    external_server=`config-get external_server`
    if [[ "${external_server}" != "" ]]; then
        external_id=`echo ${external_server} | cut -d : -f 1`
        external_group=`echo ${external_server} | cut -d : -f 2`
        if [[ "${external_group}" == "default" ]]; then
            external_group=`config-get default_group`
        fi
        external_weight=`echo ${external_server} | cut -d : -f 3`
        external_host=`echo ${external_server} | cut -d : -f 4`
        external_port1=`echo ${external_server} | cut -d : -f 5`
        external_port2=`echo ${external_server} | cut -d : -f 6`
        sed -i "/^server.${external_id}/d" $zk_conf
        sed -i "/^weight.${external_id}/d" $zk_conf
        echo "server.${external_id}=${external_host}:${external_port1}:${external_port2}" >> $zk_conf
	update_group ${external_id}	
        echo "weight.${external_id}=${external_weight}" >> $zk_conf
        open-port ${external_port1}
        open-port ${external_port2}
    fi
}

update_quorum() {
    default_group=`config-get default_group`
    default_weight=`config-get default_weight`
    # Purge out existing quorum config to deal with departure of
    # ZK nodes
    sed -i "/^server./d" $zk_conf
    sed -i "/^weight./d" $zk_conf
    sed -i "/^group./d" $zk_conf
    # Add this node back into the list
    juju-log "Adding this unit to the quorum..."
    unit_no=`echo $JUJU_UNIT_NAME | cut -d / -f 2`
    hostname=`unit-get private-address`
    server_arr[0]="server.${unit_no}=${hostname}:2888:3888"
    weight_arr[0]="weight.${unit_no}=${default_weight}"
    group_arr[0]=${unit_no}
    # Re-create based on current relations
    for member in `relation-list`
    do  
        juju-log "Adding $member to quorum"
        member_id=`echo ${member} | cut -d / -f 2`
        member_hostname=`relation-get private-address ${member}`
	server_arr=("${server_arr[@]}" "server.${member_id}=${member_hostname}:2888:3888")
	weight_arr=("${weight_arr[@]}" "weight.${member_id}=${default_weight}")
	group_arr=("${group_arr[@]}" "${member_id}")
    done
    # Dump the new config
    # servers
    for server_line in "${server_arr[@]}"
    do
        echo "${server_line}" >> ${zk_conf}
    done
    # weight
    for member_line in "${weight_arr[@]}"
    do
        echo "${member_line}" >> $zk_conf
    done
    # group
    echo "group.${default_group}=${group_arr[@]}" | tr " " ":" >> $zk_conf
}

setup_zk_interface() {
    juju-log "Setup ZooKeeper Client"
    relation-set port=$zk_port
}

COMMAND=`basename $0`

case $COMMAND in
    install)
        configure_sources
        install_zookeeper
        configure_zookeeper
        update_external_server
        restart_zookeeper
        ;;
    config-changed)
        update_external_server
        restart_zookeeper
        ;;
    start)
        start_zookeeper
        ;;
    stop)
        stop_zookeeper
        ;;
    quorum-relation-joined)
        ;;
    quorum-relation-changed|quorum-relation-departed)
        update_quorum
        update_external_server
        restart_zookeeper
        ;;
    zookeeper-relation-joined)
        setup_zk_interface
        ;;
    upgrade-charm)
        configure_zookeeper
        update_external_server
        restart_zookeeper
        ;;
    *)
        juju-log "Command not recognised"
        ;;
esac