#!/usr/bin/env python
# Copyright 2017 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""A script to merge multiple source xml files into a single histograms.xml."""

import argparse
import os
import xml.dom.minidom

import setup_modules  # pylint: disable=unused-import

import chromium_src.tools.metrics.common.xml_utils as xml_utils
import chromium_src.tools.metrics.histograms.expand_owners as expand_owners
import chromium_src.tools.metrics.histograms.histogram_configuration_model as histogram_configuration_model
import chromium_src.tools.metrics.histograms.histogram_paths as histogram_paths
import chromium_src.tools.metrics.histograms.populate_enums as populate_enums


def GetElementsByTagName(trees, tag, depth=2):
  """Gets all elements with the specified tag from a set of DOM trees.

  Args:
    trees: A list of DOM trees.
    tag: The tag of the elements to find.
    depth: The depth in the trees by which a match should be found.

  Returns:
    A list of DOM nodes with the specified tag.
  """
  iterator = xml_utils.IterElementsWithTag
  return list(e for t in trees for e in iterator(t, tag, depth))


def CombineEnumsSections(doc, trees):
  """Combines multiple <enums> from the passed in DOM trees into one.

  If trees contain ukm events, populates a list of ints to the
  "UkmEventNameHash" enum where each value is a ukm event name hash truncated
  to 31 bits and each label is the corresponding event name.

  Args:
    doc: The document where the new single <enums> section will be created.
    trees: A list of DOM trees.

  Returns:
    A single <enums> DOM node.
  """
  enums_node = doc.createElement('enums')
  # Pass depth=3 as default depth=2 won't find enum tags that are 3 levels deep.
  for enum in GetElementsByTagName(trees, 'enum', depth=3):
    xml.dom.minidom._append_child(enums_node, enum)

  ukm_events = GetElementsByTagName(
    GetElementsByTagName(trees, 'ukm-configuration'), 'event'
  )
  if ukm_events:
    populate_enums.PopulateEnumsWithUkmEvents(doc, enums_node, ukm_events)
  return enums_node


def CombineHistogramsSorted(doc, trees):
  """Sorts histograms related nodes by name and returns the combined nodes.

  This function sorts nodes including <histogram>, <variant> and
  <histogram_suffix>. Then it returns one <histograms> that contains the
  sorted <histogram> and <variant> nodes and the other <histogram_suffixes_list>
  node containing all <histogram_suffixes> nodes.

  Args:
    doc: The document to create the node in.
    trees: A list of DOM trees.

  Returns:
    A list containing the combined <histograms> node and the combined
    <histogram_suffix_list> node.
  """
  # Create the combined <histograms> tag.
  combined_histograms = doc.createElement('histograms')

  def SortByLowerCaseName(node):
    return node.getAttribute('name').lower()

  variants_nodes = GetElementsByTagName(trees, 'variants', depth=3)
  sorted_variants = sorted(variants_nodes, key=SortByLowerCaseName)

  histogram_nodes = GetElementsByTagName(trees, 'histogram', depth=3)
  sorted_histograms = sorted(histogram_nodes, key=SortByLowerCaseName)

  for variants in sorted_variants:
    # Use unsafe version of `appendChild` function here because the safe one
    # takes a lot longer (10000x) to append all children. The unsafe version
    # is ok here because:
    #   1. the node to be appended is a clean node.
    #   2. The unsafe version only does fewer checks but not changing any
    #     behavior and it's documented to be usable if performance matters.
    #     See https://github.com/python/cpython/blob/2.7/Lib/xml/dom/minidom.py#L276.
    xml.dom.minidom._append_child(combined_histograms, variants)

  for histogram in sorted_histograms:
    xml.dom.minidom._append_child(combined_histograms, histogram)

  # Create the combined <histogram_suffixes_list> tag.
  combined_histogram_suffixes_list = doc.createElement(
    'histogram_suffixes_list'
  )

  histogram_suffixes_nodes = GetElementsByTagName(
    trees, 'histogram_suffixes', depth=3
  )
  sorted_histogram_suffixes = sorted(
    histogram_suffixes_nodes, key=SortByLowerCaseName
  )

  for histogram_suffixes in sorted_histogram_suffixes:
    xml.dom.minidom._append_child(
      combined_histogram_suffixes_list, histogram_suffixes
    )

  return [combined_histograms, combined_histogram_suffixes_list]


def MakeNodeWithChildren(doc, tag, children):
  """Creates a DOM node with specified tag and child nodes.

  Args:
    doc: The document to create the node in.
    tag: The tag to create the node with.
    children: A list of DOM nodes to add as children.

  Returns:
    A DOM node.
  """
  node = doc.createElement(tag)
  for child in children:
    node.appendChild(child)
  return node


def MergeTrees(trees, should_expand_owners):
  """Merges a list of histograms.xml DOM trees.

  Args:
    trees: A list of histograms.xml DOM trees.
    should_expand_owners: Whether we want to expand owners for histograms.

  Returns:
    A merged DOM tree.
  """
  doc = xml.dom.minidom.Document()
  doc.appendChild(
    MakeNodeWithChildren(
      doc,
      'histogram-configuration',
      [CombineEnumsSections(doc, trees)]
      +
      # Sort the <histogram> and <histogram_suffixes> nodes by name and
      # return the combined nodes.
      CombineHistogramsSorted(doc, trees),
    )
  )
  # Only perform fancy operations after |doc| becomes stable. This helps improve
  # the runtime performance.
  if should_expand_owners:
    for histograms in doc.getElementsByTagName('histograms'):
      expand_owners.ExpandHistogramsOWNERS(histograms)
  return doc


def _AddComponentFromMetadataFile(tree, filename):
  """Adds the component from the metadata file to the DOM tree.

  Args:
    tree: A histogram.xml DOM tree.
    filename: The name of the metadata file.

  Returns:
    The updated tree with the component (optionally) added.
  """
  component = expand_owners.ExtractComponentViaDirmd(os.path.dirname(filename))
  if component:
    histograms = tree.getElementsByTagName('histograms')
    if histograms:
      iter_matches = xml_utils.IterElementsWithTag
      for histogram in iter_matches(histograms[0], 'histogram'):
        expand_owners.AddHistogramComponent(histogram, component)
  return tree


def _BuildDOMTreeWithComponentMetadata(filename_or_file):
  """Builds the DOM tree for the given file.

  Args:
    filename_or_file: The string filename or the file handle for histograms.xml.

  Returns:
    The histograms.xml DOM tree with (optional) component metadata.
  """
  tree = xml.dom.minidom.parse(filename_or_file)
  if isinstance(filename_or_file, str):
    # If we can find a metadata file in the same directory, we try to extract
    # a component from it.
    metadata_filename = os.path.join(
      os.path.dirname(filename_or_file), 'DIR_METADATA'
    )
    if os.path.exists(metadata_filename):
      return _AddComponentFromMetadataFile(tree, metadata_filename)
  return tree


def MergeFiles(
  filenames=[], files=[], expand_owners_and_extract_components=False
):
  """Merges a list of histograms.xml files.

  Args:
    filenames: A list of histograms.xml filenames.
    files: A list of histograms.xml file-like objects.
    expand_owners_and_extract_components: Whether we want to expand owners and
      extract components. By default, it's false because most of the callers
      don't care about the owners or components for each metadata.

  Returns:
    A merged DOM tree.
  """
  # minidom.parse() takes both files and filenames:
  all_files = files + filenames
  trees = [
    _BuildDOMTreeWithComponentMetadata(f)
    if expand_owners_and_extract_components
    else xml.dom.minidom.parse(f)
    for f in all_files
  ]
  return MergeTrees(
    trees, should_expand_owners=expand_owners_and_extract_components
  )


def PrettyPrintMergedFiles(filenames=[], files=[]):
  return histogram_configuration_model.PrettifyTree(
    MergeFiles(
      filenames=filenames,
      files=files,
      expand_owners_and_extract_components=True,
    )
  )


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('--output', required=True)
  args = parser.parse_args()
  with open(args.output, 'w', encoding='utf-8', newline='\n') as f:
    # This is run by
    # https://source.chromium.org/chromium/chromium/src/+/main:tools/metrics/BUILD.gn;drc=573e48309695102dec2da1e8f806c18c3200d414;l=5
    # to send the merged histograms.xml to the server side. Providing |UKM_XML|
    # here is not to merge ukm.xml but to populate `UkmEventNameHash` enum
    # values.
    f.write(
      PrettyPrintMergedFiles(
        histogram_paths.ALL_XMLS + [histogram_paths.UKM_XML]
      )
    )


if __name__ == '__main__':
  main()