#!/usr/bin/env python3
#
# 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)

DISPLAY = [
    'distro', # Display Operating System
    'uname:n', # Display Machine Hostname
    'uname:r',  # Display Kernel Version
    'uptime',  # Display System Uptime
    'sensors:coretemp-*', # Display System Core tempature
    'wm',  # Display Window Manager
    'de', # Display Desktop Environment
    'packages', # Display Number of Packages Installed
    'ram', # Display RAM Usage
    'uname:p', # Display CPU Model
    'env:shell', # Display Current Shell
    'env:editor', # Display Editor, using EDITOR env 
    'mpd:songs',
#    'fs:/boot', # Display /boot Partition Usage
#    'fs:/home', # Display /home Partition Usage
#    'fs:/MOUNT/POINT', # Display * Partition, Edit To Your Needs
    'fs:/', # Display / Partition Usage
]


# 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

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'''
}

FUNC_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',
    'xmonad': 'xmonad',
    '': 'None',
}

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

class display(object):
    __slots__ = ['command_line', 'arg1', 'arg2', 'arg3', 'stdindata',
                 'proccess', '_parent', 'config']
    
    command_line = ''
    arg1 = ''
    arg2 = ''
    arg3 = ''
    stdindata = ''
    
    def __init__(self, config, parent=None):
        self.config = config
        self._parent = parent
    
    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.proccess = Popen(cmd.split(), stdin=PIPE, stdout=PIPE,
                                      stderr=PIPE)
            except Exception as e:
                pass
    
    def render(self):
        (stdoutdata, stderrdata) = self.proccess.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'))

class fsDisplay(display):
    command_line = "df -TPh {arg1}"
    
    conversions = {
        'binary': {
               'KB': 2 ** 10,
               'MB': 2 ** 20,
               'GB': 2 ** 30,
               'TB': 2 ** 40,
               },
        'si': {
            'KB': 10 ** 3,
            'MB': 10 ** 6,
            'GB': 10 ** 9,
            'TB': 10 ** 12,
            },
    }
    
    def __init__(self, config, path='/', parent=None):
        self.arg1 = path
        self.config = config
        self._parent = parent
    
    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').lower()
        conversions = self.conversions[conversion_type]
        
        mount = '/root' if self.arg1 == '/' else self.arg1 
        title = mount.split('/')[-1].title() + " FS"
        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:
            pass
        else:
            used = self.color_me(used, persentage)
        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, config, sensors='', parent=None):
        self.config = config
        self._parent = parent
        arg_from_conf = config.get('sensor', 'sensor', 'coretemp-*')
        arg_from_arg = sensors
        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, config, env='SHELL', parent=None):
        self.arg1 = env
        self.config = config
        self._parent = parent
    
    def render(self):
        argvalue = getenv(self.arg1.upper())
        return ('$' + self.arg1.upper(), argvalue)

class unameDisplay(display):
    command_line = "uname {arg1}"
    
    def __init__(self, config, flag=False, parent=None):
        self.config = config
        self._parent = parent
        arg_from_conf = config.get('uname', 'argument', '')
        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:
            rc = open("/etc/pacman.conf")
        except IOError:
            distro = render(unameDisplay("o"))[1]
        else:
            distro = "Arch Linux"
        distro = ' '.join([distro, render(unameDisplay('m'))[1]])
        return "OS", distro

class proccessCheck(display):
    __slots__ = display.__slots__ + ['_processes']
    command_line = "ps -u {arg1}"
    
    render = lambda self: self
    
    def __init__(self, config, parent=None):
        self.arg1 = getuser()
        self.config = config
        self._parent = parent
    
    def run_command(self):
        super().run_command()
        out = str(self.proccess.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):
        wm = ''
        for key in WM_DICT.keys():
            if processes(key):
                wm = key
                break
        return "WM", WM_DICT[wm]

class deDisplay(display):
    def render(self):
        de = ''
        for key in DE_DICT.keys():
            if processes(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, config, stat='songs', parent=None):
        self.stat = stat
        self.hostname = config.get('mpd', 'host', 'localhost')
        self.port = int(config.get('mpd', 'port', 8800))
        self.arg1 = hostname
        self.arg2 = port
        self.method = config.get('mpd', 'method', 'mpc')
        self.config = config
        self._parent = parent
    
    def run_command(self):
        if self.method == 'python3mpd':
            self.proccess = mpd.MPDClient()
            try:
                self.proccess.connect(self.hostname, self.port)
            except Exception:
                self.method == 'mpc'
                super().run_command()
        else:
            super().run_command()
    
    def render(self):
        if self.method == 'python3mpd':
            try:
                return '{statname} in MPD database'.format(
                statname=self.stat.title()), self.proccess.stats()[self.stat]
            except Exception:
                return False
        else:
            if hasattr(self, "proccess"):
                return super().render()
            else:
                return None
    
    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 ValueError:
            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',
                         'modules':
                         """
                         os, hostname, kernel, uptime, wm, de, pacman, ram, cpu,
                         env:editor, df:/, mpd:albums
                         """
                         },
                }
    
    def __init__(self):
        super(ArcheyConfigParser, self).__init__()

    def read(self, file_location="$XDG_CONFIG_HOME/archey3.cfg"):
        """
        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))
        loaded = super(ArcheyConfigParser, self).read(config_location)
        if file_location == "$XDG_CONFIG_HOME/archey3.cfg":
            self.load_default_config()
            self.write_config(file_location)
        if not loaded:
            #Try with default
            loaded = super(ArcheyConfigParser, self).read()

    def load_default_config(self):
        """
        Loads the config options stored at self.defaults.
        """
        for section, values in self.defaults.iteritems():
            for option, value in values.iteritems():
                #strip any excess spaces
                while value.find('  ') != -1:
                    value = value.replace('  ', ' ')
                self.set(section, option, value)

    def write_config(self, location):
        """
        Writes the current config to the given location.
        """
        with open(location, 'wb') 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))

def render(instance):
    try:
        instance.run_command()
        return instance.render()
    except:
        return None

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

class Archey(object):
    def __init__(self, display=None, colorscheme='blue'):
        self.display = display or []
        for key in COLORS.keys():
            if key == colorscheme:
                self.colorcode = COLORS[key]
        
        self.distro_name = ' '.join(render(distroCheck())[1].split()[:-1])
    
    def render(self):
        results = self.prepare_results()
        
        print(LOGOS[self.distro_name].format(c1=self.color(1),
                                             c2=self.color(2),
                                             results = results
                                            ))
    
    def prepare_results(self):
        outputs = []
        # Run functions found in 'display' array.
        for func in self.display:
            if len(func.split(':')) > 1:
                cls = eval(FUNC_MAPPINGS[func.split(':')[0]])
                args = tuple(func.split(':')[1].split(','))
            else:
                cls = eval(FUNC_MAPPINGS[func])
                args = ()
            
            out = render(cls(*args, parent=self))
            if isinstance(out, list):
                for o in out:
                    outputs.append(o)
            elif out:
                outputs.append(out)
            else:
                continue
        
        for i in range(18):
            if i < len(outputs):
                outputs[i] = self.format_item(outputs[i])
            else:
                outputs.append('')
        
        return outputs

    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 main():
    global processes
    processes = render(proccessCheck())
    
    parser = OptionParser(
        usage='%prog [-c COLOR] [-s, --screenshot] [-d, --display]', 
        description="""Archey is a utility to take """,
        version="%prog 0.2-1")
    parser.add_option('-c',
            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('-d', '--display',
            action='store', dest='display',
                      help="""Define info to be displayed by archey. DISPLAY
                      must be in the format "archey_display_module:arg1,arg2,arg
                      3.second_display_module". Availible display modules are un
                      ame, fs, env, uptime, sensors, ram, de and wm.
                      See the DISPLAY list in the archey source file for more
                      info.""")
    (options, args) = parser.parse_args()
    
    if options.display:
        global DISPLAY
        DISPLAY = options.display.split('.')
    
    archey = Archey(display=DISPLAY, colorscheme=options.color)
    archey.render()
    
    if options.screenshot:
        screenshot()

if __name__ == "__main__":
    main()