# Copyright 2021 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import os
import posixpath
import sys
import json
import itertools
from typing import Any
import unittest

import gpu_path_util
from gpu_tests import common_browser_args as cba
from gpu_tests import common_typing as ct
from gpu_tests import gpu_integration_test
from gpu_tests.util import host_information

html_path = os.path.join(
  gpu_path_util.CHROMIUM_SRC_DIR, 'content', 'test', 'data', 'gpu', 'webcodecs'
)
data_path = os.path.join(
  gpu_path_util.CHROMIUM_SRC_DIR, 'media', 'test', 'data'
)
four_colors_img_path = os.path.join(data_path, 'four-colors.y4m')

frame_sources = [
  'camera',
  'capture',
  'offscreen',
  'arraybuffer',
  'hw_decoder',
  'sw_decoder',
]
hbd_frame_sources = ['hbd_arraybuffer']
video_codecs = [
  'avc1.42001E',
  'hvc1.1.6.L123.00',
  'vp8',
  'vp09.00.10.08',
  'av01.0.04M.08',
]
accelerations = ['prefer-hardware', 'prefer-software']


class WebCodecsIntegrationTest(gpu_integration_test.GpuIntegrationTest):
  @classmethod
  def Name(cls) -> str:
    return 'webcodecs'

  @classmethod
  def _SuiteSupportsParallelTests(cls) -> bool:
    return True

  def _GetSerialGlobs(self) -> set[str]:
    serial_globs = set()
    if host_information.IsWindows() and host_information.IsNvidiaGpu():
      serial_globs |= {
        # crbug.com/1473480. Windows + NVIDIA has a maximum parallel encode
        # limit of 2, so serialize hardware encoding tests on Windows.
        'WebCodecs_*prefer-hardware*',
      }
    return serial_globs

  def _GetSerialTests(self) -> set[str]:
    serial_tests = set()
    if host_information.IsWindows() and host_information.IsArmCpu():
      serial_tests |= {
        # crbug.com/323824490. Seems to flakily lose the D3D11 device when
        # run in parallel.
        'WebCodecs_FrameSizeChange_vp09.00.10.08_hw_decoder',
      }
    return serial_tests

  # pylint: disable=too-many-branches

  @classmethod
  def GenerateGpuTests(cls, options: ct.ParsedCmdArgs) -> ct.TestGenerator:
    tests = itertools.chain(
      cls.GenerateFrameTests(),
      cls.GenerateVideoTests(),
      cls.GenerateAudioTests(),
      cls.BitrateTests(),
      cls.GenerateD3D12EncodingTests(),
    )
    yield from tests

  @classmethod
  def GenerateFrameTests(cls) -> ct.TestGenerator:
    for source_type in frame_sources:
      yield (
        f'WebCodecs_DrawImage_{source_type}',
        'draw-image.html',
        [
          {
            'source_type': source_type,
          }
        ],
      )
      yield (
        f'WebCodecs_TexImage2d_{source_type}',
        'tex-image-2d.html',
        [
          {
            'source_type': source_type,
          }
        ],
      )
      yield (
        f'WebCodecs_copyTo_{source_type}',
        'copyTo.html',
        [
          {
            'source_type': source_type,
          }
        ],
      )
      yield (
        f'WebCodecs_convertToRGB_{source_type}',
        'convert-to-rgb.html',
        [
          {
            'source_type': source_type,
          }
        ],
      )
    for source_type in ['offscreen', 'hw_decoder', 'sw_decoder']:
      for codec in video_codecs:
        for acc in accelerations:
          yield (
            f'WebCodecs_cropEncode_{source_type}_{codec}_{acc}',
            'crop-encode.html',
            [
              {
                'source_type': source_type,
                'codec': codec,
                'acceleration': acc,
              }
            ],
          )
    for source_type in hbd_frame_sources:
      yield (
        f'WebCodecs_DrawImage_{source_type}',
        'draw-image.html',
        [
          {
            'source_type': source_type,
          }
        ],
      )

  @classmethod
  def GenerateAudioTests(cls) -> ct.TestGenerator:
    yield (
      'WebCodecs_AudioEncoding_AAC_LC',
      'audio-encode-decode.html',
      [
        {
          'codec': 'mp4a.67',
          'sample_rate': 48000,
          'channels': 2,
          'aac_format': 'aac',
        }
      ],
    )
    yield (
      'WebCodecs_AudioEncoding_AAC_LC_ADTS',
      'audio-encode-decode.html',
      [
        {
          'codec': 'mp4a.67',
          'sample_rate': 48000,
          'channels': 2,
          'aac_format': 'adts',
        }
      ],
    )

  @classmethod
  def BitrateTests(cls) -> ct.TestGenerator:
    high_res_codecs = [
      'avc1.420034',
      'hvc1.1.6.L123.00',
      'vp8',
      'vp09.00.10.08',
      'av01.0.04M.08',
    ]
    for codec, acc, bitrate_mode, bitrate in itertools.product(
      high_res_codecs,
      accelerations,
      ['constant', 'variable'],
      [1500000, 2000000, 3000000],
    ):
      yield (
        f'WebCodecs_EncodingRateControl_{codec}_{acc}_{bitrate_mode}_{bitrate}',
        'encoding-rate-control.html',
        [
          {
            'codec': codec,
            'acceleration': acc,
            'bitrate_mode': bitrate_mode,
            'bitrate': bitrate,
          }
        ],
      )

  @classmethod
  def GenerateVideoTests(cls) -> ct.TestGenerator:
    yield (
      'WebCodecs_WebRTCPeerConnection_Window',
      'webrtc-peer-connection.html',
      [
        {
          'use_worker': False,
        }
      ],
    )
    yield (
      'WebCodecs_WebRTCPeerConnection_Worker',
      'webrtc-peer-connection.html',
      [
        {
          'use_worker': True,
        }
      ],
    )
    yield (
      'WebCodecs_Terminate_Worker',
      'terminate-worker.html',
      [
        {
          'source_type': 'offscreen',
        }
      ],
    )

    source_type = 'offscreen'
    acc = 'prefer-hardware'
    for codec in ['avc1.42001E', 'hvc1.1.6.L93.B0']:
      yield (
        f'WebCodecs_PerFrameQpEncoding_{source_type}_{codec}_{acc}',
        'frame-qp-encoding.html',
        [
          {
            'source_type': source_type,
            'codec': codec,
            'acceleration': acc,
          }
        ],
      )

    codec = 'av01.0.04M.08'
    acc = 'prefer-software'
    for layers in range(4):
      yield (
        f'WebCodecs_ManualSVC_{codec}_{acc}_layers_{layers}',
        'manual-svc.html',
        [
          {
            'codec': codec,
            'acceleration': acc,
            'layers': layers,
          }
        ],
      )

    for source_type, codec, acc in itertools.product(
      frame_sources, video_codecs, accelerations
    ):
      yield (
        f'WebCodecs_EncodeDecode_{source_type}_{codec}_{acc}',
        'encode-decode.html',
        [
          {
            'source_type': source_type,
            'codec': codec,
            'acceleration': acc,
          }
        ],
      )

    # Also verify we can deal with the encoder's frame delay that we can
    # encounter for the AVC High profile (avc1.64).
    for source_type, codec, acc in itertools.product(
      frame_sources, video_codecs + ['avc1.64001E'], accelerations
    ):
      yield (
        f'WebCodecs_Encode_{source_type}_{codec}_{acc}',
        'encode.html',
        [
          {
            'source_type': source_type,
            'codec': codec,
            'acceleration': acc,
          }
        ],
      )

    resolutions = ['1920x1080', '3840x2160', '7680x3840']
    framerates = [30, 60, 120, 240]
    # Use at least level 6.2 (H.264/H.265/VP9), or level 6.3 (AV1) mimetypes to
    # test 8k 120fps support.
    codecs = [
      'avc1.64003E',
      'hvc1.1.6.L186.B0',
      'vp09.00.62.08',
      'av01.1.19M.08',
    ]
    for resolution, framerate, codec in itertools.product(
      resolutions, framerates, codecs
    ):
      acc = 'prefer-hardware'
      latency_mode = 'quality'
      yield (
        f'WebCodecs_EncodingFramerateResolutions_'
        f'{resolution}_{framerate}_{codec}_{acc}_{latency_mode}',
        'encoding-framerate-resolutions.html',
        [
          {
            'resolution': resolution,
            'framerate': framerate,
            'codec': codec,
            'acceleration': acc,
            'latency_mode': latency_mode,
          }
        ],
      )

    for codec, acc, bitrate_mode, latency_mode in itertools.product(
      video_codecs,
      accelerations,
      ['constant', 'variable'],
      ['realtime', 'quality'],
    ):
      source_type = 'offscreen'
      content_hint = 'motion'
      yield (
        f'WebCodecs_EncodingModes_'
        f'{source_type}_{codec}_{acc}_{bitrate_mode}_{latency_mode}',
        'encoding-modes.html',
        [
          {
            'source_type': source_type,
            'codec': codec,
            'acceleration': acc,
            'bitrate_mode': bitrate_mode,
            'latency_mode': latency_mode,
            'content_hint': content_hint,
          }
        ],
      )

    for codec, content_hint in itertools.product(
      video_codecs, ['detail', 'text', 'motion']
    ):
      source_type = 'offscreen'
      acc = 'prefer-hardware'
      bitrate_mode = 'constant'
      latency_mode = 'realtime'
      yield (
        f'WebCodecs_ContentHint_{codec}_{content_hint}',
        'encoding-modes.html',
        [
          {
            'source_type': source_type,
            'codec': codec,
            'acceleration': acc,
            'bitrate_mode': bitrate_mode,
            'latency_mode': latency_mode,
            'content_hint': content_hint,
          }
        ],
      )

    for codec, acc, layers in itertools.product(
      video_codecs, accelerations, [2, 3]
    ):
      yield (
        f'WebCodecs_SVC_{codec}_{acc}_layers_{layers}',
        'svc.html',
        [
          {
            'codec': codec,
            'acceleration': acc,
            'layers': layers,
          }
        ],
      )

    for codec, acc in itertools.product(video_codecs, accelerations):
      yield (
        f'WebCodecs_EncodeColorSpace_{codec}_{acc}',
        'encode-color-space.html',
        [
          {
            'codec': codec,
            'acceleration': acc,
          }
        ],
      )
      yield (
        f'WebCodecs_MixedSourceEncoding_{codec}_{acc}',
        'mixed-source-encoding.html',
        [
          {
            'codec': codec,
            'acceleration': acc,
          }
        ],
      )

    for codec, source_type in itertools.product(video_codecs, frame_sources):
      yield (
        f'WebCodecs_FrameSizeChange_{codec}_{source_type}',
        'frame-size-change.html',
        [
          {
            'codec': codec,
            'source_type': source_type,
          }
        ],
      )

  @classmethod
  def GenerateD3D12EncodingTests(cls) -> ct.TestGenerator:
    source_type = 'offscreen'
    acc = 'prefer-hardware'

    for source_type, codec in itertools.product(frame_sources, video_codecs):
      yield (
        f'WebCodecs_D3D12_EncodeDecode_{source_type}_{codec}',
        'encode-decode.html',
        [
          {
            'source_type': source_type,
            'codec': codec,
            'acceleration': acc,
            'runs_on_d3d12_encoder': True,
          }
        ],
      )

    for source_type, codec in itertools.product(frame_sources, video_codecs):
      yield (
        f'WebCodecs_D3D12_Encode_{source_type}_{codec}',
        'encode.html',
        [
          {
            'source_type': source_type,
            'codec': codec,
            'acceleration': acc,
            'runs_on_d3d12_encoder': True,
          }
        ],
      )

    for codec, bitrate_mode, latency_mode in itertools.product(
      video_codecs, ['constant', 'variable'], ['realtime', 'quality']
    ):
      source_type = 'offscreen'
      content_hint = 'motion'
      yield (
        f'WebCodecs_D3D12_EncodingModes_'
        f'{source_type}_{codec}_{bitrate_mode}_{latency_mode}',
        'encoding-modes.html',
        [
          {
            'source_type': source_type,
            'codec': codec,
            'acceleration': acc,
            'bitrate_mode': bitrate_mode,
            'latency_mode': latency_mode,
            'content_hint': content_hint,
            'runs_on_d3d12_encoder': True,
          }
        ],
      )

    for codec, content_hint in itertools.product(
      video_codecs, ['detail', 'text', 'motion']
    ):
      source_type = 'offscreen'
      bitrate_mode = 'constant'
      latency_mode = 'realtime'
      yield (
        f'WebCodecs_D3D12_ContentHint_{codec}_{content_hint}',
        'encoding-modes.html',
        [
          {
            'source_type': source_type,
            'codec': codec,
            'acceleration': acc,
            'bitrate_mode': bitrate_mode,
            'latency_mode': latency_mode,
            'content_hint': content_hint,
            'runs_on_d3d12_encoder': True,
          }
        ],
      )

    for codec, layers in itertools.product(video_codecs, [2, 3]):
      yield (
        f'WebCodecs_D3D12_SVC_{codec}_layers_{layers}',
        'svc.html',
        [
          {
            'codec': codec,
            'acceleration': acc,
            'layers': layers,
            'runs_on_d3d12_encoder': True,
          }
        ],
      )

    for codec in video_codecs:
      yield (
        f'WebCodecs_D3D12_EncodeColorSpace_{codec}',
        'encode-color-space.html',
        [
          {
            'codec': codec,
            'acceleration': acc,
            'runs_on_d3d12_encoder': True,
          }
        ],
      )
      yield (
        f'WebCodecs_D3D12_MixedSourceEncoding_{codec}',
        'mixed-source-encoding.html',
        [
          {
            'codec': codec,
            'acceleration': acc,
            'runs_on_d3d12_encoder': True,
          }
        ],
      )

    for codec, source_type in itertools.product(video_codecs, frame_sources):
      yield (
        f'WebCodecs_D3D12_FrameSizeChange_{codec}_{source_type}',
        'frame-size-change.html',
        [
          {
            'codec': codec,
            'source_type': source_type,
            'runs_on_d3d12_encoder': True,
          }
        ],
      )

  # pylint: enable=too-many-branches

  def RunActualGpuTest(self, test_path: str, args: ct.TestArgs) -> None:
    url = self.UrlOfStaticFilePath(posixpath.join(html_path, test_path))
    arg_obj = args[0]
    os_name = self.platform.GetOSName()
    enable_d3d12_encoder = arg_obj.get('runs_on_d3d12_encoder', False)
    if enable_d3d12_encoder and not host_information.IsWindows():
      self.skipTest('Skipping D3D12 based test.')
      return

    args = self.GetBrowserArguments(enable_d3d12_encoder)
    self.RestartBrowserIfNecessaryWithArgs(args)

    arg_obj['validate_camera_frames'] = self.CameraCanShowFourColors(os_name)
    tab = self.tab
    tab.Navigate(url)
    tab.action_runner.WaitForJavaScriptCondition(
      'document.readyState == "complete"'
    )
    tab.EvaluateJavaScript('TEST.run(' + json.dumps(arg_obj) + ')')
    tab.action_runner.WaitForJavaScriptCondition('TEST.finished', timeout=60)
    if tab.EvaluateJavaScript('TEST.skipped'):
      self.skipTest('Skipping test:' + tab.EvaluateJavaScript('TEST.summary()'))
    if not tab.EvaluateJavaScript('TEST.success'):
      self.fail('Test failure:' + tab.EvaluateJavaScript('TEST.summary()'))

  @staticmethod
  def CameraCanShowFourColors(os_name: str) -> bool:
    return os_name not in ('android', 'chromeos', 'fuchsia')

  @classmethod
  def GetBrowserArguments(cls, enable_d3d12_encoder: bool) -> list[str]:
    args = [
      '--use-fake-device-for-media-stream',
      '--use-fake-ui-for-media-stream',
      '--enable-blink-features=SharedArrayBuffer',
      cba.ENABLE_EXPERIMENTAL_WEB_PLATFORM_FEATURES,
    ]

    if enable_d3d12_encoder:
      args.append('--enable-features=D3D12VideoEncodeAccelerator')
      args.append('--disable-features=MediaFoundationVideoEncodeAccelerator')

    if cls.CameraCanShowFourColors(cls.platform.GetOSName()):
      args.append('--use-file-for-fake-video-capture=' + four_colors_img_path)
    return args

  @classmethod
  def SetUpProcess(cls) -> None:
    super(WebCodecsIntegrationTest, cls).SetUpProcess()

    args = []
    # If we don't call CustomizeBrowserArgs cls.platform is None
    cls.CustomizeBrowserArgs(args)
    args = cls.GetBrowserArguments(False)
    cls.CustomizeBrowserArgs(args)

    cls.StartBrowser()
    cls.SetStaticServerDirs([html_path, data_path])

  @classmethod
  def ExpectationsFiles(cls) -> list[str]:
    return [
      os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        'test_expectations',
        'webcodecs_expectations.txt',
      )
    ]


def load_tests(
  loader: unittest.TestLoader, tests: Any, pattern: Any
) -> unittest.TestSuite:
  del loader, tests, pattern  # Unused.
  return gpu_integration_test.LoadAllTestsInModule(sys.modules[__name__])