# -*- coding: utf-8 -*- #
# Copyright 2023 Google LLC. All Rights Reserved.
#
# 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.
"""Task for listing, sorting, and writing files for rsync."""

from __future__ import annotations

import errno
import heapq
import itertools
import os
import threading

from googlecloudsdk.api_lib.storage import cloud_api
from googlecloudsdk.command_lib.storage import errors
from googlecloudsdk.command_lib.storage import folder_util
from googlecloudsdk.command_lib.storage import regex_util
from googlecloudsdk.command_lib.storage import rsync_command_util
from googlecloudsdk.command_lib.storage import storage_url
from googlecloudsdk.command_lib.storage import wildcard_iterator
from googlecloudsdk.command_lib.storage.resources import resource_reference
from googlecloudsdk.command_lib.storage.tasks import task
from googlecloudsdk.core import log
from googlecloudsdk.core import properties
from googlecloudsdk.core.util import files


def sorting_key_for_csv_line(csv_line):
  """Returns the sorting key for the chunk CSV line.

  This key is based on only the first field of the CSV line, which is the URL of
  the resource. Since we use comma as a delimiter, we can't use the
  entire CSV line as the key since there are unicode characters before the comma
  like #, $, " which if present in the csv line can cause sorting issues.

  Args:
    csv_line (str): The CSV line to get the sorting key for.

  Returns:
    str: The sorting key for the CSV line.
  """
  fields = rsync_command_util.get_fields_from_csv_line(csv_line)
  return fields[0]


def get_resource_chunk(
    file_iterator: iter, chunk_size: int, gstmp_path: str | None = None
) -> list[resource_reference.Resource]:
  """Returns a chunk of resources from the file iterator.

  This method iterates through the file iterator and returns a chunk of
  resources. If gstmp_path is not None and
  use_rsync_unmatched_gstmp_handling is enabled, the gstmp file entry will be
  written to the gstmp_path, and the resource will not be included in the
  returned list.

  Args:
    file_iterator: An iterator yielding resource objects.
    chunk_size (int): The maximum number of items to retrieve in the chunk.
    gstmp_path (str|None): If not None, store the temporary list file in this
      path.

  Returns:
    list: A list of resources from the iterator, up to chunk_size.
  """
  if gstmp_path is None or not (
      properties.VALUES.storage.use_rsync_unmatched_gstmp_handling.GetBool()
  ):
    return list(itertools.islice(file_iterator, chunk_size))

  # Use gstmp_path to store temporary list file if
  # use_rsync_unmatched_gstmp_handling is enabled.
  resource_chunk = []
  resource_chunk_size = 0
  with files.FileWriter(
      gstmp_path, create_path=True, append=True
  ) as file_writer:
    while resource_chunk_size < chunk_size:
      try:
        file = next(file_iterator)
        if file.storage_url.url_string.endswith(
            storage_url.TEMPORARY_FILE_SUFFIX
        ):
          file_writer.write(
              rsync_command_util.get_csv_line_from_resource(file) + '\n'
          )
          continue
        resource_chunk.append(file)
        resource_chunk_size += 1
      except StopIteration:
        break
  return resource_chunk


class GetSortedContainerContentsTask(task.Task):
  """Updates a local file's POSIX metadata."""

  def __init__(
      self,
      container,
      output_path,
      exclude_pattern_strings=None,
      managed_folders_only=False,
      ignore_symlinks=True,
      recurse=False,
      gstmp_path=None,
  ):
    """Initializes task.

    Args:
      container (Resource): Contains path of files to fetch.
      output_path (str): Where to write final sorted file list.
      exclude_pattern_strings (List[str]|None): Ignore resources whose paths
        matched these regex patterns.
      managed_folders_only (bool): If True, populates the file with managed
        folders. Otherwise, populates the file with object resources.
      ignore_symlinks (bool): Should FileWildcardIterator skip symlinks.
      recurse (bool): Gather nested items in container.
      gstmp_path (str|None): If not None, store the gstmp file list in this
        path.
    """
    super(GetSortedContainerContentsTask, self).__init__()
    self._container_query_path = container.storage_url.join(
        '**' if recurse else '*'
    ).url_string
    self._output_path = output_path
    self._gstmp_path = gstmp_path
    if self._gstmp_path:
      # We will overwrite the gstmp file list if it exists.
      rsync_command_util.try_to_delete_file(self._gstmp_path)

    if exclude_pattern_strings:
      container_url_trailing_delimiter = container.storage_url.join('')
      if isinstance(container_url_trailing_delimiter, storage_url.FileUrl):
        # Remove 'file://' prefix.
        container_prefix = container_url_trailing_delimiter.resource_name
      else:
        container_prefix = (
            container_url_trailing_delimiter.versionless_url_string
        )
      self._exclude_patterns = regex_util.Patterns(
          exclude_pattern_strings,
          # Confirm container URL ends in a delimiter.
          ignore_prefix_length=len(container_prefix),
      )
    else:
      self._exclude_patterns = None

    self._managed_folders_only = managed_folders_only
    self._ignore_symlinks = ignore_symlinks

    self._worker_id = 'process {} thread {}'.format(
        os.getpid(), threading.get_ident()
    )

  def execute(self, task_status_queue=None):
    del task_status_queue  # Unused.

    if self._managed_folders_only:
      managed_folder_setting = (
          folder_util.ManagedFolderSetting.LIST_WITHOUT_OBJECTS
      )
    else:
      managed_folder_setting = folder_util.ManagedFolderSetting.DO_NOT_LIST

    file_iterator = iter(
        wildcard_iterator.get_wildcard_iterator(
            self._container_query_path,
            exclude_patterns=self._exclude_patterns,
            fetch_encrypted_object_hashes=(
                properties.VALUES.storage.check_hashes.Get()
                != properties.CheckHashes.NEVER.value
            ),
            fields_scope=cloud_api.FieldsScope.RSYNC,
            files_only=not self._managed_folders_only,
            force_include_hidden_files=True,
            ignore_symlinks=self._ignore_symlinks,
            managed_folder_setting=managed_folder_setting,
        )
    )
    chunk_count = file_count = 0
    chunk_file_paths = []
    chunk_file_readers = []
    chunk_size = properties.VALUES.storage.rsync_list_chunk_size.GetInt()
    try:
      while True:
        resources_chunk = get_resource_chunk(
            file_iterator, chunk_size, self._gstmp_path
        )
        if not resources_chunk:
          break
        chunk_count += 1
        file_count += len(resources_chunk)
        log.status.Print(
            'At {}, worker {} listed {}...'.format(
                self._container_query_path, self._worker_id, file_count
            )
        )

        chunk_file_paths.append(
            rsync_command_util.get_hashed_list_file_path(
                self._container_query_path,
                chunk_count,
                is_managed_folder_list=self._managed_folders_only,
            )
        )
        if properties.VALUES.storage.use_url_based_rsync_sorting.GetBool():
          sorted_encoded_chunk = sorted(
              [
                  rsync_command_util.get_csv_line_from_resource(x)
                  for x in resources_chunk
              ],
              key=sorting_key_for_csv_line,
          )
        else:
          sorted_encoded_chunk = sorted([
              rsync_command_util.get_csv_line_from_resource(x)
              for x in resources_chunk
          ])

        sorted_encoded_chunk.append('')  # Add trailing newline.
        files.WriteFileContents(
            chunk_file_paths[-1],
            '\n'.join(sorted_encoded_chunk),
        )

      chunk_file_readers = [files.FileReader(path) for path in chunk_file_paths]
      with files.FileWriter(self._output_path, create_path=True) as file_writer:
        if properties.VALUES.storage.use_url_based_rsync_sorting.GetBool():
          file_writer.writelines(
              heapq.merge(*chunk_file_readers, key=sorting_key_for_csv_line)
          )
        else:
          file_writer.writelines(heapq.merge(*chunk_file_readers))

    except OSError as e:
      if e.errno == errno.EMFILE:
        raise errors.Error(
            'Too many open chunk files. Try increasing the'
            ' size with `gcloud config set storage/rsync_list_chunk_size`.'
            ' The current size is {}.'.format(chunk_size)
        )
      raise e

    finally:
      for reader in chunk_file_readers:
        try:
          reader.close()
        except Exception as e:  # pylint:disable=broad-except
          log.debug('Failed to close file reader {}: {}'.format(reader.name, e))
      for path in chunk_file_paths:
        rsync_command_util.try_to_delete_file(path)

  def __eq__(self, other):
    if not isinstance(other, type(self)):
      return NotImplemented
    return (
        self._container_query_path == other._container_query_path
        and self._exclude_patterns == other._exclude_patterns
        and self._managed_folders_only == other._managed_folders_only
        and self._ignore_symlinks == other._ignore_symlinks
        and self._output_path == other._output_path
    )
