#!/usr/bin/env python
# Copyright (c) 2013 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.

import optparse
import sys
import logging
import os
import json
import xml.etree.ElementTree as ET
from urlparse import urlparse

from snakebite.client import Client
from snakebite.errors import FileNotFoundException
from snakebite.errors import DirectoryException
from snakebite.errors import FileException
from snakebite.formatter import format_listing
from snakebite.formatter import format_results
from snakebite.formatter import format_counts
from snakebite.formatter import format_fs_stats
from snakebite.formatter import format_stat
from snakebite.formatter import format_du

log = logging.getLogger(__name__)


def exitError(error):
    print str(error)
    sys.exit(-1)


class MyParser(optparse.OptionParser):
    def format_epilog(self, formatter):
        return self.epilog


class Commands(object):
    methods = {}


def command(args="", descr="", allowed_opts="", visible=True):
    def wrap(f):
        Commands.methods[f.func_name] = {"method": f,
                                         "args": args,
                                         "descr": descr,
                                         "allowed_opts": allowed_opts,
                                         "visible": visible}
    return wrap


class SnakebiteCli(object):
    def __init__(self):
        (self.cmd, self.opts, self.args) = self._parse_opts()
        self._read_config()
        self._setup_client(self.opts.namenode, int(self.opts.port))
        self._setup_logging()

    def _read_config(self):
        ''' Check if the last argument contains hdfs://'''
        if self.args and 'hdfs://' in self.args[-1]:
            parse_result = urlparse(self.args[-1])

            # Set config
            self.opts.namenode = parse_result.hostname
            self.opts.port = parse_result.port
            self.args[-1] = parse_result.path
            return

        if self.opts.namenode and self.opts.port:
            return

        ''' Try to read the config from ~/.snakebiterc and if that doesn't exist, check $HADOOP_HOME/core-site.xml
        and create a ~/.snakebiterc from that.
        '''
        config_file = os.path.join(os.path.expanduser('~'), '.snakebiterc')

        try_paths = ['/etc/hadoop/conf/core-site.xml',
                     '/usr/local/etc/hadoop/conf/core-site.xml',
                     '/usr/local/hadoop/conf/core-site.xml']

        if os.path.exists(config_file):
            config = json.loads(open(os.path.join(os.path.expanduser('~'), '.snakebiterc')).read())
            self.opts.namenode = config['namenode']
            self.opts.port = config['port']
        elif os.environ.get('HADOOP_HOME'):
            hdfs_conf = os.path.join(os.environ['HADOOP_HOME'], 'conf', 'core-site.xml')
            self._read_hadoop_config(hdfs_conf, config_file)
        else:
            # Try to find other paths
            for hdfs_conf in try_paths:
                self._read_hadoop_config(hdfs_conf, config_file)
                # Bail out on the first find
                if self.opts.namenode and self.opts.port:
                    continue

        if self.opts.namenode and self.opts.port:
            return
        else:
            print "No ~/.snakebiterc found, no HADOOP_HOME set and no -n and -p provided"
            print "Tried to find core-site.xml in:"
            for hdfs_conf in try_paths:
                print " - %s" % hdfs_conf
            print "\nYou can manually create ~/.snakebite with the following content:"
            print "{'namenode': 'ip/hostname', 'port': 54310}"
            sys.exit(1)

    def _read_hadoop_config(self, hdfs_conf, config_file):
        if os.path.exists(hdfs_conf):
            tree = ET.parse(hdfs_conf)
            root = tree.getroot()
            for p in root.findall("./property"):
                if p.findall('name')[0].text == 'fs.defaultFS':
                    parse_result = urlparse(p.findall('value')[0].text)

                    # Set config
                    self.opts.namenode = parse_result.hostname
                    self.opts.port = parse_result.port

                    # Write config to file
                    f = open(config_file, "w")
                    f.write(json.dumps({"namenode": self.opts.namenode, "port": self.opts.port}))
                    f.close()

    def _parse_opts(self):
        usage = "usage: %prog [options] cmd [args]"

        epilog = "\nCommands:\n"
        epilog += "\n".join(sorted(["  %-30s %s" % ("%s %s" % (k, v['args']), v['descr']) for k, v in Commands.methods.iteritems() if v['visible']]))
        epilog += "\n\n"

        parser = MyParser(usage=usage, epilog=epilog)
        self.parser = parser
        # Generic options
        parser.add_option('-D', '--debug', help='Show debug information',
            action='store_true')
        parser.add_option('-j', '--json', help='JSON output', action='store_true')
        parser.add_option('-n', '--namenode', help='namenode host')
        parser.add_option('-p', '--port', help='namenode RPC port')

        parser.add_option('-R', '--recurse', help='recurse into subdirectories',
            action='store_true')
        parser.add_option('-d', '--directory', help='show only the path and no children / check if path is a dir',
            action='store_true')
        parser.add_option('-H', '--human', help='human readable output',
            action='store_true')
        parser.add_option('-s', '--summary', help='print summarized output',
            action='store_true')
        parser.add_option('-z', '--zero', help='check for zero length',
            action='store_true')
        parser.add_option('-e', '--exists', help='check if file exists',
            action='store_true')

        opts, args = parser.parse_args()

        if len(args) == 0:
            parser.print_help()
            sys.exit(-1)

        command = args.pop(0)

        return (command, opts, args)

    def _setup_logging(self):
        if self.opts.debug:
            loglevel = logging.DEBUG
        else:
            loglevel = logging.INFO
        logging.basicConfig(level=loglevel)

    def _setup_client(self, host, port):
        self.client = Client(host, port)

    def execute(self):
        if not Commands.methods.get(self.cmd):
            self.parser.print_help()
            sys.exit(-1)
        try:
            return Commands.methods[self.cmd]['method'](self)
        except FileNotFoundException, e:
            exitError(e)
        except DirectoryException, e:
            exitError(e)
        except FileException, e:
            exitError(e)

    @command(visible=False)
    def commands(self):
        print "\n".join(sorted([k for k, v in Commands.methods.iteritems() if v['visible']]))

    @command(visible=False)
    def complete(self):
        self.opts.summary = True
        for line in self._listing():
            print line.replace(" ", "\\\\ ")

    @command(args="[path]", descr="list a path", allowed_opts=["d", "R", "s"])
    def ls(self):
        for line in self._listing():
            print line

    def _listing(self):
                # Mimicking hadoop client behaviour
        if self.opts.directory:
            include_children = False
            recurse = False
            include_toplevel = True
        else:
            include_children = True
            include_toplevel = False
            recurse = self.opts.recurse

        listing = self.client.ls(self.args, recurse=recurse,
                                  include_toplevel=include_toplevel,
                                  include_children=include_children)

        for line in  format_listing(listing, json_output=self.opts.json,
                                             human_readable=self.opts.human,
                                             recursive=recurse,
                                             summary=self.opts.summary):
            yield line

    @command(args="[paths]", descr="create directories")
    def mkdir(self):
        creations = self.client.mkdir(self.args)
        for line in format_results(creations, json_output=self.opts.json):
            print line

    @command(args="[paths]", descr="create directories and their parents")
    def mkdirp(self):
        creations = self.client.mkdir(self.args, create_parent=True)
        for line in format_results(creations, json_output=self.opts.json):
            print line

    @command(args="<owner:grp> [paths]", descr="change owner", allowed_opts=["R"])
    def chown(self):
        owner = self.args.pop(0)
        try:
            mods = self.client.chown(self.args, owner, recurse=self.opts.recurse)
            for line in format_results(mods, json_output=self.opts.json):
                print line
        except FileNotFoundException, e:
            exitError(e)

    @command(args="<mode> [paths]", descr="change file mode (octal)", allowed_opts=["R"])
    def chmod(self):
        mode = int(self.args.pop(0), 8)
        mods = self.client.chmod(self.args, mode, recurse=self.opts.recurse)
        for line in format_results(mods, json_output=self.opts.json):
            print line

    @command(args="<grp> [paths]", descr="change group", allowed_opts=["R"])
    def chgrp(self):
        grp = self.args.pop(0)
        mods = self.client.chgrp(self.args, grp, recurse=self.opts.recurse)
        for line in format_results(mods, json_output=self.opts.json):
            print line

    @command(args="[paths]", descr="display stats for paths")
    def count(self):
        counts = self.client.count(self.args)
        for line in format_counts(counts, json_output=self.opts.json,
                                          human_readable=self.opts.human):
            print line

    @command(args="", descr="display fs stats")
    def df(self):
        result = self.client.df()
        for line in format_fs_stats(result, json_output=self.opts.json,
                                            human_readable=self.opts.human):
            print line

    @command(args="[paths]", descr="display disk usage statistics", allowed_opts=["s"])
    def du(self):
        if self.opts.summary:
            include_children = False
            include_toplevel = True
        else:
            include_children = True
            include_toplevel = False
        result = self.client.du(self.args, include_toplevel=include_toplevel, include_children=include_children)
        for line in format_du(result, json_output=self.opts.json, human_readable=self.opts.human):
            print line

    @command(args="[paths] dst", descr="move paths to destination")
    def mv(self):
        paths = self.args[:-1]
        dst = self.args[-1]
        result = self.client.rename(paths, dst)
        for line in format_results(result, json_output=self.opts.json):
            print line

    @command(args="[paths]", descr="remove paths", allowed_opts=["R"])
    def rm(self):
        result = self.client.delete(self.args, recurse=self.opts.recurse)
        for line in format_results(result, json_output=self.opts.json):
            print line

    @command(args="[paths]", descr="creates a file of zero length")
    def touchz(self):
        result = self.client.touchz(self.args)
        for line in format_results(result, json_output=self.opts.json):
            print line

    @command(args="", descr="show server information")
    def serverdefaults(self):
        print self.client.serverdefaults()

    @command(args="[dirs]", descr="delete a directory")
    def rmdir(self):
        result = self.client.rmdir(self.args)
        for line in format_results(result, json_output=self.opts.json):
            print line

    @command(args="<rep> [paths]", descr="set replication factor", allowed_opts=['R'])
    def setrep(self):
        rep_factor = int(self.args.pop(0))
        result = self.client.setrep(self.args, rep_factor, recurse=self.opts.recurse)
        for line in format_results(result, json_output=self.opts.json):
            print line

    @command(args="<cmd>", descr="show cmd usage")
    def usage(self):
        if len(self.args) != 1:
            self.parser.print_help()
            sys.exit(-1)
        sub_cmd = self.args[0]

        cmd_entry = Commands.methods.get(sub_cmd)
        if not cmd_entry:
            self.parser.print_help()
            sys.exit(-1)

        cmd_args = []
        allowed_opts = cmd_entry.get('allowed_opts')
        if allowed_opts:
            cmd_args += ["[-%s]" % o for o in allowed_opts]
        args = cmd_entry.get('args')
        if args:
            cmd_args.append(args)

        print "snakebite [general options] %s %s" % (sub_cmd, " ".join(cmd_args))

    @command(args="[paths]", descr="stat information")
    def stat(self):
        print format_stat(self.client.stat(self.args))

    @command(args="path", descr="test a path", allowed_opts=['d', 'z', 'e'])
    def test(self):
        path = self.args[0]
        if self.client.test(path, exists=self.opts.exists, directory=self.opts.directory, zero_length=self.opts.zero):
            sys.exit(0)
        else:
            sys.exit(1)

cliclient = SnakebiteCli()
cliclient.execute()
