#!/usr/bin/env python
#
# archey3 [version 0.2-2]
#
# Copyright 2010 Melik Manukyan <melik@archlinux.us>
# Copyright 2010-2011 Laurie Clark-Michalek <bluepeppers@archlinux.us>
# Distributed under the terms of the GNU General Public License v3.
# See http://www.gnu.org/licenses/gpl.txt for the full license text.
#
# Simple python script to display an Archlinux logo in ASCII art 
# Along with basic system information.

# Display [Comment/Uncomment to Enable/Disable information.]
# Protocols:
#    uname:x = return output of `uname -x` (see UNAME_FLAG_MEANINGS for more info)
#    sensors:x = return output of `sensors x`
#    env:x = return value of env variable x
#    fs:x = return space of partition at x
#    mpd:stat,hostname,port = returns value of an mpd stat (options: artists|albums|songs)

# Import libraries
import subprocess, optparse, re, sys, configparser
from subprocess import Popen, PIPE
from optparse import OptionParser
from getpass import getuser
from time import ctime, sleep
from os import getenv
import re
import os.path
from logbook import Logger

UNAME_FLAG_MEANINGS = {
    'a': 'System Infomation',
    's': 'Kernel Name',
    'n': 'Hostname',
    'r': 'Kernel Release',
    'v': 'Kernel Version',
    'm': 'Machine Hardware name',
    'p': 'Processor Type',
    'i': 'Hardware Platform',
}

LOGOS = {'Arch Linux': '''{c1}
{c1}               +                {results[0]}
{c1}               #                {results[1]}
{c1}              ###               {results[2]}
{c1}             #####              {results[3]}
{c1}             ######             {results[4]}
{c1}            ; #####;            {results[5]}
{c1}           +##.#####            {results[6]}
{c1}          +##########           {results[7]}
{c1}         ######{c2}#####{c1}##;         {results[8]}
{c1}        ###{c2}############{c1}+        {results[9]}
{c1}       #{c2}######   #######        {results[10]}
{c2}     .######;     ;###;`\".      {results[11]}
{c2}    .#######;     ;#####.       {results[12]}
{c2}    #########.   .########`     {results[13]}
{c2}   ######'           '######    {results[14]}
{c2}  ;####                 ####;   {results[15]}
{c2}  ##'                     '##  {results[16]} 
{c2} #'                         `#  {results[17]}
\x1b[0m'''
}

CLASS_MAPPINGS = {
    'distro': 'distroCheck',
    'uname': 'unameDisplay',
    'uptime': 'uptimeDisplay',
    'sensors': 'sensorDisplay',
    'wm': 'wmDisplay',
    'de': 'deDisplay',
    'packages': 'packageDisplay',
    'ram': 'ramDisplay',
    'env': 'envDisplay',
    'fs': 'fsDisplay',
    'mpd': 'mpdDisplay',
}


DE_DICT = {'gnome-session': 'GNOME',
    'ksmserver': 'KDE',
    'xfce4-session': 'Xfce 4.6',
    'lxsession': 'LXDE',
    '': 'None',
}
        
WM_DICT = {
    'awesome': 'Awesome',
    'beryl': 'Beryl',
    'blackbox': 'Blackbox',
    'compiz': 'Compiz',
    'dwm': 'DWM',
    'enlightenment': 'Enlightenment',
    'fluxbox': 'Fluxbox',
    'fvwm': 'FVWM',
    'i3': 'i3',
    'icewm': 'IceWM',
    'kwin': 'KWin',
    'metacity': 'Metacity',
    'musca': 'Musca',
    'openbox': 'Openbox',
    'pekwm': 'PekWM',
    'ratpoison': 'ratpoison',
    'scrotwm': 'ScrotWM',
    'subtle': 'subtle',
    'wmaker': 'Window Maker',
    'wmfs': 'Wmfs',
    'wmii': 'wmii',
    'xfwm4': 'Xfwm',
    re.compile('xmonad-*'): 'xmonad',
    '': 'None',
}

COLORS = {
    'black': '0',
    'red': '1',
    'green': '2',
    'yellow': '3',
    'blue': '4',
    'magenta': '5',
    'cyan': '6',
    'white': '7'
}

class display(object):
    command_line = ''
    arg1 = ''
    arg2 = ''
    arg3 = ''
    stdindata = ''
    
    def __init__(self, args, config, logger, parent=None):
        self.config = config
        self.logger = logger
        self._parent = parent
    
    @staticmethod
    def call_command(command):
        """
        Calls a command, waits for it to exit and returns all text from stdout.
        Discards all other information.
        """
        proc = Popen(command.split(), stdout=PIPE)
        proc.wait()
        return proc.communicate()[0].decode()
    
    def run_command(self):
        if self.command_line:
            if '{arg3}' in self.command_line:
                cmd = self.command_line.format(arg1=self.arg1, arg2=self.arg2,
                                               arg3=self.arg3)
            elif '{arg2}' in self.command_line:
                cmd = self.command_line.format(arg1=self.arg1, arg2=self.arg2)
            elif '{arg1}' in self.command_line:
                cmd = self.command_line.format(arg1=self.arg1)
            else:
                cmd = self.command_line
            
            try:
                self.process = Popen(cmd.split(), stdin=PIPE, stdout=PIPE,
                                      stderr=PIPE)
            except Exception as e:
                pass
    
    def render(self):
        (stdoutdata, stderrdata) = self.process.communicate(self.stdindata
                                                             or None)
        
        return self.format_output(stdoutdata.decode())
        
    def color_me(self, output, number=None, low=30, low_color='green',
                 medium=60, medium_color='yellow', high_color='red'):
        if number is None and output.isdigit():
            number = int(output)
        elif number is None:
            return output
        
        if number <= low:
            color = low_color
        elif low < number <= medium:
            color = medium_color
        elif medium < number:
            color = high_color
        
        return '{0}{1}{2}'.format(self._parent.color(color), output,
                                  self._parent.color('clear'))
    
    regex_class = re.compile("").__class__
    def process_exists(self, key):
        global PROCESSES
        if isinstance(key, self.regex_class):
            for proc in PROCESSES._processes:
                if key.search(proc):
                    return True
        return PROCESSES(key)


class fsDisplay(display):
    command_line = "df -TPh {arg1}"
    
    conversions = {
        'binary': {
               'K': 2 ** 10,
               'M': 2 ** 20,
               'G': 2 ** 30,
               'T': 2 ** 40,
               },
        'si': {
            'K': 10 ** 3,
            'M': 10 ** 6,
            'G': 10 ** 9,
            'T': 10 ** 12,
            },
    }
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        try:
            self.arg1 = kwargs["args"][0]
        except IndexError:
            self.logger.error(
                "Did not any arguments, require one, the fs to display")
            raise
    
    def format_output(self, instring):
        values = [line for line in instring.split('\n') if line][1].split()
        used = values[3]
        total = values[2]
        fstype = values[1]
        conversion_type = self.config.get('fs', 'unit', fallback="si").lower()
        conversions = self.conversions[conversion_type]
        
        mount = '/root' if self.arg1 == '/' else self.arg1 
        title = mount.split('/')[-1].title()
        try:
            #convert to straight int
            used_ = int(used[:-1]) * conversions[used[-1].upper()]
            total_ = int(total[:-1]) * conversions[total[-1].upper()]
            persentage = used_ / total_ * 100
        except Exception as e:
            self.logger.error(
                "Could not colorize output, errored with {0}".format(e))
        else:
            used = self.color_me(used, persentage)
        
        if self.config.getboolean("fs", "persentage", fallback=True):
            part = '{used} / {total} ({persentage}%) ({fstype})'.format(
                used=used, total=total, persentage=int(persentage),
                fstype=fstype)
        else:
            part = '{used} / {total} ({fstype})'.format(
                used=used, total=total, fstype=fstype)
        return title, part

class ramDisplay(display):
    command_line = "free -m"
    
    def format_output(self, instring):
        ram = ''.join(line for line in str(instring).split('\n') if\
                      line.startswith('Mem:')).split()
        used = int(ram[2]) - int(ram[5]) - int(ram[6])
        total = int(ram[1])
        title = 'RAM'
        try:
            persentage = (used / total * 100)
        except:
            used += ' MB'
        else:
            used = self.color_me(number=persentage, output=str(used) + ' MB')
        part = '{used} / {total} MB'.format(used=used, total=total)
        return title, part
        
class sensorDisplay(display):
    command_line = "sensors {arg1}"
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        arg_from_conf = self.config.get('sensor', 'sensor',
                                        fallback='coretemp-*')
        try:
            arg_from_arg = kwargs["args"][0]
        except IndexError:
            self.logger.error(
               "Did not get any arguments, require one, the sensor to display.")
            raise
        
        if arg_from_arg:
            self.arg1 = arg_from_arg
        else:
            self.arg1 = arg_from_conf
    
    def format_output(self, instring):
        tempinfo = instring.split('\n')[2::4]
        
        out = []
        for line in tempinfo:
            info = [re.sub("\s\s+", "", line) for line in line.split('  ') if\
                    line]
            value = info[1]
            intvalue = int(value[:3])
            if intvalue > 45:
                temp = (self._parent.color("red") + info[1] +
                        self._parent.color("clear"))
            elif intvalue in range(30, 45):
                temp = (self._parent.color("magenta") + info[1] +
                        self._parent.color("clear"))
            else:
                temp = (self._parent.color("green") + info[1] +
                        self._parent.color("clear"))
            out.append((info[0], temp))
        return out

class envDisplay(display):
    def __init__(self, **kwargs):
        try:
            self.arg1 = kwargs["args"][0]
        except IndexError:
            self.logger.error("Did not get any arguments, require one," +
                              " the env variable to display.")
            raise
        
        super().__init__(**kwargs)
    
    def render(self):
        argvalue = getenv(self.arg1.upper())
        return ('$' + self.arg1.upper(), argvalue)

class unameDisplay(display):
    command_line = "uname {arg1}"
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        
        try:
            flag = kwargs["args"][0]
        except IndexError:
            self.logger.error("Did not get any arguments, require one," +
                              " the flag to pass to uname")
            raise
        
        arg_from_conf = self.config.get('uname', 'argument', fallback="")
        arg_from_arg = flag
        if arg_from_arg:
            self.arg1 = '-' + arg_from_arg
        elif arg_from_conf:
            self.arg1 = '-' + arg_from_conf
        else:
            self.arg1 = ''
    
    def format_output(self, instring):
        return (UNAME_FLAG_MEANINGS[self.arg1[1]], instring)

class uptimeDisplay(display):
    def render(self):
        with open("/proc/uptime") as upfile:
            raw = upfile.read()
            fuptime = int(raw.split('.')[0])
            
            day = int(fuptime / 86400)
            fuptime = fuptime % 86400
            hour = int(fuptime / 3600)
            fuptime = fuptime % 3600
            minute = int(fuptime / 60)
            uptime = '{daystring}{hours}:{mins:02d}'.format(
                daystring='{days} day{s}, '.format(days=day, s=('s' if day > 1
                                                     else '')) if day else '',
                hours = hour, mins = minute
            )
            return "Uptime", uptime

class packageDisplay(display):
    command_line = "pacman -Q"
    
    def format_output(self, instring):
        return "Packages", len(instring.split('\n'))

class distroCheck(display):
    def render(self):
        try:
            _ = open("/etc/pacman.conf")
        except IOError:
            distro = self.call_command("uname -o")
        else:
            distro = "Arch Linux"
        distro = '{0} {1}'.format(distro, self.call_command("uname -m"))
        return "OS", distro

class processCheck(display):
    command_line = "ps -u {arg1}"
    
    render = lambda self: self
    
    def __init__(self, **kwargs):
        self.arg1 = getuser()
        super().__init__(**kwargs)
    
    def run_command(self):
        super().run_command()
        out = str(self.process.communicate()[0])
        self._processes = set([line.split()[3] for line in out.split('\\n') if\
                              len(line.split()) == 4])
    
    def __call__(self, proc):
        if proc in self._processes:
            return True
        return False
        
class wmDisplay(display):
    def render(self):
        if self.config.get('wm', 'manual', fallback=False):
            return "WM", self.config.get('wm', 'manual')
        wm = ''
        for key in WM_DICT.keys():
            if self.process_exists(key):
                wm = key
                break
        return "WM", WM_DICT[wm]

class deDisplay(display):
    def render(self):
        if self.config.get('de', 'manual', fallback=False):
            return "DE", self.config.get('de', 'manual')
        de = ''
        for key in DE_DICT.keys():
            if self.process_exists(key):
                de = key
                break
        return "DE", DE_DICT[de]

class mpdDisplay(display):
    """
    Displays certain stat about MPD database. If mpd not installed, output
    nothing.
    """
    command_line = "mpc stats --host {arg1} --port {arg2}"
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        
        try:
            self.stat = kwargs["args"][0]
        except IndexError:
            self.logger.error("Did not get any arguments, require one," +
                              " the stat to display.")
        self.arg1 = self.config.get('mpd', 'host', fallback='localhost')
        self.arg2 = self.config.getint('mpd', 'port', fallback=6600)
    
    def format_output(self, instring):
        lines = instring.split('\n')
        stats = {}
        try:
            stats['artists'] = lines[0].split(':')[1].strip()
            stats['albums'] = lines[1].split(':')[1].strip()
            stats['songs'] = lines[2].split(':')[1].strip()
        #if people don't have mpc installed then return None)
        except:
            return False
        
        return ('{statname} in MPD database'.format(statname=self.stat.title()),
               stats[self.stat])

#------------ Config    -----------

class ArcheyConfigParser(configparser.SafeConfigParser):
    """
    A parser for the archey config file.
    """

    defaults = {'core': {'align': 'top',
                         'color': 'blue',
                         'display_modules':
                         """\
distro(), uname(n), distro(r), uptime(), wm(), de(), packages(), ram(),\
 uname(p), env(editor), df(/), mpd(albums)"""
                         },
                }
    
    def read(self, file_location=None):
        """
        Loads the config options stored in at file_location. If file_location
        does not exist, it will attempt to load from the default config location
        ($XDG_CONFIG_HOME/archey3.cfg). If that does not exist, it will write a
        default config file to $XDG_CONFIG_HOME/archey3.cfg.
        """
        
        config_location = os.path.expandvars(os.path.expanduser(
                file_location or "$XDG_CONFIG_HOME/archey3.cfg"))
        
        loaded = super(ArcheyConfigParser, self).read(config_location)
        
        if file_location == None and not loaded:
            self.load_default_config()
            self.write_config(config_location)
        if not loaded:
            #Try with default
            loaded = super(ArcheyConfigParser, self).read()
        return loaded

    def load_default_config(self):
        """
        Loads the config options stored at self.defaults.
        """
        for section, values in self.defaults.items():
            if not self.has_section(section):
                self.add_section(section)
            
            for option, value in values.items():
                #strip any excess spaces
                value = re.sub("( +)", " ", value)
                self.set(section, option, value)

    def write_config(self, location):
        """
        Writes the current config to the given location.
        """
        with open(location, 'w') as configfile:
            self.write(configfile)
    

#------------ Functions -----------

def screenshot():
        print('Screenshotting in')
        for x in sorted(range(1,6), reverse=True):
                print('%s' % x, end='')
                sys.stdout.flush()
                sleep(1.0/3)
                for x in range(3):
                    print('.', end='')
                    sys.stdout.flush()
                    sleep(1.0/3)
                    
        print('Say Cheese!')
        sys.stdout.flush()
        try:
            subprocess.check_call(['import', '-window', 'root',
                                   ctime().replace(' ','_')+'.jpg'])
        except subprocess.CalledProcessError as e:
            print('Screenshot failed with return code {0}.'.format(
                e.returncode))

#------------ Display object ---------

class Archey(object):
    DISPLAY_PARSING_REGEX = "(?P<func>\w+)\((|(?P<args>[\w, /]+))\)"
    
    def __init__(self, config, options):
        self.config = config
        self.log_level = int(options.log_level)
        self.logger = Logger("Core", self.log_level)
        
        self.display = config.get("core", "display_modules")
        colorscheme = options.color or config.get("core", "color")
        for key in COLORS.keys():
            if key == colorscheme:
                self.colorcode = COLORS[key]
        
        global PROCESSES
        PROCESSES = self.render_class(processCheck, ())
        
        self.distro_name = ' '.join(
            self.render_class(distroCheck, ())[1].split()[:-1])
    
    def render(self):
        results = self.prepare_results()
        
        return LOGOS[self.distro_name].format(c1=self.color(1),
                                              c2=self.color(2),
                                              results = results
                                              )
    
    def prepare_results(self):
        """
        Renders all classes found in the display array, and then returns them
        as a list. The returned list will be exactly 18 items long, with any
        left over spaces being filled with empty strings.
        """
        outputs = []
        # Run functions found in 'display' array.
        for func_name, args in self.parse_display():
            cls = eval(CLASS_MAPPINGS[func_name])
            
            line = self.render_class(cls, args)
            if hasattr(line, "__iter__") and len(line) != 2:
                outputs.extend(line)
            elif line:
                outputs.append(line)
        
        outputs = [self.format_item(line) for line in outputs] 
        return outputs + [""] * (18 - len(outputs))
    
    def parse_display(self):
        """
        Iterates over the display attribute of the Archey class, and tries to
        parse them using the DISPLAY_PARSING_REGEX.
        """
        for func in self.display.split(","):
            func = func.strip()
            
            info = re.match(self.DISPLAY_PARSING_REGEX, func)
            if not info:
                self.logger.error(
                    "Could not parse display string {0}".format(func))
                continue
            
            groups = info.groupdict()
            if groups["args"]:
                args = [arg.strip() for arg in groups["args"].split(",")]
            else:
                args = ()
            
            yield groups["func"], args
        raise StopIteration

    def format_item(self, item):
        title = item[0].rstrip(':')
        data = str(item[1]).rstrip()
        
        #if we're dealing with a fraction
        if len(data.split('/')) == 2:
            numerator = data.split('/')[0]
            numerator = (self.color(1, bold=True) + numerator +
                         self.color('clear'))
            denominator = data.split('/')[1]
            data = '/'.join((numerator, denominator))
        
        return "{color}{title}:{clear} {data}".format(
                color=self.color(1),
                title=title,
                data=data,
                clear=self.color("clear")
        )
        
    def color(self, code, bold=False):
        if code == 2:
            bold = True
        first_bitty_bit = '\x1b[{0};'.format(int(not bold))
        if code in range(3):
            second_bitty_bit = '3{0}m'.format(self.colorcode)
        elif code == "clear":
            return '\x1b[0m'
        else:
            second_bitty_bit = '3{0}m'.format(COLORS[code])
       
        return first_bitty_bit + second_bitty_bit
    
    def render_class(self, cls, args):
        """
        Returns the result of the run_command method for the class passed.
        """
        logger = Logger(cls.__name__, self.log_level)
        try:
            instance = cls(args=args, config=self.config, logger=logger,
                           parent=self)
        except Exception as e:
            self.logger.error(
                "Could not instantiate {0}, failed with error {1}".format(
                    cls.__name__, e))
            return
        try:
            instance.run_command()
            return instance.render()
        except Exception as e:
            self.logger.error(
                "Could not render line for {0}, failed with error {1}".format(
                    cls.__name__, e))

def main():
    parser = OptionParser(
        usage='%prog',
        description="""%prog is a utility to display system info and take\
 screenshots""",
        version="%prog 0.3")
    parser.add_option('-c', '--color',
            action='store', default='blue', type='choice', dest='color',
            choices=('black',
                     'red',
                     'green',
                     'yellow',
                     'blue',
                     'magenta',
                     'cyan',
                     'white'),
            help="""choose a color: black, red, green, yellow, blue, magenta,\
            cyan, white [Default: blue]""")
    parser.add_option('-s', '--screenshot',
            action='store_true', dest='screenshot', help='Take a screenshot')
    parser.add_option('--config',
            action='store', dest='config', default=None,
            help="Set the location of the config file to load.")
    parser.add_option('--debug',
            action='store', dest='log_level', default=1)
    (options, args) = parser.parse_args()
    
    config = ArcheyConfigParser()
    config.read(options.config)
    
    archey = Archey(config=config, options=options)
    print(archey.render())
    
    if options.screenshot:
        screenshot()

if __name__ == "__main__":
    main()