# Copyright 2019 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.

import logging
import os
import sys

from io import open

import psutil

from google.protobuf import json_format as jsonpb
from google.protobuf import text_format as textpb

from PB.go.chromium.org.luci.buildbucket.proto import common
from PB.go.chromium.org.luci.buildbucket.proto.build import Build

from ....third_party import luci_context
from ....util import fix_json_object

from ...engine import RecipeEngine
from ...global_shutdown import install_signal_handlers
from ...step_runner.subproc import SubprocessStepRunner
from ...stream.invariants import StreamEngineInvariants
from ...stream.luci import LUCIStreamEngine
from ...warn.record import NULL_WARNING_RECORDER

from . import RunBuildContractViolation


LOG = logging.getLogger(__name__)


def _contract_in_env(key):
  if key not in os.environ:
    raise RunBuildContractViolation('Expected $%s in environment.' % key)
  return os.environ[key]


def _contract_in_luci_context(section_name, key):
  section = luci_context.read(section_name)
  if section is None or key not in section:
    raise RunBuildContractViolation('Expected %r in $LUCI_CONTEXT[%r].' % (
      key, section_name))
  return section[key]


def _synth_properties(build, current_properties):
  # TODO(iannucci): expose this data natively as a Build message.
  synth_props = {
    '$recipe_engine/runtime': {
      'is_experimental': build.input.experimental,
    },
    '$recipe_engine/buildbucket': {
      'build': jsonpb.MessageToDict(build),
    },
    '$recipe_engine/path': {
      'temp_dir': _contract_in_env('TMP'),
      'cache_dir': _contract_in_luci_context('luciexe', 'cache_dir'),
    },
  }

  # TODO(iannucci): These are all deprecated and have proper apis.
  #
  # When we have the warnings functionality in recipes, update the properties
  # module to issue warnings for accessing these.
  if 'buildername' not in current_properties and build.builder.builder:
    synth_props['buildername'] = build.builder.builder
  if 'buildnumber' not in current_properties and build.number:
    synth_props['buildnumber'] = build.number
  if 'bot_id' not in current_properties and 'SWARMING_BOT_ID' in os.environ:
    synth_props['bot_id'] = os.environ['SWARMING_BOT_ID']
  LOG.info('Synthesized properties: %r', synth_props)
  return synth_props

def _tweak_env():
  # These tweaks are recipe-engine-specific tweaks to be compatible with the
  # behavior of `recipes.py run`.
  os.environ['PYTHONUNBUFFERED'] = '1'
  os.environ['PYTHONIOENCODING'] = 'UTF-8'


def main(args):
  with install_signal_handlers():
    return _main_impl(args)


def _main_impl(args):
  LOG.info('luciexe started, parsing Build message from stdin.')
  build = Build()
  build.ParseFromString(open(sys.stdin.fileno(), 'rb').read())
  LOG.info('finished parsing Build message')
  LOG.debug('build proto: %s', jsonpb.MessageToJson(build))

  properties = jsonpb.MessageToDict(build.input.properties)
  properties.update(_synth_properties(build, properties))
  properties = fix_json_object(properties)

  _tweak_env()

  luciexe_engine = LUCIStreamEngine(args.build_proto_stream_jsonpb)

  raw_result = None
  with StreamEngineInvariants.wrap(luciexe_engine) as stream_engine:
    try:
      raw_result, _ = RecipeEngine.run_steps(
        args.recipe_deps, properties, stream_engine,
        SubprocessStepRunner(), NULL_WARNING_RECORDER,
        os.environ, os.getcwd(), luci_context.read_full(),
        psutil.cpu_count(), psutil.virtual_memory().total)
      stream_engine.write_result(raw_result)
    except:
      LOG.exception("RecipeEngine.run_steps uncaught exception.")
      raise

  if args.output:
    try:
      with open(args.output, 'wb') as f:
        final_build = luciexe_engine.current_build_proto
        if args.output.endswith('.pb') :
          f.write(final_build.SerializeToString(deterministic=True))
        elif args.output.endswith('.json'):
          f.write(jsonpb.MessageToJson(
            final_build,
            preserving_proto_field_name=True,
            sort_keys=True,
            indent=2,
          ))
        elif args.output.endswith('.textpb'):
          f.write(textpb.MessageToString(final_build))
    except:
      LOG.exception("Error while writing final build to output file.")
      raise

  return 0 if luciexe_engine.was_successful else 1

@@ -3804,51 +3804,8 @@
 ry:%0A
-      with open(args.output, 'wb') as f:%0A  
     
@@ -3861,18 +3861,16 @@
 o%0A      
- 
- 
 if args.
@@ -3902,26 +3902,22 @@
         
-  f.write(
+out = 
 final_bu
@@ -3957,26 +3957,23 @@
 ic=True)
-)
 %0A      
-  
 elif arg
@@ -4008,26 +4008,22 @@
         
-  f.write(
+out = 
 jsonpb.M
@@ -4159,25 +4159,26 @@
 nt=2
-,%0A          ))%0A  
+).encode('utf-8')%0A
     
@@ -4229,18 +4229,14 @@
     
-  f.write(
+out = 
 text
@@ -4266,16 +4266,182 @@
 l_build)
+.encode('utf-8')%0A      else:%0A        raise ValueError('unknown output file extension: %25r' %25 args.output)%0A%0A      with open(args.output, 'wb') as f:%0A        f.write(out
 )%0A    ex

write final build after encode to bytes