#!/usr/bin/env python

import sys, os, time, traceback
import whisper
import threading
import re
from optparse import OptionParser
from Queue import Queue

exited_workers = [] #janky but will work for this simple case. see use near bottom os script
now = int(time.time())

class ExitMessage(object): pass

class Messenger(object):
  def __init__(self,msgQ):
    self.msgQ = msgQ
 
  def write(self,data):
    self.msgQ.put(data)

def messenger(msgQ):
  NAME = "[Message Thread]"
  while True:
    msg = msgQ.get(True)
    if msg == ExitMessage:
      sys.__stdout__.write( "%s exiting.\n" % (NAME))
      return
    try:  sys.__stdout__.write(msg)
    except: pass

def worker_loop(tid,workQ):
  while True:
    NAME = "[Thread:%s]" % tid
    data = workQ.get(True)
    if data == ExitMessage:
      print "%s exiting." % NAME
      exited_workers.append(tid)
      return
    path,archives,options = data
    try:
      worker(path,archives,options,tid)
      print "%s Resize complete: %s" % (NAME,path)
    except SystemExit, exc:
      pass
    except Exception, exc:
      print "%s Error `%s` was encountered processing %s" %(NAME,exc.message,path)

def worker(path,new_archives,options,tid=None):
  global now
  if tid == None: NAME = ""
  else: NAME = "[Thread:%s]" % tid
  info = whisper.info(path)
  old_archives = info['archives']
  # sort by precision, lowest to highest
  old_archives.sort(key=lambda a: a['secondsPerPoint'], reverse=True)

  if options.xFilesFactor is None:
    xff = info['xFilesFactor']
  else:
    xff = options.xFilesFactor
  
  if options.aggregationMethod is None:
    aggregationMethod = info['aggregationMethod']
  else:
    aggregationMethod = options.aggregationMethod
  
  print '%s Retrieving all data from the archives' % NAME
  for archive in old_archives:
    fromTime = now - archive['retention'] + archive['secondsPerPoint']
    untilTime = now
    timeinfo,values = whisper.fetch(path, fromTime, untilTime)
    archive['data'] = (timeinfo,values)
  
  if options.newfile is None:
    tmpfile = path + '.tmp'
    if os.path.exists(tmpfile):
      print '%s Removing previous temporary database file: %s' % (NAME,tmpfile)
      os.unlink(tmpfile)
    newfile = tmpfile
  else:
    newfile = options.newfile
  
  print '%s Creating new whisper database: %s' % (NAME,newfile)
  whisper.create(newfile, new_archives, xFilesFactor=xff, aggregationMethod=aggregationMethod)
  size = os.stat(newfile).st_size
  print '%s Created: %s (%d bytes)' % (NAME,newfile,size)
  
  print '%s Migrating data...' % NAME
  for archive in old_archives:
    timeinfo, values = archive['data']
    datapoints = zip( range(*timeinfo), values )
    datapoints = filter(lambda p: p[1] is not None, datapoints)
    whisper.update_many(newfile, datapoints)
  
  if options.newfile is not None:
    sys.exit(0)
  
  backup = path + '.bak'
  print '%s Renaming old database to: %s' % (NAME,backup)
  os.rename(path, backup)
  
  try:
    print '%s Renaming new database to: %s' % (NAME,path)
    os.rename(tmpfile, path)
  except:
    traceback.print_exc()
    print '\n%s Operation failed, restoring backup' % NAME
    os.rename(backup, path)
    sys.exit(1)
  
  if options.nobackup:
    print "%s Unlinking backup: %s" % (NAME,backup)
    os.unlink(backup)


def collect_arguments():
  option_parser = OptionParser(
    usage='''%prog path timePerPoint:timeToStore [timePerPoint:timeToStore]*

timePerPoint and timeToStore specify lengths of time, for example:

60:1440      60 seconds per datapoint, 1440 datapoints = 1 day of retention
15m:8        15 minutes per datapoint, 8 datapoints = 2 hours of retention
1h:7d        1 hour per datapoint, 7 days of retention
12h:2y       12 hours per datapoint, 2 years of retention
''')

  option_parser.add_option(
    '--xFilesFactor', default=None,
    type='float', help="Change the xFilesFactor")
  option_parser.add_option(
    '--aggregationMethod', default=None,
    type='string', help="Change the aggregation function (%s)" %
    ', '.join(whisper.aggregationMethods))
  option_parser.add_option(
    '--force', default=False, action='store_true',
    help="Perform a destructive change")
  option_parser.add_option(
    '--newfile', default=None, action='store',
    help="Create a new database file without removing the existing one")
  option_parser.add_option(
    '--nobackup', action='store_true',
    help='Delete the .bak file after successful execution')
  option_parser.add_option(
    '--threads', type="int",default=1,
    help='Number of threads to use if <path> points at a directory instead of a file. Directory recursion is implied.')
  option_parser.add_option(
    '--include', type="string",default=".*",
    help='Regular expresion to use for deciding which whisper files should be included in this resize. by default all are included.')
  option_parser.add_option(
    '--exclude', type="string",default='^$',
    help='Regular expresion to use for deciding which whisper files should be exluded in this resize. by default none are excluded.')


  (options, args) = option_parser.parse_args()
  if len(args) < 2:
    option_parser.print_usage()
    sys.exit(1)
  return (options,args)

def get_queues():
  return Queue(),Queue()

def start_messenger(msgQ):
  msg_writer = Messenger(msgQ)
  sys.stdout = msg_writer
    
def start_workers(options,msgQ,workQ):
  workers = []
  mt = threading.Thread(target=messenger,args=(msgQ,))
  mt.start()

  for i in range(options.threads):
    w = threading.Thread(target=worker_loop,args=(i,workQ))
    w.start()
    workers.append(w)
  return workers

def get_work(workQ,path,new_archives,options):
  inc_regex = re.compile(options.include)
  exc_regex = re.compile(options.exclude)
  for root,dirs,files in os.walk(path):
    for f in files:
      if not inc_regex.match(f): continue
      if exc_regex.match(f): continue
      full_path = os.path.join(root,f)
      msg = (full_path,new_archives,options)
      workQ.put( msg )

def exit_workers(workers,workQ,msgQ):
  for i in workers: workQ.put(ExitMessage)
  while len(exited_workers) != len(workers): time.sleep(1)
  msgQ.put(ExitMessage)

def main():
  options,args = collect_arguments()
  
  path = args[0]
  new_archives = [whisper.parseRetentionDef(retentionDef)
                  for retentionDef in args[1:]]
    
  if os.path.isfile(path):
    worker(path,new_archives,options)
  else:
    msgQ,workQ = get_queues()
    start_messenger(msgQ)
    workers = start_workers(options,msgQ,workQ)
    get_work(workQ,path,new_archives,options)
    exit_workers(workers,workQ,msgQ)

main()