# Copyright 2020 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
from __future__ import absolute_import
import base64
import json
import logging
import os

import requests  # pylint: disable=import-error
from lib.results import result_types

HTML_SUMMARY_MAX = 4096

_HTML_SUMMARY_ARTIFACT = '<text-artifact artifact-id="HTML Summary" />'
_TEST_LOG_ARTIFACT = '<text-artifact artifact-id="Test Log" />'

# Maps result_types to the luci test-result.proto.
# https://godoc.org/go.chromium.org/luci/resultdb/proto/v1#TestStatus
RESULT_MAP = {
    result_types.UNKNOWN: 'ABORT',
    result_types.PASS: 'PASS',
    result_types.FAIL: 'FAIL',
    result_types.CRASH: 'CRASH',
    result_types.TIMEOUT: 'ABORT',
    result_types.SKIP: 'SKIP',
    result_types.NOTRUN: 'SKIP',
}


def TryInitClient():
  """Tries to initialize a result_sink_client object.

  Assumes that rdb stream is already running.

  Returns:
    A ResultSinkClient for the result_sink server else returns None.
  """
  try:
    with open(os.environ['LUCI_CONTEXT']) as f:
      sink = json.load(f)['result_sink']
      return ResultSinkClient(sink)
  except KeyError:
    return None


class ResultSinkClient(object):
  """A class to store the sink's post configurations and make post requests.

  This assumes that the rdb stream has been called already and that the
  server is listening.
  """

  def __init__(self, context):
    base_url = 'http://%s/prpc/luci.resultsink.v1.Sink' % context['address']
    self.test_results_url = base_url + '/ReportTestResults'
    self.report_artifacts_url = base_url + '/ReportInvocationLevelArtifacts'
    self.update_invocation_url = base_url + '/UpdateInvocation'

    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json',
        'Authorization': 'ResultSink %s' % context['auth_token'],
    }
    self.session = requests.Session()
    self.session.headers.update(headers)

  def __enter__(self):
    return self

  def __exit__(self, exc_type, exc_value, traceback):
    self.close()

  def close(self):
    """Closes the session backing the sink."""
    self.session.close()

  def Post(self,
           test_id,
           status,
           duration,
           test_log,
           test_file,
           test_id_structured=None,
           variant=None,
           artifacts=None,
           failure_reason=None,
           html_artifact=None,
           tags=None):
    """Uploads the test result to the ResultSink server.

    This assumes that the rdb stream has been called already and that
    server is ready listening.

    Args:
      test_id: A string representing the test's name.
      status: A string representing if the test passed, failed, etc...
      duration: An int representing time in ms.
      test_log: A string representing the test's output.
      test_file: A string representing the file location of the test.
      test_id_structured: A dictionary containing structured test id fields.
      variant: An optional dict of variant key value pairs as the
          additional variant sent from test runners, which can override
          or add to the variants passed to `rdb stream` command.
      artifacts: An optional dict of artifacts to attach to the test.
      failure_reason: An optional string with the reason why the test failed.
          Should be None if the test did not fail.
      html_artifact: An optional html-formatted string to prepend to the test's
          log. Useful to encode click-able URL links in the test log, since that
          won't be formatted in the test_log.
      tags: An optional list of tuple of key name and value to prepend to the
          test's tags.

    Returns:
      N/A
    """
    assert status in RESULT_MAP
    expected = status in (result_types.PASS, result_types.SKIP)
    result_db_status = RESULT_MAP[status]

    tr = {
        'expected':
        expected,
        'status':
        result_db_status,
        'tags': [
            {
                'key': 'test_name',
                'value': test_id,
            },
            {
                # Status before getting mapped to result_db statuses.
                'key': 'raw_status',
                'value': status,
            }
        ],
        'testId':
        test_id,
        'testMetadata': {
            'name': test_id,
        }
    }

    if test_id_structured:
      tr['testIdStructured'] = test_id_structured

    if tags:
      tr['tags'].extend({
          'key': key_name,
          'value': value
      } for (key_name, value) in tags)

    if variant:
      tr['variant'] = {'def': variant}

    artifacts = artifacts or {}
    tr['summaryHtml'] = html_artifact if html_artifact else ''

    # If over max supported length of html summary, replace with artifact
    # upload.
    if (test_log
        and len(tr['summaryHtml']) + len(_TEST_LOG_ARTIFACT) > HTML_SUMMARY_MAX
        or len(tr['summaryHtml']) > HTML_SUMMARY_MAX):
      b64_summary = base64.b64encode(tr['summaryHtml'].encode()).decode()
      artifacts.update({'HTML Summary': {'contents': b64_summary}})
      tr['summaryHtml'] = _HTML_SUMMARY_ARTIFACT

    if test_log:
      # Upload the original log without any modifications.
      b64_log = base64.b64encode(test_log.encode()).decode()
      artifacts.update({'Test Log': {'contents': b64_log}})
      tr['summaryHtml'] += _TEST_LOG_ARTIFACT

    if artifacts:
      tr['artifacts'] = artifacts
    if failure_reason:
      tr['failureReason'] = {
          'primaryErrorMessage': _TruncateToUTF8Bytes(failure_reason, 1024)
      }

    if duration is not None:
      # Duration must be formatted to avoid scientific notation in case
      # number is too small or too large. Result_db takes seconds, not ms.
      # Need to use float() otherwise it does substitution first then divides.
      tr['duration'] = '%.9fs' % float(duration / 1000.0)

    if test_file and str(test_file).startswith('//'):
      tr['testMetadata']['location'] = {
          'file_name': test_file,
          'repo': 'https://chromium.googlesource.com/chromium/src',
      }

    res = self.session.post(url=self.test_results_url,
                            data=json.dumps({'testResults': [tr]}))
    res.raise_for_status()

  def ReportInvocationLevelArtifacts(self, artifacts):
    """Uploads invocation-level artifacts to the ResultSink server.

    This is for artifacts that don't apply to a single test but to the test
    invocation as a whole (eg: system logs).

    Args:
      artifacts: A dict of artifacts to attach to the invocation.
    """
    req = {'artifacts': artifacts}
    res = self.session.post(url=self.report_artifacts_url, data=json.dumps(req))
    res.raise_for_status()

  def UpdateInvocation(self, invocation, update_mask):
    """Update the invocation to the ResultSink server.

    Details can be found in the proto luci.resultsink.v1.UpdateInvocationRequest

    Args:
      invocation: a dict representation of luci.resultsink.v1.Invocation proto
      update_mask: a string representation of google.protobuf.FieldMask proto
    """
    req = {
        'invocation': invocation,
        'updateMask': update_mask,
    }
    res = self.session.post(url=self.update_invocation_url,
                            data=json.dumps(req))
    res.raise_for_status()

  def UpdateInvocationExtendedProperties(self, extended_properties, keys=None):
    """Update the extended_properties field of an invocation.

    Details can be found in the "extended_properties" field of the proto
    luci.resultdb.v1.Invocation.

    Args:
      extended_properties: a dict containing the content of extended_properties.
        The value in the dict shall be a dict containing a "@type" key
        representing the data schema, and corresponding data.
      keys: (Optional) a list of keys in extended_properties to add, replace,
        or remove. If a key exists in "keys", but not in "extended_properties",
        this is considered as deleting the key from the resultdb record side
        If None, the keys in "extended_properties" dict will be used.
    """
    # Sink server by default decodes payload with protojson, i.e. codecJSONV2
    # in https://source.chromium.org/search?q=f:server.go%20func:requestCodec
    # which requires loweCamelCase names in the json request.
    # For the value for update mask, see "JSON Encoding of Field Masks" in
    # https://protobuf.dev/reference/protobuf/google.protobuf/#field-masks
    invocation = {'extendedProperties': extended_properties}
    if not keys:
      keys = extended_properties.keys()
    mask_paths = ['extendedProperties.%s' % _ToCamelCase(key) for key in keys]
    update_mask = ','.join(mask_paths)
    self.UpdateInvocation(invocation, update_mask)


def _ToCamelCase(s):
  """Converts the string s from snake_case to lowerCamelCase."""

  elems = s.split('_')
  return elems[0] + ''.join(elem.capitalize() for elem in elems[1:])


def _TruncateToUTF8Bytes(s, length):
  """ Truncates a string to a given number of bytes when encoded as UTF-8.

  Ensures the given string does not take more than length bytes when encoded
  as UTF-8. Adds trailing ellipsis (...) if truncation occurred. A truncated
  string may end up encoding to a length slightly shorter than length because
  only whole Unicode codepoints are dropped.

  Args:
    s: The string to truncate.
    length: the length (in bytes) to truncate to.
  """
  try:
    encoded = s.encode('utf-8')
  # When encode throws UnicodeDecodeError in py2, it usually means the str is
  # already encoded and has non-ascii chars. So skip re-encoding it.
  except UnicodeDecodeError:
    encoded = s
  if len(encoded) > length:
    # Truncate, leaving space for trailing ellipsis (...).
    encoded = encoded[:length - 3]
    # Truncating the string encoded as UTF-8 may have left the final codepoint
    # only partially present. Pass 'ignore' to acknowledge and ensure this is
    # dropped.
    return encoded.decode('utf-8', 'ignore') + "..."
  return s