#!/usr/bin/env python
#
# open-port - opens a port on a service unit or for all units in a service

import argparse
import zookeeper
import logging

from juju.environment.config import EnvironmentsConfig
from juju.state.service import parse_service_name, ServiceStateManager
from juju.hooks.cli import parse_port_protocol

from twisted.internet import reactor
from twisted.internet.defer import inlineCallbacks

log = logging.getLogger("jitsu.port")


@inlineCallbacks
def open_port(provider, service_or_unit, port):
    client = yield provider.connect()

    all_units = False
    try:
        service_name, sequence = parse_service_name(service_or_unit)
    except ValueError:
        service_name = service_or_unit
        all_units = True

    services = ServiceStateManager(client)
    service_state = yield services.get_service_state(service_name)

    if all_units:
        units = yield service_state.get_unit_names()
    else:
        units = [service_or_unit]

    for unit_name in units:
        unit = yield service_state.get_unit_state(unit_name)
        yield unit.open_port(*port)
        log.info("Opened port %r on %s" % (port, unit_name))


@inlineCallbacks
def run_one(func, *args, **kw):
    try:
        yield func(*args, **kw)
    finally:
        reactor.stop()


def main():
    parser = argparse.ArgumentParser(
        description="expose a service unit(s) port")
    parser.add_argument("-e", "--environment", default=None,
                        help="Environment to act upon (uses default else)")
    parser.add_argument("service_or_unit", nargs=1,
                        help="Service or unit whose ports to open")
    parser.add_argument("port", nargs=1,
                        help="Port to open")
    options = parser.parse_args()

    log_options = {
        "level": logging.DEBUG,
        "format": "%(asctime)s %(name)s:%(levelname)s %(message)s"}

    logging.basicConfig(**log_options)
    port = parse_port_protocol(options.port[0])
    service_or_unit = options.service_or_unit[0]

    env_config = EnvironmentsConfig()
    env_config.load_or_write_sample()

    if options.environment is None:
        environment = env_config.get_default()
    else:
        environment = env_config.get(options.environment)
    provider = environment.get_machine_provider()

    zookeeper.set_debug_level(0)

    reactor.callWhenRunning(
        run_one, open_port,
        provider, service_or_unit, port)
    reactor.run()

if __name__ == '__main__':
    main()
