#!/usr/bin/env python
"""
A fast alternative status output. It caches the api endpoint, and
operates purely against the api server environment watch and its
cache.

How does it work
----------------

 - Most of the speed gains are from caching the api endpoint (output
   from juju api-endpoints) per environment. Endpoint cache is
   invalidated automatically.

 - It uses the 'allwatcher' environment watch to retrieve the current
   state of the environment. The environment watch internally has
   a cache that it updates every 5s, so no need to query out from
   mongo.

 - It does not augment that information with any provider api calls.

Caveats
-------

- Since it does not query provider details for machines it does not
  aim for 100% compatiblity. Its main goal is speed. Unit addresses
  are available in the output per their recorded state value.


Installation
------------

# python dependencies
$ sudo apt-get install python-virtualenv

$ virtualenv juju-tools
$ source bin/juju-tools
$ easy_install pyyaml jujuclient

$ chmod +x juju-faststat
$ ./juju-faststat


if juju-faststat is on your path, it can be invoked as a juju cli plugin or ie.
$ juju faststat

cli options are the same as status, unit and service filtering is not
supported atm.
"""

import argparse
import logging
import json
import sys
import subprocess
import yaml
import os

import websocket

from jujuclient import Environment, StatusTranslator, Watcher

JUJU_HOME = os.environ.get(
    "JUJU_HOME", os.path.abspath(os.path.expanduser("~/.juju")))

log = logging.getLogger("faststat")

try:
    SafeLoader, SafeDumper = yaml.CSafeLoader, yaml.CSafeDumper
except ImportError:
    SafeLoader, SafeDumper = yaml.SafeLoader, yaml.SafeDumper


def yaml_dump(value):
    return yaml.dump(value, Dumper=SafeDumper, default_flow_style=False)


def yaml_load(value):
    return yaml.load(value, Loader=SafeLoader)


def setup_parser():
    parser = argparse.ArgumentParser(
        usage="usage: juju stat [options] [pattern ...]",
        description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        "-e", "--environment",
        help="juju environment to operate in")
    parser.add_argument(
        "-f", "--format", default="yaml",
        choices=["json", "yaml"],
        help="specify output format")
    parser.add_argument(
        "-o", "--output", default="",
        help="specify an output file")
    parser.add_argument(
        "-d", "--description", action="store_true",
        help=argparse.SUPPRESS)
    parser.add_argument(
        "--debug", "--debug", action="store_true",
        help=argparse.SUPPRESS)
    return parser


class EnvEndpoint(object):

    def __init__(self, name):
        self.name = name

    def get_endpoint(self, use_cache=True):
        endpoint = None
        if use_cache:
            endpoint = self._get_cached_endpoint()
        if endpoint is None:
            log.debug("Fetching endpoint")
            use_cache = False
            endpoint = self._get_endpoint()
        if use_cache is False:
            self._set_cached_endpoint(endpoint)
        return "wss://%s" % endpoint, use_cache

    def _get_endpoint(self):
        params = ["juju", "api-endpoints"]
        if self.name:
            params.extend(["-e", self.name])
        output = subprocess.check_output(params)
        return output.splitlines()[0]

    def _get_cache_path(self):
        cache_dir = os.path.join(JUJU_HOME, "cache")
        return os.path.join(cache_dir, "%s-api-endpoint" % self.name)

    def _set_cached_endpoint(self, endpoint):
        cache_dir = os.path.join(JUJU_HOME, "cache")
        try:
            os.mkdir(cache_dir)
        except OSError:
            pass
        cache_path = self._get_cache_path()

        if endpoint is None:
            return os.remove(cache_path)

        with open(cache_path, "w") as fh:
            fh.write(endpoint)

    def _get_cached_endpoint(self):
        cache_path = self._get_cache_path()
        if not os.path.exists(cache_path):
            return
        with open(cache_path) as fh:
            return fh.read()


def get_credentials(env):
    if not env and os.environ.get("JUJU_ENV"):
        env = os.environ.get("JUJU_ENV")
    config_path = os.path.join(JUJU_HOME, "environments.yaml")
    with open(config_path) as fh:
        content = fh.read()
    config = yaml_load(content)
    envs = config['environments']
    env = env or config.get('default')
    if not env in envs:
        raise ValueError(
            "Unknown environment specified %s valid are %s" % (
                env, envs.keys()))
    return env, envs.get(env, {}).get('admin-secret')


def main():
    parser = setup_parser()
    options = parser.parse_args()

    if options.description:
        print "A fast status output."
        sys.exit(0)

    if options.debug:
        logging.basicConfig(
            level=logging.DEBUG,
            format="%(asctime)s %(levelname)s %(name)s %(msg)s")

    env_name, credentials = get_credentials(options.environment)
    env_api = EnvEndpoint(env_name)

    endpoint, using_cache = env_api.get_endpoint()
    if using_cache:
        websocket.setdefaulttimeout(7)

    log.info("Connecting to environment %r @ %s" % (env_name, endpoint))
    try:
        env = Environment(endpoint)
    except Exception, e:  # Quite a variety of exceptions possible.
        if not using_cache:
            raise
        log.info("Stale cache, reloading endpoint")
        endpoint, _ = env_api.get_endpoint(use_cache=False)
        log.info("Connect to environment %r @ %s" % (env_name, endpoint))
        env = Environment(endpoint)

    log.info("Environment login")
    env.login(credentials)

    log.debug("Fetching status information")
    watch = Watcher(env.conn)
    status = StatusTranslator().run(watch)

    out = options.output and open(options.output, "w") or sys.stdout
    if options.format == "json":
        print >> out, json.dumps(status, indent=2)
    else:
        print >> out, yaml_dump(status)


if __name__ == '__main__':
    main()