# Copyright (c) 2018 Google LLC
#
# 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.
"""A number of common spirv result checks coded in mixin classes.

A test case can use these checks by declaring their enclosing mixin classes
as superclass and providing the expected_* variables required by the check_*()
methods in the mixin classes.
"""
import difflib
import functools
import os
import re
import subprocess
import traceback
from spirv_test_framework import SpirvTest
from builtins import bytes

def convert_to_unix_line_endings(source):
  """Converts all line endings in source to be unix line endings."""
  result = source.replace('\r\n', '\n').replace('\r', '\n')
  return result


def substitute_file_extension(filename, extension):
  """Substitutes file extension, respecting known shader extensions.

    foo.vert -> foo.vert.[extension] [similarly for .frag, .comp, etc.]
    foo.glsl -> foo.[extension]
    foo.unknown -> foo.[extension]
    foo -> foo.[extension]
    """
  if filename[-5:] not in [
      '.vert', '.frag', '.tesc', '.tese', '.geom', '.comp', '.spvasm'
  ]:
    return filename.rsplit('.', 1)[0] + '.' + extension
  else:
    return filename + '.' + extension


def get_object_filename(source_filename):
  """Gets the object filename for the given source file."""
  return substitute_file_extension(source_filename, 'spv')


def get_assembly_filename(source_filename):
  """Gets the assembly filename for the given source file."""
  return substitute_file_extension(source_filename, 'spvasm')


def verify_file_non_empty(filename):
  """Checks that a given file exists and is not empty."""
  if not os.path.isfile(filename):
    return False, 'Cannot find file: ' + filename
  if not os.path.getsize(filename):
    return False, 'Empty file: ' + filename
  return True, ''


class ReturnCodeIsZero(SpirvTest):
  """Mixin class for checking that the return code is zero."""

  def check_return_code_is_zero(self, status):
    if status.returncode:
      return False, 'Non-zero return code: {ret}\n'.format(
          ret=status.returncode)
    return True, ''


class NoOutputOnStdout(SpirvTest):
  """Mixin class for checking that there is no output on stdout."""

  def check_no_output_on_stdout(self, status):
    if status.stdout:
      return False, 'Non empty stdout: {out}\n'.format(out=status.stdout)
    return True, ''


class NoOutputOnStderr(SpirvTest):
  """Mixin class for checking that there is no output on stderr."""

  def check_no_output_on_stderr(self, status):
    if status.stderr:
      return False, 'Non empty stderr: {err}\n'.format(err=status.stderr)
    return True, ''


class SuccessfulReturn(ReturnCodeIsZero, NoOutputOnStdout, NoOutputOnStderr):
  """Mixin class for checking that return code is zero and no output on
    stdout and stderr."""
  pass


class NoGeneratedFiles(SpirvTest):
  """Mixin class for checking that there is no file generated."""

  def check_no_generated_files(self, status):
    all_files = os.listdir(status.directory)
    input_files = status.input_filenames
    if all([f.startswith(status.directory) for f in input_files]):
      all_files = [os.path.join(status.directory, f) for f in all_files]
    generated_files = set(all_files) - set(input_files)
    if len(generated_files) == 0:
      return True, ''
    else:
      return False, 'Extra files generated: {}'.format(generated_files)


class CorrectBinaryLengthAndPreamble(SpirvTest):
  """Provides methods for verifying preamble for a SPIR-V binary."""

  def verify_binary_length_and_header(self, binary, spv_version=0x10000):
    """Checks that the given SPIR-V binary has valid length and header.

        Returns:
            False, error string if anything is invalid
            True, '' otherwise
        Args:
            binary: a bytes object containing the SPIR-V binary
            spv_version: target SPIR-V version number, with same encoding
                 as the version word in a SPIR-V header.
        """

    def read_word(binary, index, little_endian):
      """Reads the index-th word from the given binary file."""
      word = binary[index * 4:(index + 1) * 4]
      if little_endian:
        word = reversed(word)
      return functools.reduce(lambda w, b: (w << 8) | b, word, 0)

    def check_endianness(binary):
      """Checks the endianness of the given SPIR-V binary.

            Returns:
              True if it's little endian, False if it's big endian.
              None if magic number is wrong.
            """
      first_word = read_word(binary, 0, True)
      if first_word == 0x07230203:
        return True
      first_word = read_word(binary, 0, False)
      if first_word == 0x07230203:
        return False
      return None

    num_bytes = len(binary)
    if num_bytes % 4 != 0:
      return False, ('Incorrect SPV binary: size should be a multiple'
                     ' of words')
    if num_bytes < 20:
      return False, 'Incorrect SPV binary: size less than 5 words'

    preamble = binary[0:19]
    little_endian = check_endianness(preamble)
    # SPIR-V module magic number
    if little_endian is None:
      return False, 'Incorrect SPV binary: wrong magic number'

    # SPIR-V version number
    version = read_word(preamble, 1, little_endian)
    # TODO(dneto): Recent Glslang uses version word 0 for opengl_compat
    # profile

    if version != spv_version and version != 0:
      return False, 'Incorrect SPV binary: wrong version number'
    # Shaderc-over-Glslang (0x000d....) or
    # SPIRV-Tools (0x0007....) generator number
    if read_word(preamble, 2, little_endian) != 0x000d0007 and \
            read_word(preamble, 2, little_endian) != 0x00070000:
      return False, ('Incorrect SPV binary: wrong generator magic ' 'number')
    # reserved for instruction schema
    if read_word(preamble, 4, little_endian) != 0:
      return False, 'Incorrect SPV binary: the 5th byte should be 0'

    return True, ''


class CorrectObjectFilePreamble(CorrectBinaryLengthAndPreamble):
  """Provides methods for verifying preamble for a SPV object file."""

  def verify_object_file_preamble(self, filename, spv_version=0x10000):
    """Checks that the given SPIR-V binary file has correct preamble."""

    success, message = verify_file_non_empty(filename)
    if not success:
      return False, message

    with open(filename, 'rb') as object_file:
      object_file.seek(0, os.SEEK_END)
      num_bytes = object_file.tell()

      object_file.seek(0)

      binary = bytes(object_file.read())
      return self.verify_binary_length_and_header(binary, spv_version)

    return True, ''


class CorrectAssemblyFilePreamble(SpirvTest):
  """Provides methods for verifying preamble for a SPV assembly file."""

  def verify_assembly_file_preamble(self, filename):
    success, message = verify_file_non_empty(filename)
    if not success:
      return False, message

    with open(filename) as assembly_file:
      line1 = assembly_file.readline()
      line2 = assembly_file.readline()
      line3 = assembly_file.readline()

    if (line1 != '; SPIR-V\n' or line2 != '; Version: 1.0\n' or
        (not line3.startswith('; Generator: Google Shaderc over Glslang;'))):
      return False, 'Incorrect SPV assembly'

    return True, ''


class ValidObjectFile(SuccessfulReturn, CorrectObjectFilePreamble):
  """Mixin class for checking that every input file generates a valid SPIR-V 1.0
    object file following the object file naming rule, and there is no output on
    stdout/stderr."""

  def check_object_file_preamble(self, status):
    for input_filename in status.input_filenames:
      object_filename = get_object_filename(input_filename)
      success, message = self.verify_object_file_preamble(
          os.path.join(status.directory, object_filename))
      if not success:
        return False, message
    return True, ''


class ValidObjectFile1_3(ReturnCodeIsZero, CorrectObjectFilePreamble):
  """Mixin class for checking that every input file generates a valid SPIR-V 1.3
    object file following the object file naming rule, and there is no output on
    stdout/stderr."""

  def check_object_file_preamble(self, status):
    for input_filename in status.input_filenames:
      object_filename = get_object_filename(input_filename)
      success, message = self.verify_object_file_preamble(
          os.path.join(status.directory, object_filename), 0x10300)
      if not success:
        return False, message
    return True, ''


class ValidObjectFileWithAssemblySubstr(SuccessfulReturn,
                                        CorrectObjectFilePreamble):
  """Mixin class for checking that every input file generates a valid object

    file following the object file naming rule, there is no output on
    stdout/stderr, and the disassmbly contains a specified substring per
    input.
  """

  def check_object_file_disassembly(self, status):
    for an_input in status.inputs:
      object_filename = get_object_filename(an_input.filename)
      obj_file = str(os.path.join(status.directory, object_filename))
      success, message = self.verify_object_file_preamble(obj_file)
      if not success:
        return False, message
      cmd = [status.test_manager.disassembler_path, '--no-color', obj_file]
      process = subprocess.Popen(
          args=cmd,
          stdin=subprocess.PIPE,
          stdout=subprocess.PIPE,
          stderr=subprocess.PIPE,
          cwd=status.directory)
      output = process.communicate(None)
      disassembly = output[0]
      if not isinstance(an_input.assembly_substr, str):
        return False, 'Missing assembly_substr member'
      if an_input.assembly_substr not in disassembly:
        return False, ('Incorrect disassembly output:\n{asm}\n'
                       'Expected substring not found:\n{exp}'.format(
                           asm=disassembly, exp=an_input.assembly_substr))
    return True, ''


class ValidNamedObjectFile(SuccessfulReturn, CorrectObjectFilePreamble):
  """Mixin class for checking that a list of object files with the given
    names are correctly generated, and there is no output on stdout/stderr.

    To mix in this class, subclasses need to provide expected_object_filenames
    as the expected object filenames.
    """

  def check_object_file_preamble(self, status):
    for object_filename in self.expected_object_filenames:
      success, message = self.verify_object_file_preamble(
          os.path.join(status.directory, object_filename))
      if not success:
        return False, message
    return True, ''


class ValidFileContents(SpirvTest):
  """Mixin class to test that a specific file contains specific text
    To mix in this class, subclasses need to provide expected_file_contents as
    the contents of the file and target_filename to determine the location."""

  def check_file(self, status):
    target_filename = os.path.join(status.directory, self.target_filename)
    if not os.path.isfile(target_filename):
      return False, 'Cannot find file: ' + target_filename
    with open(target_filename, 'r') as target_file:
      file_contents = target_file.read()
      if isinstance(self.expected_file_contents, str):
        if file_contents == self.expected_file_contents:
          return True, ''
        return False, ('Incorrect file output: \n{act}\n'
                       'Expected:\n{exp}'
                       'With diff:\n{diff}'.format(
                           act=file_contents,
                           exp=self.expected_file_contents,
                           diff='\n'.join(
                               list(
                                   difflib.unified_diff(
                                       self.expected_file_contents.split('\n'),
                                       file_contents.split('\n'),
                                       fromfile='expected_output',
                                       tofile='actual_output')))))
      elif isinstance(self.expected_file_contents, type(re.compile(''))):
        if self.expected_file_contents.search(file_contents):
          return True, ''
        return False, ('Incorrect file output: \n{act}\n'
                       'Expected matching regex pattern:\n{exp}'.format(
                           act=file_contents,
                           exp=self.expected_file_contents.pattern))
    return False, (
        'Could not open target file ' + target_filename + ' for reading')


class ValidAssemblyFile(SuccessfulReturn, CorrectAssemblyFilePreamble):
  """Mixin class for checking that every input file generates a valid assembly
    file following the assembly file naming rule, and there is no output on
    stdout/stderr."""

  def check_assembly_file_preamble(self, status):
    for input_filename in status.input_filenames:
      assembly_filename = get_assembly_filename(input_filename)
      success, message = self.verify_assembly_file_preamble(
          os.path.join(status.directory, assembly_filename))
      if not success: