# pylint: disable=g-direct-third-party-import
# Copyright 2018 The Bazel Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import io
import os
import tempfile
import unittest
from unittest import mock

from google.protobuf import proto

from scripts.aquery_differ import aquery_differ
from src.main.protobuf import analysis_v2_pb2


def make_aquery_output(action_objs, artifact_objs, path_fragment_objs):
  action_graph = analysis_v2_pb2.ActionGraphContainer()

  for path_fragment_obj in path_fragment_objs:
    path_fragment = action_graph.path_fragments.add()
    path_fragment.id = path_fragment_obj["id"]
    path_fragment.label = path_fragment_obj["label"]
    if "parent_id" in path_fragment_obj:
      path_fragment.parent_id = path_fragment_obj["parent_id"]

  for artifact_obj in artifact_objs:
    artifact = action_graph.artifacts.add()
    artifact.id = artifact_obj["id"]
    artifact.path_fragment_id = artifact_obj["path_fragment_id"]

  for action_obj in action_objs:
    action = action_graph.actions.add()
    action.output_ids.extend(action_obj["output_ids"])
    action.arguments.extend(action_obj["arguments"])

    if "input_dep_set_ids" in action_obj:
      action.input_dep_set_ids.extend(action_obj["input_dep_set_ids"])

  return action_graph


def make_aquery_output_with_dep_set(
    action_objs, artifact_objs, path_fragment_objs, dep_set_objs
):
  action_graph = make_aquery_output(
      action_objs, artifact_objs, path_fragment_objs
  )

  for ds in dep_set_objs:
    dep_set = action_graph.dep_set_of_files.add()
    dep_set.id = ds["id"]
    dep_set.direct_artifact_ids.extend(ds["direct_artifact_ids"])
    dep_set.transitive_dep_set_ids.extend(ds["transitive_dep_set_ids"])

  return action_graph


class CmdLineDifferTest(unittest.TestCase):

  def test_no_difference(self):
    action_graph = make_aquery_output(
        action_objs=[
            {"arguments": ["-a", "-b"], "output_ids": [1, 2]},
            {"arguments": ["-c"], "output_ids": [3]},
        ],
        artifact_objs=[
            {"id": 1, "path_fragment_id": 2},
            {"id": 2, "path_fragment_id": 3},
            {"id": 3, "path_fragment_id": 4},
        ],
        path_fragment_objs=[
            {"id": 1, "label": "root"},
            {"id": 2, "label": "foo", "parent_id": 1},
            {"id": 3, "label": "bar", "parent_id": 1},
            {"id": 4, "label": "baz", "parent_id": 1},
        ],
    )
    mock_stdout = io.StringIO()
    attrs = ["cmdline"]
    with mock.patch("sys.stdout", mock_stdout):
      aquery_differ._aquery_diff(
          action_graph, action_graph, attrs, "before", "after"
      )
      self.assertEqual(mock_stdout.getvalue(), "No difference\n")

  def test_no_difference_different_output_files_order(self):
    first = make_aquery_output(
        action_objs=[
            {"arguments": ["-a", "-b"], "output_ids": [1, 2]},
        ],
        artifact_objs=[
            {"id": 1, "path_fragment_id": 2},
            {"id": 2, "path_fragment_id": 3},
        ],
        path_fragment_objs=[
            {"id": 1, "label": "root"},
            {"id": 2, "label": "foo", "parent_id": 1},
            {"id": 3, "label": "bar", "parent_id": 1},
        ],
    )

    second = make_aquery_output(
        action_objs=[
            {"arguments": ["-a", "-b"], "output_ids": [2, 1]},
        ],
        artifact_objs=[
            {"id": 1, "path_fragment_id": 2},
            {"id": 2, "path_fragment_id": 3},
        ],
        path_fragment_objs=[
            {"id": 1, "label": "root"},
            {"id": 2, "label": "foo", "parent_id": 1},
            {"id": 3, "label": "bar", "parent_id": 1},
        ],
    )

    mock_stdout = io.StringIO()
    attrs = ["cmdline"]
    with mock.patch("sys.stdout", mock_stdout):
      aquery_differ._aquery_diff(first, second, attrs, "before", "after")
      self.assertEqual(mock_stdout.getvalue(), "No difference\n")

  def test_first_has_extra_output_files(self):
    first = make_aquery_output(
        action_objs=[
            {"arguments": ["-a", "-b"], "output_ids": [1, 2]},
            {"arguments": ["-c"], "output_ids": [3]},
        ],
        artifact_objs=[
            {"id": 1, "path_fragment_id": 2},
            {"id": 2, "path_fragment_id": 3},
            {"id": 3, "path_fragment_id": 4},
        ],
        path_fragment_objs=[
            {"id": 1, "label": "root"},
            {"id": 2, "label": "foo", "parent_id": 1},
            {"id": 3, "label": "bar", "parent_id": 1},
            {"id": 4, "label": "baz", "parent_id": 1},
        ],
    )
    second = make_aquery_output(
        action_objs=[
            {"arguments": ["-a", "-b"], "output_ids": [1, 2]},
        ],
        artifact_objs=[
            {"id": 1, "path_fragment_id": 2},
            {"id": 2, "path_fragment_id": 3},
        ],
        path_fragment_objs=[
            {"id": 1, "label": "root"},
            {"id": 2, "label": "foo", "parent_id": 1},
            {"id": 3, "label": "bar", "parent_id": 1},
        ],
    )

    baz_path = os.path.join("root", "baz")
    expected_error = (
        "Aquery output 'before' change contains an action "
        "that generates the following outputs that aquery "
        "output 'after' change doesn't:\n{}\n\n".format(baz_path)
    )
    mock_stdout = io.StringIO()
    attrs = ["cmdline"]
    with mock.patch("sys.stdout", mock_stdout):
      aquery_differ._aquery_diff(first, second, attrs, "before", "after")
      self.assertEqual(mock_stdout.getvalue(), expected_error)

  def test_different_command_lines(self):
    first = make_aquery_output(
        action_objs=[
            {"arguments": ["-a", "-d"], "output_ids": [1, 2]},
            {"arguments": ["-c"], "output_ids": [3]},
        ],
        artifact_objs=[
            {"id": 1, "path_fragment_id": 2},
            {"id": 2, "path_fragment_id": 3},
            {"id": 3, "path_fragment_id": 4},
        ],
        path_fragment_objs=[
            {"id": 1, "label": "root"},
            {"id": 2, "label": "foo", "parent_id": 1},
            {"id": 3, "label": "bar", "parent_id": 1},
            {"id": 4, "label": "baz", "parent_id": 1},
        ],
    )
    second = make_aquery_output(
        action_objs=[
            {"arguments": ["-a", "-b"], "output_ids": [1, 2]},
            {"arguments": ["-c", "-d"], "output_ids": [3]},
        ],
        artifact_objs=[
            {"id": 1, "path_fragment_id": 2},
            {"id": 2, "path_fragment_id": 3},
            {"id": 3, "path_fragment_id": 4},
        ],
        path_fragment_objs=[
            {"id": 1, "label": "root"},
            {"id": 2, "label": "foo", "parent_id": 1},
            {"id": 3, "label": "bar", "parent_id": 1},
            {"id": 4, "label": "baz", "parent_id": 1},
        ],
    )

    foo_path = os.path.join("root", "foo")
    bar_path = os.path.join("root", "bar")
    baz_path = os.path.join("root", "baz")

    expected_error_one = "\n".join([
        "Difference in the action that generates the following output(s):",
        "\t{}".format(baz_path),
        "--- before",
        "+++ after",
        "@@ -1 +1,2 @@",
        " -c",
        "+-d",
        "\n",
    ])
    expected_error_two = "\n".join([
        "Difference in the action that generates the following output(s):",
        "\t{}".format(bar_path),
        "\t{}".format(foo_path),
        "--- before",
        "+++ after",
        "@@ -1,2 +1,2 @@",
        " -a",
        "--d",
        "+-b",
        "\n",
    ])
    attrs = ["cmdline"]

    mock_stdout = io.StringIO()
    with mock.patch("sys.stdout", mock_stdout):
      aquery_differ._aquery_diff(first, second, attrs, "before", "after")
      self.assertIn(expected_error_one, mock_stdout.getvalue())
      self.assertIn(expected_error_two, mock_stdout.getvalue())

  def test_different_inputs(self):
    first = make_aquery_output_with_dep_set(
        action_objs=[
            {"arguments": [], "output_ids": [1, 2], "input_dep_set_ids": [2]}
        ],
        artifact_objs=[
            {"id": 1, "path_fragment_id": 2},
            {"id": 2, "path_fragment_id": 3},
        ],
        path_fragment_objs=[
            {"id": 1, "label": "root"},
            {"id": 2, "label": "foo", "parent_id": 1},
            {"id": 3, "label": "bar", "parent_id": 1},
        ],
        dep_set_objs=[
            {"id": 1, "transitive_dep_set_ids": [], "direct_artifact_ids": [1]},
            {
                "id": 2,
                # Note: Artifact 1 which is both a direct and transitive input
                # should be deduplicated in the aquery_differ output.
                "transitive_dep_set_ids": [1],
                "direct_artifact_ids": [1, 2],
            },
        ],
    )
    second = make_aquery_output_with_dep_set(
        action_objs=[
            {"arguments": [], "output_ids": [1, 2], "input_dep_set_ids": [1]},
        ],
        artifact_objs=[
            {"id": 1, "path_fragment_id": 2},
            {"id": 2, "path_fragment_id": 3},
        ],
        path_fragment_objs=[
            {"id": 1, "label": "root"},
            {"id": 2, "label": "foo", "parent_id": 1},
            {"id": 3, "label": "bar", "parent_id": 1},
        ],
        dep_set_objs=[
            {"id": 1, "transitive_dep_set_ids": [], "direct_artifact_ids": [1]}
        ],
    )

    foo_path = os.path.join("root", "foo")
    bar_path = os.path.join("root", "bar")
    expected_error_one = "\n".join([
        "Difference in the action that generates the following output(s):",
        "\t{}".format(bar_path),
        "\t{}".format(foo_path),
        "--- before",
        "+++ after",
        "@@ -1,2 +1 @@",
        "-{}".format(bar_path),
        " {}".format(foo_path),
        "\n",
    ])
    attrs = ["inputs"]

    mock_stdout = io.StringIO()
    with mock.patch("sys.stdout", mock_stdout):
      aquery_differ._aquery_diff(first, second, attrs, "before", "after")
      self.assertIn(expected_error_one, mock_stdout.getvalue())


class ReadProtoTest(unittest.TestCase):

  def test_read_streamed_proto(self):
    ag1 = analysis_v2_pb2.ActionGraphContainer()
    a1 = ag1.artifacts.add()
    a1.id = 1
    a1.path_fragment_id = 1

    ag2 = analysis_v2_pb2.ActionGraphContainer()
    a2 = ag2.artifacts.add()
    a2.id = 2
    a2.path_fragment_id = 2

    def write_delimited(f, msg):
      proto.serialize_length_prefixed(msg, f)

    with tempfile.NamedTemporaryFile(delete=False) as f:
      write_delimited(f, ag1)
      write_delimited(f, ag2)
      temp_file_path = f.name

    try:
      result_proto = aquery_differ._read_proto(temp_file_path, "streamed_proto")
      self.assertEqual(len(result_proto.artifacts), 2)
      self.assertEqual(result_proto.artifacts[0].id, 1)
      self.assertEqual(result_proto.artifacts[1].id, 2)
    finally:
      os.remove(temp_file_path)


if __name__ == "__main__":
  unittest.main()