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

import json
import logging
import os
import re
import shutil
import tempfile

from catapult_base import cloud_storage  # pylint: disable=import-error


def AssertValidCloudStorageBucket(bucket):
  is_valid = bucket in (None,
                        cloud_storage.PUBLIC_BUCKET,
                        cloud_storage.PARTNER_BUCKET,
                        cloud_storage.INTERNAL_BUCKET)
  if not is_valid:
    raise ValueError("Cloud storage privacy bucket %s is invalid" % bucket)


class ArchiveError(Exception):
  pass


class WprArchiveInfo(object):
  def __init__(self, file_path, data, bucket):
    AssertValidCloudStorageBucket(bucket)
    self._file_path = file_path
    self._base_dir = os.path.dirname(file_path)
    self._data = data
    self._bucket = bucket

    # Ensure directory exists.
    if not os.path.exists(self._base_dir):
      os.makedirs(self._base_dir)

    # Map from the relative path (as it appears in the metadata file) of the
    # .wpr file to a list of story names it supports.
    self._wpr_file_to_story_names = data['archives']

    # Map from the story name to a relative path (as it appears
    # in the metadata file) of the .wpr file.
    self._story_name_to_wpr_file = dict()
    # Find out the wpr file names for each story.
    for wpr_file in data['archives']:
      story_names = data['archives'][wpr_file]
      for story_name in story_names:
        self._story_name_to_wpr_file[story_name] = wpr_file
    self.temp_target_wpr_file_path = None

  @classmethod
  def FromFile(cls, file_path, bucket):
    if os.path.exists(file_path):
      with open(file_path, 'r') as f:
        data = json.load(f)
        return cls(file_path, data, bucket)
    return cls(file_path, {'archives': {}}, bucket)

  def DownloadArchivesIfNeeded(self):
    """Downloads archives iff the Archive has a bucket parameter and the user
    has permission to access the bucket.

    Raises cloud storage Permissions or Credentials error when there is no
    local copy of the archive and the user doesn't have permission to access
    the archive's bucket.

    Warns when a bucket is not specified or when the user doesn't have
    permission to access the archive's bucket but a local copy of the archive
    exists.
    """
    # Download all .wpr files.
    if not self._bucket:
      logging.warning('Story set in %s has no bucket specified, and '
                      'cannot be downloaded from cloud_storage.', )
      return
    assert 'archives' in self._data, 'Invalid data format in %s. \'archives\'' \
                                     ' field is needed' % self._file_path
    for archive_path in self._data['archives']:
      archive_path = self._WprFileNameToPath(archive_path)
      try:
        cloud_storage.GetIfChanged(archive_path, self._bucket)
      except (cloud_storage.CredentialsError, cloud_storage.PermissionError):
        if os.path.exists(archive_path):
          # If the archive exists, assume the user recorded their own and
          # simply warn.
          logging.warning('Need credentials to update WPR archive: %s',
                          archive_path)
        else:
          logging.error("You either aren't authenticated or don't have "
                        "permission to use the archives for this page set."
                        "\nYou may need to run gsutil config."
                        "\nYou can find instructions for gsutil config at: "
                        "http://www.chromium.org/developers/telemetry/"
                        "upload_to_cloud_storage")
          raise

  def WprFilePathForStory(self, story):
    if self.temp_target_wpr_file_path:
      return self.temp_target_wpr_file_path
    wpr_file = self._story_name_to_wpr_file.get(story.display_name, None)
    if wpr_file is None and hasattr(story, 'url'):
      # Some old pages always use the URL to identify a page rather than the
      # display_name, so try to look for that.
      wpr_file = self._story_name_to_wpr_file.get(story.url, None)
    if wpr_file:
      return self._WprFileNameToPath(wpr_file)
    return None

  def AddNewTemporaryRecording(self, temp_wpr_file_path=None):
    if temp_wpr_file_path is None:
      temp_wpr_file_handle, temp_wpr_file_path = tempfile.mkstemp()
      os.close(temp_wpr_file_handle)
    self.temp_target_wpr_file_path = temp_wpr_file_path

  def AddRecordedStories(self, stories, upload_to_cloud_storage=False):
    if not stories:
      os.remove(self.temp_target_wpr_file_path)
      return

    (target_wpr_file, target_wpr_file_path) = self._NextWprFileName()
    for story in stories:
      self._SetWprFileForStory(story.display_name, target_wpr_file)
    shutil.move(self.temp_target_wpr_file_path, target_wpr_file_path)

    # Update the hash file.
    target_wpr_file_hash = cloud_storage.CalculateHash(target_wpr_file_path)
    with open(target_wpr_file_path + '.sha1', 'wb') as f:
      f.write(target_wpr_file_hash)
      f.flush()

    self._WriteToFile()
    self._DeleteAbandonedWprFiles()

    # Upload to cloud storage
    if upload_to_cloud_storage:
      if not self._bucket:
        logging.warning('StorySet must have bucket specified to upload '
                        'stories to cloud storage.')
        return
      try:
        cloud_storage.Insert(self._bucket, target_wpr_file_hash,
                             target_wpr_file_path)
      except cloud_storage.CloudStorageError, e:
        logging.warning('Failed to upload wpr file %s to cloud storage. '
                        'Error:%s' % target_wpr_file_path, e)

  def _DeleteAbandonedWprFiles(self):
    # Update the metadata so that the abandoned wpr files don't have
    # empty story name arrays.
    abandoned_wpr_files = self._AbandonedWprFiles()
    for wpr_file in abandoned_wpr_files:
      del self._wpr_file_to_story_names[wpr_file]
      # Don't fail if we're unable to delete some of the files.
      wpr_file_path = self._WprFileNameToPath(wpr_file)
      try:
        os.remove(wpr_file_path)
      except Exception:
        logging.warning('Failed to delete file: %s' % wpr_file_path)

  def _AbandonedWprFiles(self):
    abandoned_wpr_files = []
    for wpr_file, story_names in (
        self._wpr_file_to_story_names.iteritems()):
      if not story_names:
        abandoned_wpr_files.append(wpr_file)
    return abandoned_wpr_files

  def _WriteToFile(self):
    """Writes the metadata into the file passed as constructor parameter."""
    metadata = dict()
    metadata['description'] = (
        'Describes the Web Page Replay archives for a story set. '
        'Don\'t edit by hand! Use record_wpr for updating.')
    metadata['archives'] = self._wpr_file_to_story_names.copy()
    # Don't write data for abandoned archives.
    abandoned_wpr_files = self._AbandonedWprFiles()
    for wpr_file in abandoned_wpr_files:
      del metadata['archives'][wpr_file]

    with open(self._file_path, 'w') as f:
      json.dump(metadata, f, indent=4)
      f.flush()

  def _WprFileNameToPath(self, wpr_file):
    return os.path.abspath(os.path.join(self._base_dir, wpr_file))

  def _NextWprFileName(self):
    """Creates a new file name for a wpr archive file."""
    # The names are of the format "some_thing_number.wpr". Read the numbers.
    highest_number = -1
    base = None
    for wpr_file in self._wpr_file_to_story_names:
      match = re.match(r'(?P<BASE>.*)_(?P<NUMBER>[0-9]+)\.wpr', wpr_file)
      if not match:
        raise Exception('Illegal wpr file name ' + wpr_file)
      highest_number = max(int(match.groupdict()['NUMBER']), highest_number)
      if base and match.groupdict()['BASE'] != base:
        raise Exception('Illegal wpr file name ' + wpr_file +
                        ', doesn\'t begin with ' + base)
      base = match.groupdict()['BASE']
    if not base:
      # If we're creating a completely new info file, use the base name of the
      # story set file.
      base = os.path.splitext(os.path.basename(self._file_path))[0]
    new_filename = '%s_%03d.wpr' % (base, highest_number + 1)
    return new_filename, self._WprFileNameToPath(new_filename)

  def _SetWprFileForStory(self, story_name, wpr_file):
    """For modifying the metadata when we're going to record a new archive."""
    old_wpr_file = self._story_name_to_wpr_file.get(story_name, None)
    if old_wpr_file:
      self._wpr_file_to_story_names[old_wpr_file].remove(story_name)
    self._story_name_to_wpr_file[story_name] = wpr_file
    if wpr_file not in self._wpr_file_to_story_names:
      self._wpr_file_to_story_names[wpr_file] = []
    self._wpr_file_to_story_names[wpr_file].append(story_name)