#!/usr/bin/env python

import os, sys
import happybase
import logging

from boto.sqs import connect_to_region as sqs_connect
from boto.sqs.jsonmessage import JSONMessage
from boto.s3 import connect_to_region as s3_connect
from boto.s3.key import Key

from multiprocessing import Process as TaskClass
from multiprocessing import JoinableQueue as Queue

import signal
import random

from collections import deque
import simplejson as json


logger = logging.getLogger(__name__)

# Following params need to be adjusted based on payload size, bandwidth etc
MAX_ROWS_IN_FLIGHT = 4096
TASK_QUEUE_SIZE = MAX_ROWS_IN_FLIGHT * 4


class HBaseSource:
  def __init__(self, addr, row_range, max_rows = 2048, batch_size = 256):
    self.thrift_addr = addr
    self.start_row, self.end_row = row_range
    self.max_rows = max_rows
    self.batch_size = batch_size
      
  def items(self):
    prev_last_read_key = None
    curr_last_read_key = self.start_row
    end_row = self.end_row

    while True:
      src_tbl = happybase.Connection(random.choice(self.thrift_addr)).table('crash_reports')

      nrows = 0

      try:
        logger.debug('fetch %d rows of data via thrift', self.max_rows)

        # scan fetches rows with key in the range [row_start, row_stop)
        # this necessitates the check for repeating keys as stopping condition
        # 
        logger.info("scan start")
        data = deque(src_tbl.scan(row_start = curr_last_read_key, 
                                  row_stop = end_row,
                                  columns = ['raw_data', 'processed_data', 'meta_data'],
                                  limit = self.max_rows, 
                                  batch_size = self.batch_size))
        logger.info("scan end %d rows starting at %s", len(data), data[0][0])
        while True:
          if not data:
            break

          key, val = data.popleft()
          if (key == prev_last_read_key):
            # last record from previous batch should be ignored
            continue

          yield key, val
          nrows += 1

          prev_last_read_key = curr_last_read_key
          curr_last_read_key = key
          
        logger.debug('read %d rows of data from hbase ending at %s', nrows, curr_last_read_key)
        if nrows < self.max_rows:
          print >> sys.stderr, "end of range. exiting"
          break

      except happybase.hbase.ttypes.IOError:
        logger.exception('caught exception. retrying.')

      except Exception:
        logger.exception('unrecoverable exception.')
        raise

class SourceWorker(TaskClass):
  def __init__(self, queue, source_config):
    TaskClass.__init__(self)

    self.source = HBaseSource(*source_config)
    self.queue = queue

  def run(self):
    num_rows_written = 0
    total_size_written = 0
    s3_path_tmpl = '{env}/v1/{ftype}/{uuid}'
    env = 'stage'

    for key, cols in self.source.items():
      dump_names = []
      for j in cols.keys():

        suffix = get_suffix(j)
        if not suffix:
          #logger.info('column %s ignored for key %s', j, key)
          continue

        if j.startswith('raw_data'):
          dump_names.append(suffix)

        # crashstats/stage/v1/
        # format  {{bucket}}/{{prefix}}/{{version}}/{{crash_type}}/{{crash_id}}
        skey = s3_path_tmpl.format(env = env, 
                                   uuid = key[7:],
                                   ftype = suffix)

        self.queue.put((skey, cols[j]))

        total_size_written += len(cols[j])

      self.queue.put((s3_path_tmpl.format(env = env,
                                          uuid = key[7:],
                                          ftype = 'dump_names'),
                      json.dumps(dump_names)))

      num_rows_written += 1

      if ((num_rows_written % 1000) == 0):
        logger.info("wrote %d rows, at %s", num_rows_written, key)
        logger.warn("qsize is %d", self.queue.qsize())

    print >> sys.stderr, "SourceWorker DONE", num_rows_written, total_size_written

class S3Worker(TaskClass):
  def __init__(self, s3_region, s3_bucket, task_queue, result_queue):
    signal.signal(signal.SIGINT, signal.SIG_IGN)

    TaskClass.__init__(self)
    self.task_queue = task_queue
    self.result_queue = result_queue
    self.s3_region = s3_region
    self.s3_bucket = s3_bucket

  def setup_s3(self):
    self.s3 = s3_connect(self.s3_region)
    self.bucket = self.s3.get_bucket(self.s3_bucket)

  def write_to_s3(self, key, payload):
    k = Key(self.bucket)
    k.key = key
    k.set_contents_from_string(payload)

  def run(self):
    self.setup_s3()

    while True:
      kv = self.task_queue.get()
      
      if kv is None:
        print >> sys.stderr, '%s: Exiting' % self.name
        self.task_queue.task_done()
        break

      k, v = kv
      self.write_to_s3(k, v)
      self.task_queue.task_done()
    return

def get_suffix(colname):
  suffix_map = {
    'processed_data:json' : 'processed_crash',
    'raw_data:dump' : 'dump',
    'meta_data:json' : 'raw_crash',
    'raw_data:upload_file_minidump_browser' : 'upload_file_minidump_browser',
    'raw_data:upload_file_minidump_flash1' : 'upload_file_minidump_flash1',
    'raw_data:upload_file_minidump_flash2' : 'upload_file_minidump_flash2'
    }

  if colname in suffix_map:
    return suffix_map[colname]
  elif colname.startswith('raw_data'):
    return colname.split(':', 1)[1]
  else:
    return None


def main(num_workers = 64):
  if len(sys.argv) != 3:
    show_usage_and_quit()

  queue = Queue(TASK_QUEUE_SIZE)

  # start s3 workers
  workers = [S3Worker('us-west-2', 'crashstats', queue, None) 
             for i in xrange(num_workers)]

  for i in workers:
    i.start()

  thrift_hosts = sys.argv[1].split(',')
  date = sys.argv[2]

  # start hbase workers
  key_ranges = []
  for i in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f']:
    key_ranges.append(('%s%s%s' % (i, date, i), '%s%s%sg' % (i, date, i)))

  num_hbase_workers = 1

  for i in xrange(0, len(key_ranges), num_hbase_workers):
    src_workers = []
    krng = key_ranges[i : (i + num_hbase_workers)]

    for j in range(len(krng)):
      src_workers.append(SourceWorker(queue, (thrift_hosts, krng[j])))

    for w in src_workers:
      print "starting src worker", w
      w.start()

    for w in src_workers:
      w.join()

  for i in workers:
    queue.put(None)

  queue.join()
  
def show_usage_and_quit():
  print >> sys.stderr, "Usage: %s hosts('host1,host2,host3') date(YYMMDD)" % (sys.argv[0])
  sys.exit(2)
  

if __name__ == '__main__':
    logging.basicConfig(format = '%(asctime)s %(name)s:%(levelname)s: %(message)s',
                        level = logging.INFO)

    main()