# 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.
#

"""Implementation of an Artifact{Staging,Retrieval}Service.

The staging service here can be backed by any beam filesystem.
"""

# pytype: skip-file

import concurrent.futures
import hashlib
import os
import queue
import sys
import tempfile
import threading
from io import BytesIO
from typing import Any
from typing import BinaryIO  # pylint: disable=unused-import
from typing import Callable
from typing import MutableMapping
from typing import Optional
from urllib.request import urlopen

import grpc

from apache_beam.io import filesystems
from apache_beam.io.filesystems import CompressionTypes
from apache_beam.portability import common_urns
from apache_beam.portability.api import beam_artifact_api_pb2
from apache_beam.portability.api import beam_artifact_api_pb2_grpc
from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.utils import proto_utils


class ArtifactRetrievalService(
    beam_artifact_api_pb2_grpc.ArtifactRetrievalServiceServicer):

  _DEFAULT_CHUNK_SIZE = 2 << 20

  def __init__(
      self,
      file_reader: Callable[[str], BinaryIO],
      chunk_size=None,
  ):
    self._file_reader = file_reader
    self._chunk_size = chunk_size or self._DEFAULT_CHUNK_SIZE

  def ResolveArtifacts(self, request, context=None):
    return beam_artifact_api_pb2.ResolveArtifactsResponse(
        replacements=request.artifacts)

  def GetArtifact(self, request, context=None):
    if request.artifact.type_urn == common_urns.artifact_types.FILE.urn:
      payload = proto_utils.parse_Bytes(
          request.artifact.type_payload,
          beam_runner_api_pb2.ArtifactFilePayload)
      read_handle = self._file_reader(payload.path)
    elif request.artifact.type_urn == common_urns.artifact_types.URL.urn:
      payload = proto_utils.parse_Bytes(
          request.artifact.type_payload, beam_runner_api_pb2.ArtifactUrlPayload)
      read_handle = urlopen(payload.url)
    elif request.artifact.type_urn == common_urns.artifact_types.EMBEDDED.urn:
      payload = proto_utils.parse_Bytes(
          request.artifact.type_payload,
          beam_runner_api_pb2.EmbeddedFilePayload)
      read_handle = BytesIO(payload.data)
    else:
      raise NotImplementedError(request.artifact.type_urn)

    with read_handle as fin:
      while True:
        chunk = fin.read(self._chunk_size)
        if not chunk:
          break
        yield beam_artifact_api_pb2.GetArtifactResponse(data=chunk)


class ArtifactStagingService(
    beam_artifact_api_pb2_grpc.ArtifactStagingServiceServicer):
  def __init__(
      self,
      file_writer: Callable[[str, Optional[str]], tuple[BinaryIO, str]],
  ):
    self._lock = threading.Lock()
    self._jobs_to_stage: dict[
        str,
        tuple[dict[Any, list[beam_runner_api_pb2.ArtifactInformation]],
              threading.Event]] = {}
    self._file_writer = file_writer

  def register_job(
      self,
      staging_token: str,
      dependency_sets: MutableMapping[
          Any, list[beam_runner_api_pb2.ArtifactInformation]]):
    if staging_token in self._jobs_to_stage:
      raise ValueError('Already staging %s' % staging_token)
    with self._lock:
      self._jobs_to_stage[staging_token] = (
          dict(dependency_sets), threading.Event())

  def resolved_deps(self, staging_token, timeout=None):
    with self._lock:
      dependency_sets, event = self._jobs_to_stage[staging_token]
    try:
      if not event.wait(timeout):
        raise concurrent.futures.TimeoutError()
      return dependency_sets
    finally:
      with self._lock:
        del self._jobs_to_stage[staging_token]

  def ReverseArtifactRetrievalService(self, responses, context=None):
    staging_token = next(responses).staging_token
    with self._lock:
      try:
        dependency_sets, event = self._jobs_to_stage[staging_token]
      except KeyError:
        if context:
          context.set_code(grpc.StatusCode.NOT_FOUND)
          context.set_details('No such staging token: %r' % staging_token)
        raise

    requests = _QueueIter()

    class ForwardingRetrievalService(object):
      def ResolveArtifactss(self, request):
        requests.put(
            beam_artifact_api_pb2.ArtifactRequestWrapper(
                resolve_artifact=request))
        return next(responses).resolve_artifact_response

      def GetArtifact(self, request):
        requests.put(
            beam_artifact_api_pb2.ArtifactRequestWrapper(get_artifact=request))
        while True:
          response = next(responses)
          yield response.get_artifact_response
          if response.is_last:
            break

    def resolve():
      try:
        for key, dependencies in dependency_sets.items():
          dependency_sets[key] = list(
              resolve_as_files(
                  ForwardingRetrievalService(), lambda name: self._file_writer(
                      os.path.join(staging_token, name)),
                  dependencies))
        requests.done()
      except:  # pylint: disable=bare-except
        requests.abort()
        raise
      finally:
        event.set()

    t = threading.Thread(target=resolve)
    t.daemon = True
    t.start()

    return requests


def resolve_as_files(retrieval_service, file_writer, dependencies):
  """Translates a set of dependencies into file-based dependencies."""
  # Resolve until nothing changes.  This ensures that they can be fetched.
  resolution = retrieval_service.ResolveArtifactss(
      beam_artifact_api_pb2.ResolveArtifactsRequest(
          artifacts=dependencies,
          # Anything fetchable will do.
          # TODO(robertwb): Take advantage of shared filesystems, urls.
          preferred_urns=[],
      ))
  dependencies = resolution.replacements

  # Fetch each of the dependencies, using file_writer to store them as
  # file-based artifacts.
  # TODO(robertwb): Consider parallelizing the actual writes.
  for dep in dependencies:
    if dep.role_urn == common_urns.artifact_roles.STAGING_TO.urn:
      base_name = os.path.basename(
          proto_utils.parse_Bytes(
              dep.role_payload,
              beam_runner_api_pb2.ArtifactStagingToRolePayload).staged_name)
    else:
      base_name = None
    unique_name = '-'.join(
        filter(
            None,
            [hashlib.sha256(dep.SerializeToString()).hexdigest(), base_name]))
    file_handle, path = file_writer(unique_name)
    with file_handle as fout:
      for chunk in retrieval_service.GetArtifact(
          beam_artifact_api_pb2.GetArtifactRequest(artifact=dep)):
        fout.write(chunk.data)
    yield beam_runner_api_pb2.ArtifactInformation(
        type_urn=common_urns.artifact_types.FILE.urn,
        type_payload=beam_runner_api_pb2.ArtifactFilePayload(
            path=path).SerializeToString(),
        role_urn=dep.role_urn,
        role_payload=dep.role_payload)


def offer_artifacts(
    artifact_staging_service, artifact_retrieval_service, staging_token):
  """Offers a set of artifacts to an artifact staging service, via the
  ReverseArtifactRetrievalService API.

  The given artifact_retrieval_service should be able to resolve/get all
  artifacts relevant to this job.
  """
  responses = _QueueIter()
  responses.put(
      beam_artifact_api_pb2.ArtifactResponseWrapper(
          staging_token=staging_token))
  requests = artifact_staging_service.ReverseArtifactRetrievalService(responses)
  try:
    for request in requests:
      if request.HasField('resolve_artifact'):
        responses.put(
            beam_artifact_api_pb2.ArtifactResponseWrapper(
                resolve_artifact_response=artifact_retrieval_service.
                ResolveArtifacts(request.resolve_artifact)))
      elif request.HasField('get_artifact'):
        for chunk in artifact_retrieval_service.GetArtifact(
            request.get_artifact):
          responses.put(
              beam_artifact_api_pb2.ArtifactResponseWrapper(
                  get_artifact_response=chunk))
        responses.put(
            beam_artifact_api_pb2.ArtifactResponseWrapper(
                get_artifact_response=beam_artifact_api_pb2.GetArtifactResponse(
                    data=b''),
                is_last=True))
    responses.done()
  except:  # pylint: disable=bare-except
    responses.abort()
    raise


class BeamFilesystemHandler(object):
  def __init__(self, root):
    self._root = root

  def file_reader(self, path):
    return filesystems.FileSystems.open(
        path, compression_type=CompressionTypes.UNCOMPRESSED)

  def file_writer(self, name=None):
    full_path = filesystems.FileSystems.join(self._root, name)
    return filesystems.FileSystems.create(
        full_path, compression_type=CompressionTypes.UNCOMPRESSED), full_path


def resolve_artifacts(artifacts, service, dest_dir):
  if not artifacts:
    return artifacts
  else:
    return [
        maybe_store_artifact(artifact, service,
                             dest_dir) for artifact in service.ResolveArtifacts(
                                 beam_artifact_api_pb2.ResolveArtifactsRequest(
                                     artifacts=artifacts)).replacements
    ]


def maybe_store_artifact(artifact, service, dest_dir):
  if artifact.type_urn in (common_urns.artifact_types.URL.urn,
                           common_urns.artifact_types.EMBEDDED.urn):
    return artifact
  elif artifact.type_urn == common_urns.artifact_types.FILE.urn:
    payload = beam_runner_api_pb2.ArtifactFilePayload.FromString(
        artifact.type_payload)
    # pylint: disable=condition-evals-to-constant
    if os.path.exists(
        payload.path) and payload.sha256 and payload.sha256 == sha256(
            payload.path) and False:
      return artifact
    else:
      return store_artifact(artifact, service, dest_dir)
  else:
    return store_artifact(artifact, service, dest_dir)


def store_artifact(artifact, service, dest_dir):
  hasher = hashlib.sha256()
  with tempfile.NamedTemporaryFile(dir=dest_dir, delete=False) as fout:
    for block in service.GetArtifact(
        beam_artifact_api_pb2.GetArtifactRequest(artifact=artifact)):
      hasher.update(block.data)
      fout.write(block.data)
  return beam_runner_api_pb2.ArtifactInformation(
      type_urn=common_urns.artifact_types.FILE.urn,
      type_payload=beam_runner_api_pb2.ArtifactFilePayload(
          path=fout.name, sha256=hasher.hexdigest()).SerializeToString(),
      role_urn=artifact.role_urn,
      role_payload=artifact.role_payload)


def sha256(path):
  hasher = hashlib.sha256()
  with open(path, 'rb') as fin:
    for block in iter(lambda: fin.read(4 << 20), b''):
      hasher.update(block)
  return hasher.hexdigest()


class _QueueIter(object):

  _END = object()

  def __init__(self):
    self._queue = queue.Queue()

  def put(self, item):
    self._queue.put(item)

  def done(self):
    self._queue.put(self._END)
    self._queue.put(StopIteration)

  def abort(self, exn=None):
    if exn is None:
      exn = sys.exc_info()[1]
    self._queue.put(self._END)
    self._queue.put(exn)

  def __iter__(self):
    return self

  def __next__(self):
    item = self._queue.get()
    if item is self._END:
      raise self._queue.get()
    else:
      return item
