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

from collections.abc import Callable
import dataclasses
from typing import Any
import unittest
from unittest import mock

from telemetry.internal.platform import gpu_info

from gpu_tests import common_typing as ct
from gpu_tests import gpu_helper


# pylint: disable=too-many-arguments
def CreateGpuDeviceDict(
  vendor_id: int | None = None,
  device_id: int | None = None,
  sub_sys_id: int | None = None,
  revision: int | None = None,
  vendor_string: str | None = None,
  device_string: str | None = None,
  driver_vendor: str | None = None,
  driver_version: str | None = None,
) -> dict[str, str | int]:
  return {
    'vendor_id': vendor_id or 0,
    'device_id': device_id or 0,
    'sub_sys_id': sub_sys_id or 0,
    'revision': revision or 0,
    'vendor_string': 'vendor_string'
    if vendor_string is None
    else vendor_string,
    'device_string': 'device_string'
    if device_string is None
    else device_string,
    'driver_vendor': 'driver_vendor'
    if driver_vendor is None
    else driver_vendor,
    'driver_version': 'driver_version'
    if driver_version is None
    else driver_version,
  }


# pylint: enable=too-many-arguments


@dataclasses.dataclass
class TagHelperTestCase:
  """Struct-like class for defining a tag helper test case."""

  expected_result: Any
  device_dict: dict[str, str | int] = ct.EmptyDict()
  aux_attributes: dict[str, Any] = ct.EmptyDict()
  feature_status: dict[str, str] = ct.EmptyDict()
  extra_browser_args: list[str] = ct.EmptyList()


class TagHelpersUnittest(unittest.TestCase):
  def runTagHelperTestWithIndex(
    self,
    tc: TagHelperTestCase,
    test_method: Callable[[gpu_info.GPUInfo | None, int], Any],
  ) -> None:
    """Helper method for running a single tag helper test case w/ index."""
    info = gpu_info.GPUInfo(
      [CreateGpuDeviceDict(**tc.device_dict)],
      tc.aux_attributes,
      tc.feature_status,
      None,
    )
    self.assertEqual(test_method(info, 0), tc.expected_result)

  def runTagHelperTest(
    self,
    tc: TagHelperTestCase,
    test_method: Callable[[gpu_info.GPUInfo | None], Any],
  ) -> None:
    """Helper method for running a single tag helper test case w/o index."""
    info = gpu_info.GPUInfo(
      [CreateGpuDeviceDict(**tc.device_dict)],
      tc.aux_attributes,
      tc.feature_status,
      None,
    )
    self.assertEqual(test_method(info), tc.expected_result)

  def testGetGpuVendorString(self) -> None:
    """Tests all code paths for the GetGpuVendorString() method."""
    cases = [
      # Explicit ID -> AMD.
      TagHelperTestCase(
        'amd',
        {
          'vendor_id': 0x1002,
          'device_string': 'ANGLE (ANGLE gpu, 1, 2)',
          'vendor_string': 'Vendor_gpu 1 2',
        },
      ),
      # Explicit ID -> Intel.
      TagHelperTestCase(
        'intel',
        {
          'vendor_id': 0x8086,
          'device_string': 'ANGLE (ANGLE gpu, 1, 2)',
          'vendor_string': 'Vendor_gpu 1 2',
        },
      ),
      # Explicit ID -> NVIDIA.
      TagHelperTestCase(
        'nvidia',
        {
          'vendor_id': 0x10DE,
          'device_string': 'ANGLE (ANGLE gpu, 1, 2)',
          'vendor_string': 'Vendor_gpu 1 2',
        },
      ),
      # ANGLE vendor string.
      TagHelperTestCase(
        'angle gpu',
        {
          'device_string': 'ANGLE (ANGLE gpu, 1, 2)',
          'vendor_string': 'Vendor_gpu 1 2',
        },
      ),
      # Vendor string.
      TagHelperTestCase('vendor_gpu', {'vendor_string': 'Vendor_gpu 1 2'}),
      # Defined info but unknown.
      TagHelperTestCase('unknown_gpu', {'vendor_string': ''}),
    ]

    for tc in cases:
      self.runTagHelperTestWithIndex(tc, gpu_helper.GetGpuVendorString)

    # Undefined info.
    self.assertEqual(gpu_helper.GetGpuVendorString(None, 0), 'unknown_gpu')

  def testGetGpuDeviceId(self) -> None:
    """Tests all code paths for the GetGpuDeviceId() method."""
    cases = [
      # Explicit device.
      TagHelperTestCase(
        0xFFFF,
        {
          'device_id': 0xFFFF,
          'device_string': 'ANGLE (Vendor, Device, Driver)',
        },
      ),
      # ANGLE device string.
      TagHelperTestCase(
        'Device', {'device_string': 'ANGLE (Vendor, Device, Driver)'}
      ),
      # Device string.
      TagHelperTestCase('Some device', {'device_string': 'Some device'}),
    ]

    for tc in cases:
      self.runTagHelperTestWithIndex(tc, gpu_helper.GetGpuDeviceId)

    # Undefined info.
    self.assertEqual(gpu_helper.GetGpuDeviceId(None, 0), 0)

  def testIntelMasks(self) -> None:
    """Tests the masking methods for determining Intel generation."""
    # Sample of real IDs taken from
    # https://dgpu-docs.intel.com/devices/hardware-table.html
    # Note that 12th gen is referred to as "Xe" or "XeHPG.
    gen_9_ids = {0x1923, 0x3184, 0x3EA4, 0x591C, 0x5A85, 0x9BC8}
    # 0x4F and 0xA7-prefixed samples missing since none were listed.
    gen_12_ids = {0x4C8A, 0x9A40, 0x4905, 0x4680, 0x5690}

    for pci_id in gen_9_ids:
      self.assertTrue(gpu_helper.IsIntelGen9(pci_id))
      self.assertFalse(gpu_helper.IsIntelGen12(pci_id))

    for pci_id in gen_12_ids:
      self.assertTrue(gpu_helper.IsIntelGen12(pci_id))
      self.assertFalse(gpu_helper.IsIntelGen9(pci_id))

  def testGetGpuDriverVendor(self) -> None:
    """Tests all code paths for the GetGpuDriverVendor() method."""
    # Explicit vendor.
    self.runTagHelperTest(
      TagHelperTestCase('vendor', {'driver_vendor': 'vendor'}),
      gpu_helper.GetGpuDriverVendor,
    )
    # Undefined info.
    self.assertEqual(gpu_helper.GetGpuDriverVendor(None), None)

  def testGetGpuDriverVersion(self) -> None:
    """Tests all code paths for the GetGpuDriverVersion() method."""
    # Explicit version.
    self.runTagHelperTest(
      TagHelperTestCase('Some version', {'driver_version': 'Some version'}),
      gpu_helper.GetGpuDriverVersion,
    )
    # Undefined info.
    self.assertEqual(gpu_helper.GetGpuDriverVersion(None), None)

  def testGetANGLERenderer(self) -> None:
    """Tests all code paths for the GetANGLERenderer() method."""
    cases = [
      # No aux attributes.
      TagHelperTestCase('angle-disabled'),
      # Non-ANGLE renderer.
      TagHelperTestCase(
        'angle-disabled', aux_attributes={'gl_renderer': 'renderer'}
      ),
      # D3D11.
      TagHelperTestCase(
        'angle-d3d11', aux_attributes={'gl_renderer': 'ANGLE Direct3D11'}
      ),
      # D3D9.
      TagHelperTestCase(
        'angle-d3d9', aux_attributes={'gl_renderer': 'ANGLE Direct3D9'}
      ),
      # OpenGL ES.
      TagHelperTestCase(
        'angle-opengles', aux_attributes={'gl_renderer': 'ANGLE OpenGL ES'}
      ),
      # OpenGL.
      TagHelperTestCase(
        'angle-opengl', aux_attributes={'gl_renderer': 'ANGLE OpenGL'}
      ),
      # Metal.
      TagHelperTestCase(
        'angle-metal', aux_attributes={'gl_renderer': 'ANGLE Metal'}
      ),
      # SwiftShader, explicitly test that it's chosen over Vulkan.
      TagHelperTestCase(
        'angle-swiftshader',
        aux_attributes={'gl_renderer': 'ANGLE Vulkan SwiftShader'},
      ),
      # Vulkan.
      TagHelperTestCase(
        'angle-vulkan', aux_attributes={'gl_renderer': 'ANGLE Vulkan'}
      ),
    ]

    for tc in cases:
      self.runTagHelperTest(tc, gpu_helper.GetANGLERenderer)

    # Undefined info.
    self.assertEqual(gpu_helper.GetANGLERenderer(None), 'angle-disabled')

  def testGetCommandDecoder(self) -> None:
    """Tests all code paths for the GetcommandDecoder() method."""
    cases = [
      # No aux attributes.
      TagHelperTestCase('no_passthrough'),
      # Validating.
      TagHelperTestCase(
        'no_passthrough', aux_attributes={'passthrough_cmd_decoder': False}
      ),
      # Passthrough.
      TagHelperTestCase(
        'passthrough', aux_attributes={'passthrough_cmd_decoder': True}
      ),
    ]

    for tc in cases:
      self.runTagHelperTest(tc, gpu_helper.GetCommandDecoder)

    # Undefined info.
    self.assertEqual(gpu_helper.GetCommandDecoder(None), 'no_passthrough')

  def testGetSkiaGraphiteStatus(self) -> None:
    """Tests all the code paths for the GetSkiaGraphiteStatus() method."""
    cases = [
      # No feature status.
      TagHelperTestCase('graphite-disabled'),
      # Feature status off.
      TagHelperTestCase(
        'graphite-disabled', feature_status={'skia_graphite': 'disabled'}
      ),
      # Feature status on.
      TagHelperTestCase(
        'graphite-enabled', feature_status={'skia_graphite': 'enabled_on'}
      ),
    ]

    for tc in cases:
      self.runTagHelperTest(tc, gpu_helper.GetSkiaGraphiteStatus)

    # Undefined info.
    self.assertEqual(
      gpu_helper.GetSkiaGraphiteStatus(None), 'graphite-disabled'
    )

  def testGetSkiaRenderer(self) -> None:
    """Tests all code paths for the GetSkiaRenderer() method."""
    cases = [
      # No feature status.
      TagHelperTestCase('renderer-software'),
      # No GPU Compositing.
      TagHelperTestCase(
        'renderer-software', feature_status={'gpu_compositing': 'disabled'}
      ),
      # No renderer.
      TagHelperTestCase(
        'renderer-software', feature_status={'gpu_compositing': 'enabled'}
      ),
      # Vulkan Skia Renderer.
      TagHelperTestCase(
        'renderer-skia-vulkan',
        feature_status={
          'gpu_compositing': 'enabled',
          'vulkan': 'enabled_on',
          'opengl': 'enabled_on',
        },
      ),
      # GL Skia Renderer.
      TagHelperTestCase(
        'renderer-skia-gl',
        feature_status={
          'gpu_compositing': 'enabled',
          'vulkan': 'enabled_off',
          'opengl': 'enabled_on',
        },
      ),
    ]

    for tc in cases:
      self.runTagHelperTest(tc, gpu_helper.GetSkiaRenderer)

    # Undefined info.
    self.assertEqual(gpu_helper.GetSkiaRenderer(None), 'renderer-software')

  def testGetDisplayServer(self) -> None:
    """Tests all code paths for the GetDisplayServer() method."""
    with mock.patch(
      'gpu_tests.util.host_information.IsLinux', return_value=True
    ):
      # Remote platforms.
      for browser_type in gpu_helper.REMOTE_BROWSER_TYPES:
        self.assertEqual(gpu_helper.GetDisplayServer(browser_type), None)
      # X.
      with mock.patch.dict('os.environ', {}, clear=True):
        self.assertEqual(gpu_helper.GetDisplayServer(''), 'display-server-x')
      # Wayland.
      with mock.patch.dict('os.environ', {'WAYLAND_DISPLAY': '1'}, clear=True):
        self.assertEqual(
          gpu_helper.GetDisplayServer(''), 'display-server-wayland'
        )

    with mock.patch(
      'gpu_tests.util.host_information.IsLinux', return_value=False
    ):
      self.assertEqual(gpu_helper.GetDisplayServer(''), None)

  def testGetAsanStatus(self) -> None:
    """Tests all code paths for the GetAsanStatus() method."""
    cases = [
      # No aux attributes.
      TagHelperTestCase('no-asan'),
      # Built without ASan.
      TagHelperTestCase('no-asan', aux_attributes={'is_asan': False}),
      # Built with ASan.
      TagHelperTestCase('asan', aux_attributes={'is_asan': True}),
    ]

    for tc in cases:
      self.runTagHelperTest(tc, gpu_helper.GetAsanStatus)

    # Undefined info.
    self.assertEqual(gpu_helper.GetAsanStatus(None), 'no-asan')

  def testGetTargetCpuStatus(self) -> None:
    """Tests all code paths for the GetTargetCpuStatus() method."""
    cases = [
      # No aux attributes.
      TagHelperTestCase('target-cpu-unknown'),
      # Target CPU specified.
      TagHelperTestCase(
        'target-cpu-32', aux_attributes={'target_cpu_bits': 32}
      ),
    ]

    for tc in cases:
      self.runTagHelperTest(tc, gpu_helper.GetTargetCpuStatus)

    # Undefined info.
    self.assertEqual(gpu_helper.GetTargetCpuStatus(None), 'target-cpu-unknown')

  def testGetClangCoverage(self) -> None:
    """Tests all code paths for the GetClangCoverage() method."""
    cases = [
      # No aux attributes.
      TagHelperTestCase('no-clang-coverage'),
      # Built without Clang coverage.
      TagHelperTestCase(
        'no-clang-coverage', aux_attributes={'is_clang_coverage': False}
      ),
      # Built with Clang coverage.
      TagHelperTestCase(
        'clang-coverage', aux_attributes={'is_clang_coverage': True}
      ),
    ]

    for tc in cases:
      self.runTagHelperTest(tc, gpu_helper.GetClangCoverage)

    # Undefined info.
    self.assertEqual(gpu_helper.GetClangCoverage(None), 'no-clang-coverage')


class ReplaceTagsUnittest(unittest.TestCase):
  def testSubstringReplacement(self) -> None:
    tags = ['some_tag', 'some-nvidia-corporation', 'another_tag']
    self.assertEqual(
      gpu_helper.ReplaceTags(tags), ['some_tag', 'some-nvidia', 'another_tag']
    )

  def testRegexReplacement(self) -> None:
    tags = [
      'some_tag',
      'google-Vulkan-1.3.0-(SwiftShader-Device-(LLVM-10.0.0)-(0x0000C0DE))',
      'another_tag',
    ]
    self.assertEqual(
      gpu_helper.ReplaceTags(tags), ['some_tag', 'google-vulkan', 'another_tag']
    )


class EvaluateVersionComparisonUnittest(unittest.TestCase):
  def testWindowsIntelUncomparableVersions(self) -> None:
    """Tests Windows Intel comparison when versions cannot be compared."""
    non_ne_operations = ('eq', 'ge', 'gt', 'le', 'lt')
    # Versions should only be comparable with 4 elements.
    self.assertTrue(
      gpu_helper.EvaluateVersionComparison(
        '1.2.3', 'ne', '1.2.3', 'win', 'intel'
      )
    )
    for op in non_ne_operations:
      self.assertFalse(
        gpu_helper.EvaluateVersionComparison(
          '1.2.3', op, '1.2.3', 'win', 'intel'
        )
      )
    self.assertTrue(
      gpu_helper.EvaluateVersionComparison(
        '1.2.3.4.5', 'ne', '1.2.3.4.5', 'win', 'intel'
      )
    )
    for op in non_ne_operations:
      self.assertFalse(
        gpu_helper.EvaluateVersionComparison(
          '1.2.3.4.5', op, '1.2.3.4.5', 'win', 'intel'
        )
      )

  def testWindowsIntelOnlyLastTwoPartsUsed(self) -> None:
    """Tests that only the last two version parts are used on Windows Intel."""
    self.assertTrue(
      gpu_helper.EvaluateVersionComparison(
        '1.2.3.4', 'eq', '2.3.3.4', 'win', 'intel'
      )
    )
    self.assertFalse(
      gpu_helper.EvaluateVersionComparison(
        '1.2.3.4', 'eq', '2.3.5.4', 'win', 'intel'
      )
    )

  def testInvalidOperation(self) -> None:
    """Tests that an error is raised when using an invalid operation."""
    with self.assertRaisesRegex(Exception, 'Invalid operation: foo'):
      gpu_helper.EvaluateVersionComparison('1.2.3.4', 'foo', '2.3.4.5')

  def testEqual(self) -> None:
    """Tests that equality operations work as expected."""
    eq_operations = ('eq', 'ge', 'le')
    for op in eq_operations:
      # Purely numerical.
      self.assertTrue(
        gpu_helper.EvaluateVersionComparison('1.2.3.4', op, '1.2.3.4')
      )

      # Numerical + suffix.
      self.assertTrue(
        gpu_helper.EvaluateVersionComparison('1a.2b.3c.4a', op, '1a.2b.3c.4a')
      )

      # Mismatched length implies 0.
      self.assertTrue(
        gpu_helper.EvaluateVersionComparison('1.2.0.0', op, '1.2')
      )
      self.assertTrue(
        gpu_helper.EvaluateVersionComparison('1.2', op, '1.2.0.0')
      )

      # Failure to parse gets skipped unless in the reference version.
      self.assertTrue(
        gpu_helper.EvaluateVersionComparison('1.2..4', op, '1.2.3.4')
      )
      with self.assertRaises(AssertionError):
        gpu_helper.EvaluateVersionComparison('1.2.3.4', op, '1.2..4')

  def testNotEqual(self) -> None:
    """Tests that the not equal operation works as expected."""
    true_cases = (
      # Purely numerical.
      ('1.2.3.4', '1.2.3'),
      ('1.2.3', '1.2.3.4'),
      # Same suffix, different numerical.
      ('2a.2b.3c.4d', '1a.2b.3c.4d'),
      ('1a.3b.3c.4d', '1a.2b.3c.4d'),
      ('1a.2b.4c.4d', '1a.2b.3c.4d'),
      ('1a.2b.3c.5d', '1a.2b.3c.4d'),
      # Same numerical, different suffix.
      ('1b.2b.3c.4d', '1a.2b.3c.4d'),
      ('1a.2c.3c.4d', '1a.2b.3c.4d'),
      ('1a.2b.3d.4d', '1a.2b.3c.4d'),
      ('1a.2b.3c.4e', '1a.2b.3c.4d'),
    )
    false_cases = (
      # Numerical.
      ('1.2.3.4', '1.2.3.4'),
      # Numerical + suffix.
      ('1a.2b.3c.4d', '1a.2b.3c.4d'),
    )
    for left, right in true_cases:
      self.assertTrue(gpu_helper.EvaluateVersionComparison(left, 'ne', right))
      self.assertTrue(gpu_helper.EvaluateVersionComparison(right, 'ne', left))
    for left, right in false_cases:
      self.assertFalse(gpu_helper.EvaluateVersionComparison(left, 'ne', right))
      self.assertFalse(gpu_helper.EvaluateVersionComparison(right, 'ne', left))

  def testGreater(self) -> None:
    """Tests that greater operations work as expected."""
    gt_operations = ('gt', 'ge')
    for op in gt_operations:
      for left, right in GetGreaterTestCases():
        self.assertTrue(gpu_helper.EvaluateVersionComparison(left, op, right))
        self.assertFalse(gpu_helper.EvaluateVersionComparison(right, op, left))

  def testLess(self) -> None:
    """Tests that less operations work as expected."""
    lt_operations = ('lt', 'le')
    for op in lt_operations:
      # Less than test cases are simply the inverse of greater than test cases,
      # so just reverse the order.
      for right, left in GetGreaterTestCases():
        self.assertTrue(gpu_helper.EvaluateVersionComparison(left, op, right))
        self.assertFalse(gpu_helper.EvaluateVersionComparison(right, op, left))


def GetGreaterTestCases() -> tuple[tuple[str, str], ...]:
  return (
    # Purely numerical.
    ('2.2.3.4', '1.2.3.4'),
    ('1.3.3.4', '1.2.3.4'),
    ('1.2.4.4', '1.2.3.4'),
    ('1.2.3.5', '1.2.3.4'),
    # Same suffix, different numerical.
    ('2a.2b.3c.4d', '1a.2b.3c.4d'),
    ('1a.3b.3c.4d', '1a.2b.3c.4d'),
    ('1a.2b.4c.4d', '1a.2b.3c.4d'),
    ('1a.2b.3c.5d', '1a.2b.3c.4d'),
    # Same numerical, different suffix.
    ('1b.2b.3c.4d', '1a.2b.3c.4d'),
    ('1a.2c.3c.4d', '1a.2b.3c.4d'),
    ('1a.2b.3d.4d', '1a.2b.3c.4d'),
    ('1a.2b.3c.4e', '1a.2b.3c.4d'),
  )


if __name__ == '__main__':
  unittest.main(verbosity=2)