#!/usr/bin/env python

import argparse
import json
import logging
import os
import textwrap
import zookeeper

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

from juju.environment.config import EnvironmentsConfig

from juju.state.service import ServiceStateManager
from juju.state.relation import RelationStateManager
from juju.state.environment import EnvironmentStateManager


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


@inlineCallbacks
def export(env):

    provider = env.get_machine_provider()
    client = yield provider.connect()

    data = {'services': [],
            'relations': []}

    rels = set()

    # Get services
    services = ServiceStateManager(client)
    relations = RelationStateManager(client)
    environment = EnvironmentStateManager(client)

    for s in (yield services.get_all_service_states()):
        units = yield s.get_unit_names()
        charm_id = yield s.get_charm_id()
        constraints = yield s.get_constraints()
        config = yield s.get_config()
        exposed = yield s.get_exposed_flag()

        log.info("processing service %s / %s", s.service_name, charm_id)

        # Really!? relation-type not available via state api
        topology = yield relations._read_topology()

        # More efficient to do this via topology, ie op per rel
        # instead inspection of both sides by each endpoint, which is
        # kinda of gross, need some state api support. But trying to
        # stick to the state api as much as possible.
        svc_rels = yield relations.get_relations_for_service(s)

        for r in svc_rels:
            relation_type = topology.get_relation_type(r.internal_relation_id)
            rdata = [
                (s.service_name,
                 relation_type,
                 r.relation_name,
                 r.relation_role,
                 r.relation_scope)]
            if r.relation_role == "peer":
                rels.add(tuple(rdata))
                continue
            rel_services = yield r.get_service_states()

            # Get the endpoint's svc rel state.
            found = None
            for ep_svc in rel_services:
                if ep_svc.service_name == s.service_name:
                    continue
                ep_rels = yield relations.get_relations_for_service(ep_svc)
                for ep_r in ep_rels:
                    if ep_r.internal_relation_id != r.internal_relation_id:
                        continue
                    found = (ep_svc, ep_r)
                    break

            if not found:
                log.info("Couldn't decipher rel %s %s",
                         s.service_name, r)
                continue

            ep_svc, ep_r = found
            rdata.append(
                (ep_svc.service_name,
                 relation_type,
                 ep_r.relation_name,
                 ep_r.relation_role,
                 ep_r.relation_scope))
            rdata.sort()
            rels.add(tuple(rdata))

        data['services'].append(
            {'name': s.service_name,
             'charm': charm_id,
             'exposed': exposed,
             'unit_count': len(units),
             'constraints': constraints.data,
             'config': dict(config.items())})

    data['relations'] = list(rels)
    data['constraints'] = (yield environment.get_constraints()).data
    print json.dumps(data, indent=2)


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


def get_environment(env_option):
    env_config = EnvironmentsConfig()
    env_config.load_or_write_sample()

    if env_option is None:
        environment = env_config.get_default()
    else:
        environment = env_config.get(env_option)

    return env_config, environment


def get_parser():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=textwrap.dedent(main.__doc__))

    parser.add_argument(
        "--environment", "-e", default=os.environ.get("JUJU_ENV"),
        help="Environment to operate on.")

    return parser


def main():
    """
    Export a juju environment.

    Prints a json formatted file containing information needed to recreate
    an environment (services, relations, config) to stdout.

    Example:
       jitsu export -e env_name

    Also respects JUJU_ENV for environment to use, defaults to default env.
    """
    parser = get_parser()
    options = parser.parse_args()

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

    zookeeper.set_debug_level(0)

    env_config, environment = get_environment(options.environment)
    reactor.callWhenRunning(run_one, export, environment)
    reactor.run()


if __name__ == '__main__':
    main()
