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

"""Manage tree status."""

from __future__ import print_function

import httplib
import json
import logging
import os
import re
import socket
import urllib
import urllib2

from chromite.cbuildbot import constants
from chromite.lib import osutils
from chromite.lib import timeout_util


CROS_TREE_STATUS_URL = 'https://chromiumos-status.appspot.com'
CROS_TREE_STATUS_JSON_URL = '%s/current?format=json' % CROS_TREE_STATUS_URL
CROS_TREE_STATUS_UPDATE_URL = '%s/status' % CROS_TREE_STATUS_URL

_USER_NAME = 'buildbot@chromium.org'
_PASSWORD_PATH = '/home/chrome-bot/.status_password_chromiumos'

# The tree status json file contains the following keywords.
TREE_STATUS_STATE = 'general_state'
TREE_STATUS_USERNAME = 'username'
TREE_STATUS_MESSAGE = 'message'
TREE_STATUS_DATE = 'date'
TREE_STATUS_CAN_COMMIT = 'can_commit_freely'

# These keywords in a status message are detected automatically to
# update the tree status.
MESSAGE_KEYWORDS = ('open', 'throt', 'close', 'maint')

# This is the delimiter to separate messages from different updates.
MESSAGE_DELIMITER = '|'


class PasswordFileDoesNotExist(Exception):
  """Raised when password file does not exist."""


class InvalidTreeStatus(Exception):
  """Raised when user wants to set an invalid tree status."""


def _GetStatusDict(status_url, raw_message=False):
  """Polls |status_url| and returns the retrieved tree status dictionary.

  This function gets a JSON response from |status_url|, and returns
  the dictionary of the tree status, if one exists and the http
  request was successful.

  The tree status dictionary contains:
    TREE_STATUS_USERNAME: User who posted the message (foo@chromium.org).
    TREE_STATUS_MESSAGE: The status message ("Tree is Open (CQ is good)").
    TREE_STATUS_CAN_COMMIT: Whether tree is commit ready ('true' or 'false').
    TREE_STATUS_STATE: one of constants.VALID_TREE_STATUSES.

  Args:
    status_url: The URL of the tree status to check.
    raw_message: Whether to return the raw message without stripping the
      "Tree is open/throttled/closed" string. Defaults to always strip.

  Returns:
    The tree status as a dictionary, if it was successfully retrieved.
    Otherwise None.
  """
  try:
    # Check for successful response code.
    response = urllib.urlopen(status_url)
    if response.getcode() == 200:
      data = json.load(response)
      if not raw_message:
        # Tree status message is usually in the form:
        #   "Tree is open/closed/throttled (reason for the tree closure)"
        # We want only the reason enclosed in the parentheses.
        # This is a best-effort parsing because user may post the message
        # in a form that we don't recognize.
        match = re.match(r'Tree is [\w\s\.]+\((.*)\)',
                         data.get(TREE_STATUS_MESSAGE, ''))
        data[TREE_STATUS_MESSAGE] = '' if not match else match.group(1)
      return data
  # We remain robust against IOError's.
  except IOError as e:
    logging.error('Could not reach %s: %r', status_url, e)


def _GetStatus(status_url):
  """Polls |status_url| and returns the retrieved tree status.

  This function gets a JSON response from |status_url|, and returns the
  value associated with the TREE_STATUS_STATE, if one exists and the
  http request was successful.

  Returns:
    The tree status, as a string, if it was successfully retrieved. Otherwise
    None.
  """
  status_dict = _GetStatusDict(status_url)
  if status_dict:
    return status_dict.get(TREE_STATUS_STATE)


def WaitForTreeStatus(status_url=None, period=1, timeout=1, throttled_ok=False):
  """Wait for tree status to be open (or throttled, if |throttled_ok|).

  Args:
    status_url: The status url to check i.e.
      'https://status.appspot.com/current?format=json'
    period: How often to poll for status updates.
    timeout: How long to wait until a tree status is discovered.
    throttled_ok: is TREE_THROTTLED an acceptable status?

  Returns:
    The most recent tree status, either constants.TREE_OPEN or
    constants.TREE_THROTTLED (if |throttled_ok|)

  Raises:
    timeout_util.TimeoutError if timeout expired before tree reached
    acceptable status.
  """
  if not status_url:
    status_url = CROS_TREE_STATUS_JSON_URL

  acceptable_states = set([constants.TREE_OPEN])
  verb = 'open'
  if throttled_ok:
    acceptable_states.add(constants.TREE_THROTTLED)
    verb = 'not be closed'

  timeout = max(timeout, 1)

  def _LogMessage(minutes_left):
    logging.info('Waiting for the tree to %s (%d minutes left)...', verb,
                 minutes_left)

  def _get_status():
    return _GetStatus(status_url)

  return timeout_util.WaitForReturnValue(
      acceptable_states, _get_status, timeout=timeout,
      period=period, side_effect_func=_LogMessage)


def IsTreeOpen(status_url=None, period=1, timeout=1, throttled_ok=False):
  """Wait for tree status to be open (or throttled, if |throttled_ok|).

  Args:
    status_url: The status url to check i.e.
      'https://status.appspot.com/current?format=json'
    period: How often to poll for status updates.
    timeout: How long to wait until a tree status is discovered.
    throttled_ok: Does TREE_THROTTLED count as open?

  Returns:
    True if the tree is open (or throttled, if |throttled_ok|). False if
    timeout expired before tree reached acceptable status.
  """
  if not status_url:
    status_url = CROS_TREE_STATUS_JSON_URL

  try:
    WaitForTreeStatus(status_url=status_url, period=period, timeout=timeout,
                      throttled_ok=throttled_ok)
  except timeout_util.TimeoutError:
    return False
  return True


def _GetPassword():
  """Returns the password for updating tree status."""
  if not os.path.exists(_PASSWORD_PATH):
    raise PasswordFileDoesNotExist(
        'Unable to retrieve password. %s does not exist',
        _PASSWORD_PATH)

  return osutils.ReadFile(_PASSWORD_PATH).strip()


def _UpdateTreeStatus(status_url, message):
  """Updates the tree status to |message|.

  Args:
    status_url: The tree status URL.
    message: The tree status text to post .
  """
  password = _GetPassword()
  params = urllib.urlencode({
      'message': message,
      'username': _USER_NAME,
      'password': password,
  })
  headers = {'Content-Type': 'application/x-www-form-urlencoded'}
  req = urllib2.Request(status_url, data=params, headers=headers)
  try:
    urllib2.urlopen(req)
  except (urllib2.URLError, httplib.HTTPException, socket.error) as e:
    logging.error('Unable to update tree status: %s', e)
    raise e
  else:
    logging.info('Updated tree status with message: %s', message)


def UpdateTreeStatus(status, message, announcer='cbuildbot', epilogue='',
                     status_url=None, dryrun=False):
  """Updates the tree status to |status| with additional |message|.

  Args:
    status: A status in constants.VALID_TREE_STATUSES.
    message: A string to display as part of the tree status.
    announcer: The announcer the message.
    epilogue: The string to append to |message|.
    status_url: The URL of the tree status to update.
    dryrun: If set, don't update the tree status.
  """
  if status_url is None:
    status_url = CROS_TREE_STATUS_UPDATE_URL

  if status not in constants.VALID_TREE_STATUSES:
    raise InvalidTreeStatus('%s is not a valid tree status.' % status)

  if status == 'maintenance':
    # This is a special case because "Tree is maintenance" is
    # grammatically incorrect.
    status = 'under maintenance'

  text_dict = {
      'status': status,
      'epilogue': epilogue,
      'announcer': announcer,
      'message': message,
      'delimiter': MESSAGE_DELIMITER
  }
  if epilogue:
    text = ('Tree is %(status)s (%(announcer)s: %(message)s %(delimiter)s '
            '%(epilogue)s)' % text_dict)
  else:
    text = 'Tree is %(status)s (%(announcer)s: %(message)s)' % text_dict

  if dryrun:
    logging.info('Would have updated the tree status with message: %s', text)
  else:
    _UpdateTreeStatus(status_url, text)


def ThrottleOrCloseTheTree(announcer, message, internal=None, buildnumber=None,
                           dryrun=False):
  """Throttle or close the tree with |message|.

  By default, this function throttles the tree with an updated
  message. If the tree is already not open, it will keep the original
  status (closed, maintenance) and only update the message. This
  ensures that we do not lower the severity of tree closure.

  In the case where the tree is not open, the previous tree status
  message is kept by prepending it to |message|, if possible. This
  ensures that the cause of the previous tree closure remains visible.

  Args:
    announcer: The announcer the message.
    message: A string to display as part of the tree status.
    internal: Whether the build is internal or not. Append the build type
      if this is set. Defaults to None.
    buildnumber: The build number to append.
    dryrun: If set, generate the message but don't update the tree status.
  """
  # Get current tree status.
  status_dict = _GetStatusDict(CROS_TREE_STATUS_JSON_URL)
  current_status = status_dict.get(TREE_STATUS_STATE)
  current_msg = status_dict.get(TREE_STATUS_MESSAGE)

  status = constants.TREE_THROTTLED
  if (constants.VALID_TREE_STATUSES.index(current_status) >
      constants.VALID_TREE_STATUSES.index(status)):
    # Maintain the current status if it is more servere than throttled.
    status = current_status

  epilogue = ''
  # Don't prepend the current status message if the tree is open.
  if current_status != constants.TREE_OPEN and current_msg:
    # Scan the current message and discard the text by the same
    # announcer.
    chunks = [x.strip() for x in current_msg.split(MESSAGE_DELIMITER)
              if '%s' % announcer not in x.strip()]
    current_msg = MESSAGE_DELIMITER.join(chunks)

    if any(x for x in MESSAGE_KEYWORDS if x.lower() in
           current_msg.lower().split()):
      # The waterfall scans the message for keywords to change the
      # tree status. Don't prepend the current status message if it
      # contains such keywords.
      logging.warning('Cannot prepend the previous tree status message because '
                      'there are keywords that may affect the tree state.')
    else:
      epilogue = current_msg

  if internal is not None:
    # 'p' stands for 'public.
    announcer += '-i' if internal else '-p'

  if buildnumber:
    announcer = '%s-%d' % (announcer, buildnumber)

  UpdateTreeStatus(status, message, announcer=announcer, epilogue=epilogue,
                   dryrun=dryrun)


def _OpenSheriffURL(sheriff_url):
  """Returns the content of |sheriff_url| or None if failed to open it."""
  try:
    response = urllib.urlopen(sheriff_url)
    if response.getcode() == 200:
      return response.read()
  except IOError as e:
    logging.error('Could not reach %s: %r', sheriff_url, e)


def GetSheriffEmailAddresses(sheriff_type):
  """Get the email addresses of the sheriffs or deputy.

  Args:
    sheriff_type: Type of the sheriff to look for. See the keys in
    constants.SHERIFF_TYPE_TO_URL.
      - 'tree': tree sheriffs
      - 'build': build deputy
      - 'lab' : lab sheriff
      - 'chrome': chrome gardener

  Returns:
    A list of email addresses.
  """
  if sheriff_type not in constants.SHERIFF_TYPE_TO_URL:
    raise ValueError('Unknown sheriff type: %s' % sheriff_type)

  urls = constants.SHERIFF_TYPE_TO_URL.get(sheriff_type)
  sheriffs = []
  for url in urls:
    # The URL displays a line: document.write('taco, burrito')
    raw_line = _OpenSheriffURL(url)
    if raw_line is not None:
      match = re.search(r'\'(.*)\'', raw_line)
      if match and match.group(1) != 'None (channel is sheriff)':
        sheriffs.extend(x.strip() for x in match.group(1).split(','))

  return ['%s%s' % (x, constants.GOOGLE_EMAIL) for x in sheriffs]


def GetHealthAlertRecipients(builder_run):
  """Returns a list of email addresses of the health alert recipients."""
  recipients = []
  for entry in builder_run.config.health_alert_recipients:
    if '@' in entry:
      # If the entry is an email address, add it to the list.
      recipients.append(entry)
    else:
      # Perform address lookup for a non-email entry.
      recipients.extend(GetSheriffEmailAddresses(entry))

  return recipients