#!/usr/bin/python
#
# Copyright 2011 Google Inc. All Rights Reserved.
"""Script to divide and merge profiles."""

import copy
import optparse
import os
import pickle
import re
import sys
import tempfile

import build_chrome_browser
import lock_machine
import run_tests
from utils import command_executer
from utils import logger


class ProfileMerger:

  def __init__(self, inputs, output, chunk_size, merge_program, multipliers):
    self._inputs = inputs
    self._output = output
    self._chunk_size = chunk_size
    self._merge_program = merge_program
    self._multipliers = multipliers
    self._ce = command_executer.GetCommandExecuter()
    self._l = logger.GetLogger()

  def _GetFilesSetForInputDir(self, input_dir):
    output_file = tempfile.mktemp()
    command = "find %s -name '*.gcda' -o -name '*.imports' > %s" % (input_dir,
                                                                    output_file)
    self._ce.RunCommand(command)
    files = open(output_file, 'r').read()
    files_set = set([])
    for f in files.splitlines():
      stripped_file = f.replace(input_dir, '', 1)
      stripped_file = stripped_file.lstrip('/')
      files_set.add(stripped_file)
    return files_set

  def _PopulateFilesSet(self):
    self._files_set = set([])
    for i in self._inputs:
      current_files_set = self._GetFilesSetForInputDir(i)
      self._files_set.update(current_files_set)

  def _GetSubset(self):
    ret = []
    for i in range(self._chunk_size):
      if not self._files_set:
        break
      ret.append(self._files_set.pop())
    return ret

  def _CopyFilesTree(self, input_dir, files, output_dir):
    for f in files:
      src_file = os.path.join(input_dir, f)
      dst_file = os.path.join(output_dir, f)
      if not os.path.isdir(os.path.dirname(dst_file)):
        command = 'mkdir -p %s' % os.path.dirname(dst_file)
        self._ce.RunCommand(command)
      command = 'cp %s %s' % (src_file, dst_file)
      self._ce.RunCommand(command)

  def _DoChunkMerge(self, current_files):
    temp_dirs = []
    for i in self._inputs:
      temp_dir = tempfile.mkdtemp()
      temp_dirs.append(temp_dir)
      self._CopyFilesTree(i, current_files, temp_dir)
    # Now do the merge.
    command = ('%s --inputs=%s --output=%s' %
               (self._merge_program, ','.join(temp_dirs), self._output))
    if self._multipliers:
      command = ('%s --multipliers=%s' % (command, self._multipliers))
    ret = self._ce.RunCommand(command)
    assert ret == 0, '%s command failed!' % command
    for temp_dir in temp_dirs:
      command = 'rm -rf %s' % temp_dir
      self._ce.RunCommand(command)

  def DoMerge(self):
    self._PopulateFilesSet()
    while True:
      current_files = self._GetSubset()
      if not current_files:
        break
      self._DoChunkMerge(current_files)


def Main(argv):
  """The main function."""
  # Common initializations
  ###  command_executer.InitCommandExecuter(True)
  command_executer.InitCommandExecuter()
  l = logger.GetLogger()
  ce = command_executer.GetCommandExecuter()
  parser = optparse.OptionParser()
  parser.add_option('--inputs',
                    dest='inputs',
                    help='Comma-separated input profile directories to merge.')
  parser.add_option('--output', dest='output', help='Output profile directory.')
  parser.add_option('--chunk_size',
                    dest='chunk_size',
                    default='50',
                    help='Chunk size to divide up the profiles into.')
  parser.add_option('--merge_program',
                    dest='merge_program',
                    default='/home/xur/bin/profile_merge_v15.par',
                    help='Merge program to use to do the actual merge.')
  parser.add_option('--multipliers',
                    dest='multipliers',
                    help='multipliers to use when merging. (optional)')

  options, _ = parser.parse_args(argv)

  if not all([options.inputs, options.output]):
    l.LogError('Must supply --inputs and --output')
    return 1

  try:
    pm = ProfileMerger(
        options.inputs.split(','), options.output, int(options.chunk_size),
        options.merge_program, options.multipliers)
    pm.DoMerge()
    retval = 0
  except:
    retval = 1
  finally:
    print 'My work is done...'
  return retval


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