#!/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)