#!/usr/bin/env python
"""
pgtune

Sample usage shown by running with "--help"
"""

import sys
import os
import datetime
import optparse
import csv
import platform

# Windows specific routines
try:
    # ctypes is only available starting in Python 2.5
    from ctypes import *
    # wintypes is only is available on Windows
    from ctypes.wintypes import *
  
    def Win32Memory():
        class memoryInfo(Structure):
            _fields_ = [
              ('dwLength', c_ulong),
              ('dwMemoryLoad', c_ulong),
              ('dwTotalPhys', c_ulong),
              ('dwAvailPhys', c_ulong),
              ('dwTotalPageFile', c_ulong),
              ('dwAvailPageFile', c_ulong),
              ('dwTotalVirtual', c_ulong),
              ('dwAvailVirtual', c_ulong)
              ]
        
        mi = memoryInfo()
        mi.dwLength = sizeof(memoryInfo)
        windll.kernel32.GlobalMemoryStatus(byref(mi))
        return mi.dwTotalPhys

except:
    # TODO For pre-2.5, and possibly replacing the above in all cases, you
    # can grab this from the registry via _winreg (standard as of 2.0) looking
    # at "HARDWARE\RESOURCEMAP\System Resources\Physical Memory"
    # see http://groups.google.com/groups?hl=en&lr=&client=firefox-a&threadm=b%25B_8.3255%24Dj6.2964%40nwrddc04.gnilink.net&rnum=2&prev=/groups%3Fhl%3Den%26lr%3D%26client%3Dfirefox-a%26q%3DHARDWARE%255CRESOURCEMAP%255CSystem%2BResources%255CPhysical%2BMemory%26btnG%3DSearch
    pass


# Memory constants
KB = 1024
MB = 1024 * KB
GB = 1024 * MB
KB_PER_MB = MB / KB
KB_PER_GB = GB / KB


def total_mem():
    try:
        if platform.system() == "Windows":
            mem = Win32Memory()
        else:
            # Should work on other, more UNIX-ish platforms
            physPages = os.sysconf("SC_PHYS_PAGES")
            pageSize = os.sysconf("SC_PAGE_SIZE")
            mem = physPages * pageSize
        return mem
    except:
        return None

def binaryround(value):
    """
    Keeps the 4 most significant binary bits, truncates the rest so
    that SHOW will be likely to use a larger divisor
    >>> binaryround(22)
    22
    >>> binaryround(1234567)
    1179648
    """
    multiplier = 1
    while value > 16:
        value = int(value / 2)
        multiplier = multiplier * 2
    return multiplier * value

class PGConfigLine(object):
    """
    Stores the value of a single line in the postgresql.conf file, with the 
    following fields:
      line_number : integer
      original_line : string
      comment_section : string
      sets_parameter : boolean
  
    If sets_parameter is True these will also be set:
      name : string
      readable : string
      raw : string  This is the actual value 
      delimiter (expectations are '' and "")
    """
    def __init__(self, line, num=0):
        self.original_line = line
        self.line_number = num
        self.sets_parameter = False
    
        # Remove comments and edge whitespace
        self.comment_section = ""
        self.name = None
        self.sets_parameter = None
        self.readable = None

    def process_line(self):
        """
        >>> line = PGConfigLine('checkpoint_completion_target = 0.9 # pgtune')
        >>> line.process_line()
        >>> line.comment_section
        '# pgtune'
        >>> line.name
        'checkpoint_completion_target'
        >>> line.readable
        '0.9'
        """
        line = self.original_line
        comment_index = line.find('#')
        if comment_index >= 0:
            line = self.original_line[0:comment_index]
            self.comment_section = self.original_line[comment_index:]
    
        line = line.strip()
        if line == "":
            return
    
        # Split into name,value pair
        if '=' not in line:
            return
    
        name, value = line.split('=', 1)
        name = name.strip()
        value = value.strip()
        self.name = name
        self.sets_parameter = True
    
        # Many types of values have ' ' characters around them, strip
        # TODO Set delimiter based on whether there is one here or not
        value = value.rstrip("'")
        value = value.lstrip("'")
    
        self.readable = value

    # Implement a Java-ish interface for this class that renames
    # Could use a python-ish property instead
    def value(self):
        return self.readable

    def is_setting(self):
        return self.sets_parameter

    def __str__(self):
        result = ['%s sets?=%s' %(self.line_number, self.sets_parameter)]
        if self.sets_parameter:
            result.append('%s=%s' %(self.name, self.value))
            # TODO:  Include comment_section, readable,raw, delimiter
        result.append('original_line:  %s' % self.original_line)
        return ' '.join(result)


class PGConfigFile(object):
    """
    Read, write, and manage a postgresql.conf file
  
    There are two main structures here:
  
    config_lines[]:  Array of PGConfigLine entries for each line in the file
    param_to_line:  Dictionary mapping parameter names to the line that set them
    """
  
    def __init__(self, filename):
        self.filename = filename
        self.param_to_line = {}
        self.config_lines = []
        self.settings = None

    def read_config_file(self):
        for i, line in enumerate(open(self.filename)):
            line = line.rstrip('\n')
            line_num = i + 1
    
            config_line = PGConfigLine(line, line_num)
            config_line.process_line()
            self.config_lines.append(config_line)
    
            if config_line.is_setting():
                # TODO Check if the line is already in the file, in which case
                # we should throw and error here suggesting that be corrected
                self.param_to_line[config_line.name] = config_line    

    def store_settings(self, settings):
        """
        Much of this class will only operate with a settings database.
        The only reason that isn't required by the constructor itself
        is that making it a second step introduces the possibility of
        detecting which version someone is running, based on what
        settings do and don't exist in their postgresql.conf
        """
        self.settings = settings
  
    def current_value(self, name):
        """
        Get the current value, assuming the default if that parameter
        isn't set
        """
        current = self.settings.boot_val(name)
        if name in self.param_to_line:
            current = self.settings.parse(name, self.param_to_line[name].value())
        current = current.strip()
        return current

    def numeric_value(self, name, value):
        """ 
        Get any numeric value the way the server will see it, so things
        are always on the same scale.  Returns None if this is not a
        numeric value.  
        TODO Maybe throw an exception instead?
        TODO Finish this implementation for integers, floats
        """
        return None
  
    def limit_checked(self, name, value):
        """
        TODO Check against min,max.  Clip to edge and issue hint
        if value is outside of server limits.
        """
        return None

    def update_setting(self, name, value):
        current = self.current_value(name)
        value = str(value).strip()
    
        # If it matches what's currently in the file, don't do anything
        if current == value:
            return
    
        # TODO Throw a HINT if you're reducing a value.  This only makes
        # sense for integer and float settings, and presumes that there
        # aren't any settings where a lower value is more aggressive
    
        # TODO Clamp the new value against the min and max for this setting
        #print name,"min=",settings.min_val(name),"max=",settings.max_val(name) 
    
        # Construct a new settings line
        text = "%s = %s # pgtune wizard %s" %(name, value, datetime.date.today())
        new_line = PGConfigLine(text)
        new_line.process_line()
    
        # Comment out any line already setting this value
        if name in self.param_to_line:
            old_line = self.param_to_line[name]
            old_line_num = old_line.line_number
            commentedLineText = "# %s" % old_line.original_line
            commentedLine = PGConfigLine(commentedLineText, old_line_num)
            commentedLine.process_line()
            # Subtract one here to adjust for zero offset of array.
            # Any future change that adds lines in-place will need to do
            # something smarter here, because the line numbers won't match 
            # the array indexes anymore
            self.config_lines[old_line_num - 1] = commentedLine
        
        self.config_lines.append(new_line)
        self.param_to_line[name] = new_line
    
    def update_if_larger(self, name, value):
        if name in self.param_to_line:
            # TODO This comparison needs all the values converted to numeric form
            # and converted to the same scale before it will work
            if (True):  #newValue > self.param_to_line[name].value():
                self.update_setting(name, value)
    
    def write(self, fout):
        fout.writelines(['%s\n' % line.original_line for line in self.config_lines])

    def debug_print_input(self):
        print "Original file:"
        for l in self.config_lines:
            print str(l)

    def debug_print_settings(self):
        print "Settings listing:"
        for k, line in self.param_to_line.items():
            print '%s = %s' %(k, line.value())


class PGSettings(object):
    """
    Read and index a delimited text dump of a typical pg_settings dump for 
    the appropriate architecture.  Maximum values are different for some
    settings on 32 and 64 bit platforms.
    
    An appropriately formatted dump can be generated with:
  
    psql postgres -c "COPY (SELECT name,setting,unit,category,short_desc,
    extra_desc,context,vartype,min_val,max_val,enumvals,boot_val
    FROM pg_settings WHERE NOT source='override')
    TO '/<path>/pg_settings-<ver>-<bits>'"
  
    Note that some of these columns (such as boot_val) are only available 
    starting in PostgreSQL 8.4
    """

    def __init__(self, settings_dir):
        self.param_to_dict = {}
        self.settings_dir = settings_dir

    def read_config_file(self):
        platform_bits = 32
        if platform.architecture()[0] == "64bit":
            platform_bits = 64
            
        # TODO Support handling versions other than 8.4
        # TODO Allow passing in platform bit size
        setting_dump_file = os.path.join(self.settings_dir, "pg_settings-8.4-%s" % platform_bits)
        setting_columns = ["name", "setting", "unit", "category", "short_desc",
                          "extra_desc", "context", "vartype", "min_val", "max_val", "enumvals",
                          "boot_val"]
        reader = csv.DictReader(open(setting_dump_file), setting_columns, delimiter="\t")
        for d in reader:
            # Convert nulls into blanks
            for key in d.keys():
                if d[key] == '\\N':
                    d[key] = ""
        
            # Memory units must be specified in some number of kB (never a larger 
            # unit).  Typically they are either "kB" for 1kB or "8kB", unless someone
            # compiled the server with a larger database or xlog block size
            # (BLCKSZ/XLOG_BLCKSZ).  This code has no notion that such a thing is
            # possible though.
            d['memory_unit'] = d['unit'].endswith('kB')
            if d['memory_unit']:
                divisor = d['unit'].rstrip('kB')
                if divisor == '':
                    divisor = "1"
                d['memory_divisor'] = int(divisor)
            else:
                d['memory_divisor'] = None
      
            self.param_to_dict[d['name']] = d
    
    def debug_print_settings(self):
        for key in self.param_to_dict.keys():
            print "key=", key, " value=", self.param_to_dict[key]

    def min_val(self, setting):
        return (self.param_to_dict[setting])['min_val']

    def max_val(self, setting):
        return (self.param_to_dict[setting])['max_val']

    def boot_val(self, setting):
        return (self.param_to_dict[setting])['boot_val']
  
    def unit(self, setting):
        return (self.param_to_dict[setting])['unit']
  
    def vartype(self, setting):
        return (self.param_to_dict[setting])['vartype']
  
    def memory_unit(self, setting):
        return (self.param_to_dict[setting])['memory_unit']
  
    def memory_divisor(self, setting):
        return (self.param_to_dict[setting])['memory_divisor']
    
    def show(self, name, value):
        formatted = value
        s = self.param_to_dict[name]
    
        if s['memory_unit']:
            # Use the same logic as the GUC code that implements "SHOW".  This uses
            # larger units only if there's no loss of resolution in displaying
            # with that value.  Therefore, if using this to output newly assigned
            # values, that value needs to be rounded appropriately if you want
            # it to show up as an even number of MB or GB
            if (value % KB_PER_GB == 0):
                value = value / KB_PER_GB
                unit = "GB"
            elif (value % KB_PER_MB == 0):
                value = value / KB_PER_MB
                unit = "MB"
            else:
                unit = "kB"
            formatted = str(value) + unit
      
        # print >> sys.stderr,"Showing",name,"with value",value,"gives",formatted
        return formatted

    def parse_int(self, name, value):
        """
        Parse an integer value into its internal form.  The main
        difficulty here is that if that integer is a memory unit, you
        need to be aware of what unit it is specified in.  1kB and 8kB
        pages are two popular ones and that is reflected in
        memory_divisor

        >>> ps = PGSettings('.')
        >>> ps.read_config_file()
        >>> ps.parse_int('max_connections', '10')
        10
        >>> ps.parse_int('shared_buffers', '960MB')
        122880
        """
        if self.memory_unit(name):
            if value.endswith('kB'):
                internal = int(value.rstrip('kB'))
                internal = internal / self.memory_divisor(name)
            elif value.endswith('MB'):
                internal = int(value.rstrip('MB'))
                internal = internal * KB_PER_MB / self.memory_divisor(name)
            elif value.endswith('GB'):
                internal = int(value.rstrip('GB'))
                internal = internal * KB_PER_GB / self.memory_divisor(name)
            else:
                internal = int(value)        
        else:        
            internal = int(value)        
    
        return internal
  
    def parse(self, name, value):
        """
        Return a string representing the internal value this setting
        would be parsed into.  This includes converting memory values
        into their internal integer representation.

        TODO It might be helpful to eventually handle all the boolean
        representations that the PostgreSQL GUC code understands,
        outputting in standard form
        """
        if self.vartype(name) == "integer":
            return str(self.parse_int(name, value))
        return value

  
def wizard_tune(config, options, settings):
    """
    We expect the following options are passed into here:
    
    db_type:  Defaults to mixed
    connections:  If missing, will set based on db_type
    totalMemory:  If missing, will detect
    """
    db_type = options.db_type.lower()
  
    # Save all settings to be updated as (setting,value) dictionary values
    s = {}
    try:
        s['max_connections'] = {'web':200, 'oltp':300, 'dw':20, 'mixed':80, 'desktop':5}[db_type]
    except KeyError:
        print "Error:  unexpected setting for db_type"
        sys.exit(1)
  
    # Now that we've screened for that, we know we've got a good db_type and
    # don't have to wrap the rest of these settings in an try block
  
    # Allow overriding the maximum connections
    if options.connections != None:
        s['max_connections'] = options.connections
  
    # Estimate memory on this system via parameter or system lookup
    total_memory = options.total_memory
    if total_memory is None:
        total_memory = total_mem()
    if total_memory is None:
        print "Error:  total memory not specified and unable to detect"
        sys.exit(1)
    
    # Memory allocation
    # Extract some values just to make the code below more compact
    # The base unit for memory types is the kB, so scale system memory to that
    mem = int(total_memory) / KB
    con = int(s['max_connections'])
  
    if total_memory >= (256 * MB):
        if False:  # platform.system()=="Windows"
            # TODO Adjust shared_buffers for Windows
            pass
        else:
            s['shared_buffers'] = {'web':mem / 4, 'oltp':mem / 4, 'dw':mem / 4,
                                   'mixed':mem / 4, 'desktop':mem / 16}[db_type]
    
        s['effective_cache_size'] = {'web':mem * 3 / 4, 'oltp':mem * 3 / 4, 'dw':mem * 3 / 4,
                                     'mixed':mem * 3 / 4, 'desktop':mem / 4}[db_type]
    
        s['work_mem'] = {'web':mem / con, 'oltp':mem / con,'dw':mem / con / 2,
                         'mixed':mem / con / 2,'desktop':mem / con / 6}[db_type]
    
        s['maintenance_work_mem'] = {'web':mem / 16, 'oltp':mem / 16,'dw':mem / 8,
                                   'mixed':mem / 16,'desktop':mem / 16}[db_type]
        
        # Cap maintenence RAM at 1GB on servers with lots of memory
        # (Remember that the setting is in terms of kB here)
        if s['maintenance_work_mem'] > (1 * MB):
            s['maintenance_work_mem'] = 1 * MB
  
    else:
        # TODO HINT about this tool not being optimal for low memory systems
        pass
  
    # Checkpoint parameters
    s['checkpoint_segments'] = {'web':8, 'oltp':16, 'dw':64,
                                'mixed':16, 'desktop':3}[db_type]
  
    s['checkpoint_completion_target'] = {'web':0.7, 'oltp':0.9, 'dw':0.9,
                                         'mixed':0.9, 'desktop':0.5}[db_type]
  
    s['wal_buffers'] = 512 * s['checkpoint_segments']
  
    # Paritioning and statistics
    s['constraint_exclusion'] = {'web':'off', 'oltp':'off', 'dw':'on', 
                                 'mixed':'on', 'desktop':'off'}[db_type]
  
    s['default_statistics_target'] = {'web':10, 'oltp':10, 'dw':100, 
                                      'mixed':50, 'desktop':10}[db_type]  
    for key in s.keys():
        value = s[key]
        # TODO Make this logic part of the config class, so this
        # function doesn't need to be passed settings
        if settings.memory_unit(key):
            value = binaryround(s[key])
        # TODO Add show method to config class for similar reasons
        config.update_setting(key, settings.show(key, value))


def read_options(program_args):
    parser = optparse.OptionParser(usage="usage: %prog [options]",
                                   version="1.0",
                                   conflict_handler="resolve")
      
    parser.add_option('-i', '--input-config', dest="input_config", default=None,
                      help="Input configuration file")
  
    parser.add_option('-o', '--output-config', dest="output_config", default=None, 
                      help="Output configuration file, defaults to standard output")
    
    parser.add_option('-M', '--memory', dest="total_memory", default=None, 
                      help="Total system memory, will attempt to detect if unspecified")
  
    parser.add_option('-T', '--type', dest="db_type", default="Mixed", 
                      help="Database type, defaults to Mixed, valid options are DW, OLTP, Web, Mixed, Desktop")
  
    parser.add_option('-c', '--connections', dest="connections", default=None, 
                      help="Maximum number of expected connections, default depends on database type")
  
    parser.add_option('-D', '--debug', action="store_true", dest="debug",
                      default="False", help="Enable debugging mode")
  
    parser.add_option('-S', '--settings', dest="settings_dir", default=None, 
                      help="Directory where settings data files are located at.  Defaults to the directory where the script is being run from")

    parser.add_option('--doctest', help='run doctests', action='store_true')
    options, args = parser.parse_args(program_args)
    
    if options.debug == True:
        print "Command line options:  ",options
        print "Command line arguments:  ",args
    
    return options, args, parser


def main(program_args):
    options, args, parser = read_options(program_args) 

    if options.doctest:
        import doctest
        doctest.testmod()
        return(0)
      
    configFile = options.input_config
    if configFile is None:
        print >> sys.stderr,"Can't do anything without an input config file; try --help"
        parser.print_help()
        return(1)
      
    config = PGConfigFile(configFile)
    config.read_config_file()
  
    if options.debug == True:  
        config.debug_print_input()
        print
        config.debug_print_settings()
  
    if options.settings_dir is None:
        options.settings_dir = os.path.abspath(os.path.dirname(sys.argv[0]))
  
    settings = PGSettings(options.settings_dir)
    settings.read_config_file()
    config.store_settings(settings)
  
    wizard_tune(config, options, settings)
    
    output_file_name = options.output_config
    if output_file_name is None:  
        fout = sys.stdout
    else:
        fout = open(output_file_name, 'w')
  
    config.write(fout)

if __name__ == '__main__':
    sys.exit(main(sys.argv))