# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#
"""Test cases for :module:`artifact_service_client`."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import hashlib
import os
import random
import shutil
import sys
import tempfile
import time
import unittest

import grpc

from apache_beam.portability.api import beam_artifact_api_pb2
from apache_beam.portability.api import beam_artifact_api_pb2_grpc
from apache_beam.runners.portability import artifact_service
from apache_beam.utils.thread_pool_executor import UnboundedThreadPoolExecutor


class AbstractArtifactServiceTest(unittest.TestCase):

  def setUp(self):
    self._staging_dir = tempfile.mkdtemp()
    self._service = self.create_service(self._staging_dir)

  def tearDown(self):
    if self._staging_dir:
      shutil.rmtree(self._staging_dir)

  def create_service(self, staging_dir):
    raise NotImplementedError(type(self))

  @staticmethod
  def put_metadata(staging_token, name, sha256=None):
    return beam_artifact_api_pb2.PutArtifactRequest(
        metadata=beam_artifact_api_pb2.PutArtifactMetadata(
            staging_session_token=staging_token,
            metadata=beam_artifact_api_pb2.ArtifactMetadata(
                name=name,
                sha256=sha256)))

  @staticmethod
  def put_data(chunk):
    return beam_artifact_api_pb2.PutArtifactRequest(
        data=beam_artifact_api_pb2.ArtifactChunk(
            data=chunk))

  @staticmethod
  def retrieve_artifact(retrieval_service, retrieval_token, name):
    return b''.join(chunk.data for chunk in retrieval_service.GetArtifact(
        beam_artifact_api_pb2.GetArtifactRequest(
            retrieval_token=retrieval_token,
            name=name)))

  def test_basic(self):
    self._run_staging(self._service, self._service)

  def test_with_grpc(self):
    server = grpc.server(UnboundedThreadPoolExecutor())
    try:
      beam_artifact_api_pb2_grpc.add_ArtifactStagingServiceServicer_to_server(
          self._service, server)
      beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server(
          self._service, server)
      port = server.add_insecure_port('[::]:0')
      server.start()
      channel = grpc.insecure_channel('localhost:%d' % port)
      self._run_staging(
          beam_artifact_api_pb2_grpc.ArtifactStagingServiceStub(
              channel),
          beam_artifact_api_pb2_grpc.ArtifactRetrievalServiceStub(
              channel))
      channel.close()
    finally:
      server.stop(1)

  def _run_staging(self, staging_service, retrieval_service):

    staging_session_token = '/session_staging_token \n\0*'

    # First stage some files.
    staging_service.PutArtifact(iter([
        self.put_metadata(staging_session_token, 'name'),
        self.put_data(b'data')]))

    staging_service.PutArtifact(iter([
        self.put_metadata(staging_session_token, 'many_chunks'),
        self.put_data(b'a'),
        self.put_data(b'b'),
        self.put_data(b'c')]))

    staging_service.PutArtifact(iter([
        self.put_metadata(staging_session_token, 'long'),
        self.put_data(b'a' * 1000)]))

    staging_service.PutArtifact(iter([
        self.put_metadata(staging_session_token,
                          'with_hash',
                          hashlib.sha256(b'data...').hexdigest()),
        self.put_data(b'data'),
        self.put_data(b'...')]))

    with self.assertRaises(Exception):
      staging_service.PutArtifact(iter([
          self.put_metadata(staging_session_token,
                            'bad_hash',
                            'bad_hash'),
          self.put_data(b'data')]))

    manifest = beam_artifact_api_pb2.Manifest(artifact=[
        beam_artifact_api_pb2.ArtifactMetadata(name='name'),
        beam_artifact_api_pb2.ArtifactMetadata(name='many_chunks'),
        beam_artifact_api_pb2.ArtifactMetadata(name='long'),
        beam_artifact_api_pb2.ArtifactMetadata(name='with_hash'),
    ])

    retrieval_token = staging_service.CommitManifest(
        beam_artifact_api_pb2.CommitManifestRequest(
            staging_session_token=staging_session_token,
            manifest=manifest)).retrieval_token

    # Now attempt to retrieve them.

    retrieved_manifest = retrieval_service.GetManifest(
        beam_artifact_api_pb2.GetManifestRequest(
            retrieval_token=retrieval_token)).manifest
    self.assertEqual(manifest, retrieved_manifest)

    self.assertEqual(
        b'data',
        self.retrieve_artifact(retrieval_service, retrieval_token, 'name'))

    self.assertEqual(
        b'abc',
        self.retrieve_artifact(
            retrieval_service, retrieval_token, 'many_chunks'))

    self.assertEqual(
        b'a' * 1000,
        self.retrieve_artifact(retrieval_service, retrieval_token, 'long'))

    self.assertEqual(
        b'data...',
        self.retrieve_artifact(retrieval_service, retrieval_token, 'with_hash'))

    with self.assertRaises(Exception):
      self.retrieve_artifact(retrieval_service, retrieval_token, 'bad_hash')

    with self.assertRaises(Exception):
      self.retrieve_artifact(retrieval_service, retrieval_token, 'missing')

  def test_concurrent_requests(self):

    num_sessions = 7
    artifacts = collections.defaultdict(list)

    def name(index):
      # Overlapping names across sessions.
      return 'name%d' % (index // num_sessions)

    def session(index):
      return 'session%d' % (index % num_sessions)

    def delayed_data(data, index, max_msecs=1):
      time.sleep(max_msecs / 1000.0 * random.random())
      return ('%s_%d' % (data, index)).encode('ascii')

    def put(index):
      artifacts[session(index)].append(
          beam_artifact_api_pb2.ArtifactMetadata(name=name(index)))
      self._service.PutArtifact([
          self.put_metadata(session(index), name(index)),
          self.put_data(delayed_data('a', index)),
          self.put_data(delayed_data('b' * 20, index, 2))])
      return session(index)

    def commit(session):
      return session, self._service.CommitManifest(
          beam_artifact_api_pb2.CommitManifestRequest(
              staging_session_token=session,
              manifest=beam_artifact_api_pb2.Manifest(
                  artifact=artifacts[session]))).retrieval_token

    def check(index):
      self.assertEqual(
          delayed_data('a', index) + delayed_data('b' * 20, index, 2),
          self.retrieve_artifact(
              self._service, tokens[session(index)], name(index)))

    # pylint: disable=range-builtin-not-iterating
    pool = UnboundedThreadPoolExecutor()
    sessions = set(pool.map(put, range(100)))
    tokens = dict(pool.map(commit, sessions))
    # List forces materialization.
    _ = list(pool.map(check, range(100)))


@unittest.skipIf(sys.version_info < (3, 6), "Requires Python 3.6+")
class ZipFileArtifactServiceTest(AbstractArtifactServiceTest):
  def create_service(self, staging_dir):
    return artifact_service.ZipFileArtifactService(
        os.path.join(staging_dir, 'test.zip'), chunk_size=10)


class BeamFilesystemArtifactServiceTest(AbstractArtifactServiceTest):
  def create_service(self, staging_dir):
    return artifact_service.BeamFilesystemArtifactService(
        staging_dir, chunk_size=10)


# Don't discover/test the abstract base class.
del AbstractArtifactServiceTest


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