#!/usr/bin/env python
# Copyright 2012 the rootpy developers
# distributed under the terms of the GNU General Public License

from rootpy.extern.argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('-l', action='store_true', dest='nointro', default=False,
        help="don't print the intro message")
parser.add_argument('-u', '--update', action='store_true', default=False,
        help="open the file in UPDATE mode (default: READ)")
parser.add_argument('filename')
args = parser.parse_args()

import os
import sys

if not os.path.isfile(args.filename):
    sys.exit("File %s does not exist" % args.filename)

import readline
import cmd
import subprocess
import re
from fnmatch import fnmatch
from glob import glob

import rootpy
rootpy.log.basic_config_colorized()
from rootpy.io import open as ropen, utils
from rootpy.io.file import _DirectoryBase, DoesNotExist
from rootpy.userdata import DATA_ROOT
from rootpy.plotting import Canvas
from rootpy import rootpy_globals as _globals


EXEC_CMD = re.compile('(?P<name>\w+)\.(?P<call>\S+)')
ASSIGN_CMD = re.compile('\w+\s*((\+=)|(-=)|(=))\s*\w+')
LOAD_CMD = re.compile('^(?P<name>\S+)(:?\s+as\s+(?P<alias>\S+))?$')

# Try and get the termcolor module - pip install termcolor
try:
    from termcolor import colored
except ImportError:
    # Make a dummy function which does not color the text
    def colored(text, *args, **kwargs):
        return text

_COLOR_MATCHER = [
    (re.compile('^TH[123][CSIDF]'), 'red'),
    (re.compile('^TTree'), 'green'),
    (re.compile('^TChain'), 'green'),
    (re.compile('^TDirectory'), 'blue'),
]


def color_key(tkey):
    ''' 
    Function which returns a colorized TKey name given its type 
    '''
    name = tkey.GetName()
    classname = tkey.GetClassName()
    for class_regex, color in _COLOR_MATCHER:
        if class_regex.match(classname):
            return colored(name, color=color)
    return name


def prompt(vars, message):
    
    prompt_message = message
    try:
        from IPython.Shell import IPShellEmbed
        ipshell = IPShellEmbed(argv=[''],
                banner=prompt_message, exit_msg="Goodbye")
        return  ipshell
    except ImportError:
        ## this doesn't quite work right, in that it doesn't go to the right env
        ## so we just fail.
        import code
        import rlcompleter
        import readline
        readline.parse_and_bind("tab: complete")
        # calling this with globals ensures we can see the environment
        if prompt_message:
            print prompt_message
        shell = code.InteractiveConsole(vars)
        return shell.interact


def ioctl_GWINSZ(fd):
    
    try:
        import fcntl, termios, struct, os
        cr = struct.unpack('hh', fcntl.ioctl(fd, termios.TIOCGWINSZ, '1234'))
    except:
        return None
    return cr


def get_terminal_size():

    cr = ioctl_GWINSZ(0) or ioctl_GWINSZ(1) or ioctl_GWINSZ(2)
    if not cr:
        try:
            fd = os.open(os.ctermid(), os.O_RDONLY)
            cr = ioctl_GWINSZ(fd)
            os.close(fd)
        except:
            pass
    if not cr:
        try:
            cr = (env['LINES'], env['COLUMNS'])
        except:
            cr = (25, 80)
    return int(cr[1]), int(cr[0])


def make_identifier(name):

    # Replace invalid characters with '_'
    name = re.sub('[^0-9a-zA-Z_]', '_', name)

    # Remove leading characters until we find a letter or underscore
    return re.sub('^[^a-zA-Z_]+', '', name)


def is_valid_identifier(name):

    return name == make_identifier(name)


class shell_cmd(cmd.Cmd, object):

    def do_shell(self, s):

        subprocess.call(s, shell=True)

    def help_shell(self):

        print "execute commands in your $SHELL (i.e. bash)"


class empty_cmd(cmd.Cmd, object):
    
    def emptyline(self):

        pass


class exit_cmd(cmd.Cmd, object):

    def can_exit(self):

        return True

    def onecmd(self, line):

        r = super(exit_cmd, self).onecmd(line)
        if r and (self.can_exit() or
           raw_input('exit anyway ? (yes/no):')=='yes'):
             return True
        return False

    def do_exit(self, s):

        return True

    def help_exit(self):

        print "Exit the interpreter."
        print "You can also use the Ctrl-D shortcut."

    def do_EOF(self, s):

        print
        return True
    
    help_EOF= help_exit


def root_glob(directory, pattern):

    matches = []
    for dirpath, dirnames, filenames in \
            utils.walk(directory, maxdepth=pattern.count(os.path.sep)):
        for dirname in dirnames:
            dirname = os.path.join(dirpath, dirname)
            if fnmatch(dirname, pattern):
                matches.append(dirname)
        for filename in filenames:
            filename = os.path.join(dirpath, filename)
            if fnmatch(filename, pattern):
                matches.append(filename)
    return matches


class ROOSH(exit_cmd, shell_cmd, empty_cmd):

    ls_parser = ArgumentParser()
    ls_parser.add_argument('-l', action='store_true',
                           dest='showtype', default=False)
    ls_parser.add_argument('files', nargs='*')
    
    mkdir_parser = ArgumentParser()
    mkdir_parser.add_argument('-p', action='store_true',
                           dest='recurse', default=False,
                           help='create parent directories as required')
    mkdir_parser.add_argument('paths', nargs='*')

    def __init__(self, filename, mode='READ', stdin=None, stdout=None):
        
        if stdin is None:
            stdin = sys.stdin
        if stdout is None:
            stdout = sys.stdout
        super(ROOSH, self).__init__(stdin=stdin, stdout=stdout)
        
        root_file = ropen(filename, mode)
        self.files = {}
        self.files[filename] = root_file
        self.pwd = root_file
        self.prev_pwd = root_file
        self.current_file = root_file
        
        self.namespace = {}
        self.__update_namespace()
        self.__update_prompt()
        
        self.canvases = {}
        self.current_canvas = None
        self.latest_default_canvas = -1
    
    def __update_prompt(self):

        dirname = os.path.basename(self.pwd._path)
        if len(dirname) > 20:
            dirname = (dirname[:10] + '..' + dirname[-10:])
        self.prompt = '%s > ' % dirname
     
    def __update_namespace(self):

        for key in self.pwd.GetListOfKeys():
            name = key.GetName()
            if not key.GetClassName().startswith('TDirectory') and \
               not key.GetClassName().startswith('TTree'):
                self.namespace[name] = self.pwd.Get(name)

    def do_env(self, s):

        for name, value in self.namespace.items():
            if name == '__builtins__':
                continue
            print name, value

    def help_env(self):

        print "print all variable names and values in current environment"
        print "object (excluding directories) contained within the"
        print "current directory are automatically included by name"
   
    def do_load(self, name):

        try:
            match = re.match(LOAD_CMD, name)
            if match:
                name = match.group('name')
                alias = match.group('alias')
                if alias is None:
                    alias = make_identifier(os.path.basename(name))
                # check that alias is a valid identifier
                elif not is_valid_identifier(alias):
                    print "%s is not a valid identifier" % alias
                    return 
                self.namespace[alias] = self.pwd.Get(name)
            else:
                self.default(name)
        except DoesNotExist, e:
            print e
    
    def complete_load(self, text, line, begidx, endidx):

        return self.completion_helper(text, line, begidx, endidx, 'TTree')
         
    def help_load(self):

        print "load the specified object into the current namespace"
        print "TTrees are not loaded automatically when cd'ing into a directory"
        print "Use this command to load TTrees or other objects not in current directory"
        print "Use 'load foo as bar' to alias the object named foo as bar"

    def do_cd(self, path):
        
        prev_pwd = self.pwd
        if path == '.':
            return
            self.prev_pwd = self.pwd
        try:
            if not path:
                self.pwd = self.current_file
            elif path == '-':
                self.pwd = self.prev_pwd
                self.do_pwd()
            else:
                self.pwd = self.pwd.GetDirectory(path)
            self.pwd.cd()
            self.__update_namespace()
            self.__update_prompt()
            self.prev_pwd = prev_pwd
        except DoesNotExist, e:
            print e
    
    def complete_cd(self, text, line, begidx, endidx):

        return self.completion_helper(
                text, line, begidx, endidx, 'TDirectoryFile')

    def help_cd(self):

        print "change the current directory"
        print "'cd -' will change to the previous directory"
        print "'cd' (with no path) will change to the root directory"
        print "All non-TTree objects are loaded"
        print "into the current namespace automatically"

    def do_ls(self, args=None):
        
        if args is None:
            args = '' 
        args = ROOSH.ls_parser.parse_args(args.split())
        if not args.files:
            args.files = ['']
        for i, path in enumerate(args.files):
            if '*' in path:
                paths = root_glob(self.pwd, path)
                if not paths:
                    paths = [path]
            else:
                paths = [path]
            for path in paths:
                _dir = self.pwd
                if path:
                    try:
                        _dir = self.pwd.Get(path)
                    except DoesNotExist, e:
                        print e
                        continue
                if isinstance(_dir, _DirectoryBase):
                    if len(args.files) > 1:
                        if i > 0:
                            print
                        print '%s:' % _dir.GetName()
                    keys = _dir.unique_keys()
                    keys.sort(key=lambda key: key.GetName())
                    things = [color_key(key) for key in keys]
                    if things:
                        self.columnize(things)
                else:
                    print path

    def complete_ls(self, text, line, begidx, endidx):

        return self.completion_helper(text, line, begidx, endidx) 
    
    def help_ls(self):

        print "list items contained in a directory"

    def do_mkdir(self, args=None):
        
        if args is None:
            args = '' 
        args = ROOSH.mkdir_parser.parse_args(args.split())
        
        for path in args.paths:
            try:
                utils.mkdir(path, args.recurse)
            except Exception, e:
                print e

    def complete_mkdir(self, text, line, begidx, endidx):

        return self.completion_helper(text, line, begidx, endidx,
                typename='TDirectoryFile')

    def completion_helper(self, text, line, begidx, endidx, typename=None):    
        
        things = []
        directory = self.pwd
        head = ''
        if begidx != endidx:
            prefix = line[begidx: endidx]
            head, prefix = os.path.split(prefix)
            if head:
                try:
                    directory = directory.GetDirectory(head)
                except DoesNotExist:
                    return []
        else:
            prefix = ''
        for key in directory.GetListOfKeys():
            if typename is not None:
                if key.GetClassName() != typename:
                    continue
            name = key.GetName()
            if prefix and not name.startswith(prefix):
                continue
            if key.GetClassName() == 'TDirectoryFile':
                things.append(os.path.join(head, '%s/' %  name))
            else:
                things.append(os.path.join(head, name))
        return things

        
    def do_pwd(self, s=None):
        
        print self.pwd._path
    
    def help_pwd(self):

        print "print the current directory"
     
    def help_help(self):

        print "'help CMD' will print help for a command"
        print "'help' will print all available commands"
    
    def do_python(self, s=None):
    
        prompt(self.namespace,'')()
    
    def help_python(self):

        print "drop into an interactive Python shell"
        print "anything loaded into your current namespace"
        print "will be handed over to Python"
    
    def do_canvas(self, name=None):
        
        if self.current_canvas is None and _globals.pad is not None:
            self.canvases['0'] = _globals.pad
            _globals.pad.SetTitle('0')
            self.latest_default_canvas += 1
        if not name:
            name = self.latest_default_canvas + 1
            while str(name) in self.canvases:
                name += 1
            self.latest_default_canvas = name
            name = str(name)
        elif name in self.canvases:
            canvas = self.canvases[name]
            canvas.cd()
            self.current_canvas = canvas
            print "switching to previous canvas '%s'" % name
            return
        print "switching to new canvas '%s'" % name
        canvas = Canvas(name=name, title=name)
        canvas.cd()
        self.current_canvas = canvas
        self.canvases[name] = canvas

    def help_canvas(self):

        print "switch to a new or previous canvas"

    def complete_canvas(self, text, line, begidx, endidx):
        
        names = []
        if begidx != endidx:
            prefix = line[begidx: endidx]
        else:
            prefix = ''
        for name in self.canvases.keys():
            if prefix and not name.startswith(prefix):
                continue
            names.append(name)
        return names
    
    def do_roosh(self, filename=None):

        if not filename:
            for name, file in self.files.items():
                if file == self.current_file:
                    print '* %s' % name
                else:
                    print '  %s' % name
            return
        if not os.path.isfile(filename):
            print "file '%s' does not exist" % filename
            return
        prev_pwd = self.pwd
        if filename not in self.files:
            print "switching to new file %s" % filename
            file = ropen(filename)
            self.files[filename] = file
        else:
            print "switching to previous file %s" % filename
            file = self.files[filename]
        self.pwd = file
        self.current_file = file
        self.pwd.cd()
        self.__update_namespace()
        self.__update_prompt()
        self.prev_pwd = prev_pwd
    
    def help_roosh(self):

        print "switch to a new or previous file"
    
    def complete_roosh(self, text, line, begidx, endidx):
        
        names = []
        if begidx != endidx:
            prefix = line[begidx: endidx]
        else:
            prefix = ''
        for name in glob(prefix + '*'):
            if os.path.isdir(name) or fnmatch(name, '*.root*'):
                names.append(name)
        if len(names) == 1 and os.path.isdir(names[0]):
            names[0] = os.path.normpath(names[0]) + os.path.sep
        return names

    def completenames(self, text, *ignored):
        
        dotext = 'do_' + text
        cmds = [a[3:] for a in self.get_names() if a.startswith(dotext)]
        objects = [key for key in self.namespace.keys() if key.startswith(text)]
        return cmds + objects

    def completedefault(self, text, line, begidx, endidx):

        return self.completion_helper(text, line, begidx, endidx)

    def default(self, line):
        
        try:
            if not re.match(ASSIGN_CMD, line):
                line = line.strip()
                if (not line.startswith('print') and
                    not line.startswith('from ') and
                    not line.startswith('import ') and
                    not line.startswith('with ') and
                    not line.startswith('if ')):
                    line = '__ = ' + line
            exec line in self.namespace
            if '__' in self.namespace:
                if self.namespace['__'] is not None:
                    print self.namespace['__']
                del self.namespace['__']
            return
        except Exception, e:
            print e
            return
        return super(ROOSH, self).default(line)


history_file = os.path.join(DATA_ROOT, 'roosh_history')
if os.path.exists(history_file):
    readline.read_history_file(history_file)
history_size = os.getenv('ROOSH_HISTORY_SIZE', 500)
readline.set_history_length(history_size)
try:
    terminal = ROOSH(
            args.filename,
            mode='UPDATE' if args.update else 'READ')
    if args.nointro:
        terminal.cmdloop()
    else:
        terminal.cmdloop("Welcome to the ROOSH terminal\ntype help for help")
    readline.write_history_file(history_file)
except Exception, e:
    readline.write_history_file(history_file)
    sys.exit(e)
