#!/usr/bin/env python

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

from txzookeeper.retry import RetryClient

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

from juju.charm.errors import CharmNotFound
from juju.charm.publisher import CharmPublisher
from juju.charm.repository import (
    CharmURL, LocalCharmRepository, RemoteCharmRepository, CS_STORE_URL)
from juju.control.utils import sync_environment_state
from juju.environment.config import EnvironmentsConfig
from juju.errors import JujuError

from juju.state.endpoint import RelationEndpoint
from juju.state.service import ServiceStateManager
from juju.state.machine import MachineStateManager
from juju.state.placement import _unassigned_placement
from juju.state.relation import RelationStateManager



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


@inlineCallbacks
def load(env_configs, env, repo_dir, data):
    """

    Check for name conflicts
     - Upload charms
     - Create services (+ config, constraints)
     - Add units
     - Add relations
    """
    provider = env.get_machine_provider()
    client = yield provider.connect()
    client = RetryClient(client)

    # Make sure the provider is ready to go.
    yield sync_environment_state(client, env_configs, env.name)

    # Get state managers
    services = ServiceStateManager(client)
    relations = RelationStateManager(client)
    machines = MachineStateManager(client)

    # First detect conflicts
    existing_services = yield services.get_all_service_states()
    env_svc_names = [e.service_name for e in existing_services]
    import_svc_names = [i['name'] for i in data['services']]
    conflict_svc_names = set(env_svc_names) & set(import_svc_names)
    if conflict_svc_names:
        raise JujuError(
            "Import has name conflicts with existing services %s" % (
                ", ".join(conflict_svc_names)))

    # Next find all the nesc charms
    if repo_dir:
        repository = LocalCharmRepository(repo_dir)
    else:
        repository = None

    store = RemoteCharmRepository(CS_STORE_URL)

    charms = []
    for s in data['services']:
        curl = CharmURL.infer(s['charm'], env.default_series)
        if s['charm'].startswith('local:'):
            if repository is None:
                raise JujuError("Local charm needed but repository not specified")
            try:
                charm = yield repository.find(curl)
            except CharmNotFound:
                # retry with less specific version
                charm = yield repository.find(
                    CharmURL.infer(
                        s['charm'][:s['charm'].rfind('-')],
                        env.default_series))
            charms.append((s['charm'], charm))
        else:
            assert s['charm'].startswith('cs:')
            charm = yield store.find(curl)
            charms.append((s['charm'], charm))

    # Upload charms to the environment
    publisher = CharmPublisher(client, provider.get_file_storage())
    for cid, c in charms:
        log.info("Publishing charm %s", cid)
        yield publisher.add_charm(
            str(CharmURL.infer(cid, env.default_series)),
            c)
    charm_states = yield publisher.publish()

    # Index by url
    charm_map = dict([(ci[0], cs) for ci, cs in zip(charms, charm_states)])
    constraint_set = yield provider.get_constraint_set()

    # Create services
    for s in data['services']:
        log.info("Creating service %s %s", s['name'], charm_map[s['charm']].id)
        constraints = constraint_set.load(s['constraints'])
        svc = yield services.add_service_state(
            s['name'],
            charm_map[s['charm']],
            constraints)
        config = yield svc.get_config()
        config.update(s['config'])
        yield config.write()

        # Create units
        log.info("Creating units")
        for i in range(s['unit_count']):
            unit = yield svc.add_unit_state()
            yield _unassigned_placement(client, machines, unit)

    # Add relations
    for r in data['relations']:
        if len(r) == 1:
            eps = [RelationEndpoint(*r[0])]
        else:
            eps = [RelationEndpoint(*r[0]), RelationEndpoint(*r[1])]
        log.info("Adding relation %s" %(
            " ".join(map(
                lambda r: "%s:%s %s" % (r.service_name, r.relation_name, r.relation_role),
                eps))))
        yield relations.add_relation_state(*eps)



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.")

    parser.add_argument(
        "--repository", "-r", default=os.environ.get("JUJU_REPOSITORY"),
        help="Reposiory for local charms.")

    parser.add_argument(
        "export_file", nargs=None,
        help="Environment export file")

    return parser


@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 main():
    """
    Load an environment export into an environment

      jitsu import <environment-export-json-file>

    The target environment must already be bootstrapped. Any conflicts
    among service names between the export and the target environment
    are fatal (and detected early).

    if the environment uses local charms then a local repository needs
    to be specified. For eample

      jitsu import --repository=$HOME/charms env.json

    Both JUJU_REPOSITORY and JUJU_ENV variables are respected.

    Load also supports reading from stdin ie.

      jitsu export -e staging | jitsu import -e production
    """
    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)

    if options.export_file == "-":
        export_file = sys.stdin
    else:
        export_path = os.path.abspath(
            os.path.expanduser(options.export_file))
        if not os.path.exists(export_path):
            raise JujuError("Invalid export file, does not exist %s" % export_path)
        export_file = open(export_path)

    try:
        export_data = json.loads(export_file.read())
    except Exception, e:
        raise JujuError(
            "Invalid environment export, malformed json %s %s" % (
                options.export_file, e))

    env_config, environment = get_environment(options.environment)

    reactor.callWhenRunning(
        run_one, load,
        env_config,
        environment,
        options.repository,
        export_data,
        )
    reactor.run()


if __name__ == '__main__':
    main()
