# Copyright 2019 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import collections
import copy
import itertools
import os
from collections import namedtuple

# A simple container for associating a shard with a swarming task request. We
# need this to be able to match up the summary.json produced by the task with
# the tests.json manifest for the shard, so that the "name" and "gn_label"
# fields in summary.json can be correctly populated using tests.json for QEMU
# shards.
# TODO(40190): Once done, maybe we can remove the shards and just keep the task
# requests.
#
# Attributes:
#   shard (api.testsharder.Shard): The input shard containing a list of tests
#     to run.
#   task_request (api.swarming.TaskRequest): The swarming request
#     corresponding to the shard.
ShardTaskRequest = collections.namedtuple(
    'ShardTaskRequest', field_names=('shard', 'task_request'))


class TestOrchestrationInputs(object):
  """Data used to orchestrate testing."""

  # Constant shared between recipes that produce and consume.
  # Property that holds the isolate hash of the test orchestration inputs.
  HASH_PROPERTY = 'test_orchestration_inputs_hash'

  # IsolatedLayout describes the directory structure to use when isolating
  # TestOrchestrationInputs. Equivalent layouts must be used to isolate
  # and download.
  IsolatedLayout = namedtuple(
      'IsolatedLayout',
      (
          'covargs',
          'ids',
          'llvm_cov',
          'llvm_profdata',
          'llvm_symbolizer',
          'minfs',
          'shards',
          'symbolize_tool',
          'task_requests',
          # TODO(garymm,joshuaseaton): Once deprecated testing codepath is eliminated we
          # can remove the tests file from this class.
          'tests'),
  )

  DEFAULT_ISOLATED_LAYOUT = IsolatedLayout(
      covargs='covargs',
      ids='ids.txt',
      llvm_cov='llvm_cov',
      llvm_profdata='llvm_profdata',
      llvm_symbolizer='llvm-symbolizer',
      minfs='minfs',
      shards='shards.json',
      symbolize_tool='symbolize_tool',
      task_requests='task_requests.json',
      tests='tests.json',
  )

  def __init__(self, covargs, ids, llvm_cov, llvm_profdata, llvm_symbolizer,
               minfs, symbolize_tool, shard_requests, tests_file):
    # TODO(garymm): Take in an API object here rather than in the
    # member functions, and hide this constructor behind a function
    # in build/api.py. That will prevent the calling recipes from
    # needing to depend on all the transitive dependencies of this
    # code. See comment in docstring of TestOrchestrationInputs.isolate()
    # for more on that.
    self.covargs = covargs
    self.ids = ids
    self.llvm_cov = llvm_cov
    self.llvm_profdata = llvm_profdata
    self.llvm_symbolizer = llvm_symbolizer
    self.minfs = minfs
    self.symbolize_tool = symbolize_tool
    self.shard_requests = shard_requests
    self.tests_file = tests_file
    self._layout = self.DEFAULT_ISOLATED_LAYOUT
    self._shards_file = None
    self._task_requests_file = None

  @classmethod
  def from_build_results(cls, build_results, shard_requests):
    return cls(
        build_results.tool('covargs'), build_results.ids,
        build_results.tool('llvm-cov'), build_results.tool('llvm-profdata'),
        build_results.llvm_symbolizer, build_results.minfs,
        build_results.tool('symbolize'), shard_requests,
        build_results.tests_file)

  @classmethod
  def download(cls, api, isolated_hash):
    """Downloads an isolated containing TestOrchestrationInputs

      Args:
        isolated_hash (string): The isolated hash to fetch.

      Returns:
        A TestOrchestrationInputs object.
      """
    with api.step.nest('download test orchestration inputs'):
      return cls._Isolator(api,
                           cls.DEFAULT_ISOLATED_LAYOUT).download(isolated_hash)

  def isolate(self, api):
    """Uploads to isolate.

    Args:
      api: magic recipe engine API object.
           Note: because of the way DEPs stuff works, the api object gets the
           deps of the calling recipe, NOT the deps listed in
           build/__init__.py.

    Returns:
      (str) Isolated hash.
    """
    with api.step.nest('isolate test orchestration inputs'):
      shards_data = []
      task_requests_data = []
      for shard_request in self.shard_requests:
        shards_data.append(shard_request.shard.render_to_jsonish())
        task_requests_data.append(shard_request.task_request.to_jsonish())
      temp_dir = api.path.mkdtemp('test-orchestration-inputs')
      self._shards_file = temp_dir.join(self._layout.shards)
      self._task_requests_file = temp_dir.join(self._layout.task_requests)
      api.file.write_json(
          'write %s' % self._layout.shards,
          data=shards_data,
          dest=self._shards_file,
          indent=2)
      api.file.write_json(
          'write %s' % self._layout.task_requests,
          data=task_requests_data,
          dest=self._task_requests_file,
          indent=2)
      return self._Isolator(api, self._layout).isolate(self)

  class _Isolator(object):
    """Isolator isolates and downloads TestOrchestrationInputs"""

    def __init__(self, api, layout):
      self._api = api
      self._layout = layout

    def download(self, isolated_hash):
      """Downloads an isolated containing TestOrchestrationInputs

      Args:
        isolated_hash (string): The isolated hash to fetch.

      Returns:
        A TestOrchestrationInputs object.
      """
      download_dir = self._api.path.mkdtemp('test-orchestration-inputs')
      self._api.isolated.download(
          'download test orchestration inputs',
          isolated_hash=isolated_hash,
          output_dir=download_dir)

      shards_file = download_dir.join(self._layout.shards)
      shards_in_json = self._api.json.read('load test shards',
                                           shards_file).json.output
      shards_in_json = shards_in_json or []
      task_requests_file = download_dir.join(self._layout.task_requests)
      task_requests_in_json = self._api.json.read(
          'load task requests', task_requests_file).json.output
      task_requests_in_json = task_requests_in_json or []
      assert len(shards_in_json) == len(task_requests_in_json), '%d vs %d' % (
          len(shards_in_json), len(task_requests_in_json))
      shard_requests = []
      for shard_json, task_request_json in itertools.izip(
          shards_in_json, task_requests_in_json):
        shard_requests.append(
            ShardTaskRequest(
                self._api.testsharder.Shard.from_jsonish(shard_json),
                self._api.swarming.task_request_from_jsonish(
                    task_request_json)))

      return TestOrchestrationInputs(
          download_dir.join(self._layout.covargs),
          download_dir.join(self._layout.ids),
          download_dir.join(self._layout.llvm_cov),
          download_dir.join(self._layout.llvm_profdata),
          download_dir.join(self._layout.llvm_symbolizer),
          download_dir.join(self._layout.minfs),
          download_dir.join(self._layout.symbolize_tool), shard_requests,
          download_dir.join(self._layout.tests))

    def isolate(self, inputs):
      """Isolates BuildArtifacts according to this Isolator's layout.

      Args:
        inputs (TestOrchestrationInputs):
          The TestOrchestrationInputs to isolate.

      Returns:
        A string isolated ID hash.
      """
      root = self._api.path.mkdtemp('test_orchestration_inputs')
      symlink_tree = self._api.file.symlink_tree(root)
      self._symlink(symlink_tree, self._layout, inputs, root)
      symlink_tree.create_links('create_links')
      isolated = self._api.isolated.isolated(root)
      isolated.add_dir(root)
      return isolated.archive('isolate')

    def _symlink(self, symlink_tree, layout, inputs, root):
      """Adds TestOrchestrationInputs to a SymlinkTree.

      Args:
        symlink_tree (SymlinkTree): The tree of symlinks to write to.
        layout (_Layout): The layout to use for the symlinks.
        inputs (TestOrchestrationInputs): provide paths to add to symlink tree.
        root (Path): Root of symlink_tree
      """
      symlink_tree.register_link(
          linkname=root.join(layout.covargs),
          target=inputs.covargs,
      )
      # Not all builds produce ids. If missing, upload an empty file.
      # Currently only the coverage profile builders consume this file,
      # and those builds always produce ids so this is fine.
      # An alternative to this approach would be to add logic to
      # conditionally isolate / download the files that are only needed for
      # coverage profile builds.
      if not self._api.path.exists(inputs.ids):
        self._api.file.write_raw('write empty build ids', inputs.ids, b'')
      symlink_tree.register_link(
          linkname=root.join(layout.ids),
          target=inputs.ids,
      )
      symlink_tree.register_link(
          linkname=root.join(layout.llvm_cov),
          target=inputs.llvm_cov,
      )
      symlink_tree.register_link(
          linkname=root.join(layout.llvm_profdata),
          target=inputs.llvm_profdata,
      )
      symlink_tree.register_link(
          linkname=root.join(layout.llvm_symbolizer),
          target=inputs.llvm_symbolizer,
      )
      symlink_tree.register_link(
          linkname=root.join(layout.minfs),
          target=inputs.minfs,
      )
      symlink_tree.register_link(
          linkname=root.join(layout.shards),
          target=inputs._shards_file,
      )
      symlink_tree.register_link(
          linkname=root.join(layout.symbolize_tool),
          target=inputs.symbolize_tool,
      )
      symlink_tree.register_link(
          linkname=root.join(layout.task_requests),
          target=inputs._task_requests_file,
      )
      symlink_tree.register_link(
          linkname=root.join(layout.tests),
          target=inputs.tests_file,
      )