#!/usr/bin/env vpython3
# Copyright 2021 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Marshal various scheduling configs. Store them into the UFS datastore
through the Google datastore API.

By default, this script converts a few config-related protos into datastore
entities:

1. ConfigBundleList from 'hw_design/generated/configs.jsonproto'
2. DutAttributeList from
  '.../chromiumos/src/config/generated/dut_attributes.jsonproto'
3. FlatConfigList from 'hw_design/generated/flattened.jsonproto'
4. DeviceStabilityList from
  '.../chromiumos/infra/config/testingconfig/generated/device_stability.cfg'

The lists are parsed and individual entities are extracted. Using the datastore
client specified, it encodes the protos as datastore entities and stores them
into the UFS datastore.
"""

import argparse
import datetime
import logging
import os

from google.cloud import datastore

from checker import io_utils
from common import proto_utils

# type constants
CB_INPUT_TYPE = 'chromiumos.config.payload.ConfigBundleList'
CB_OUTPUT_TYPE = 'chromiumos.config.payload.ConfigBundle'
DA_INPUT_TYPE = 'chromiumos.test.api.DutAttributeList'
DA_OUTPUT_TYPE = 'chromiumos.test.api.DutAttribute'
FC_INPUT_TYPE = 'chromiumos.config.payload.FlatConfigList'
FC_OUTPUT_TYPE = 'chromiumos.config.payload.FlatConfig'
DEV_STAB_INPUT_TYPE = 'chromiumos.test.dut.DeviceStabilityList'
DEV_STAB_OUTPUT_TYPE = 'chromiumos.test.dut.DeviceStability'

# UFS services
UFS_DEV_PROJECT = 'unified-fleet-system-dev'
UFS_PROD_PROJECT = 'unified-fleet-system'

# datastore constants
CONFIG_BUNDLE_KIND = 'ConfigBundle'
DUT_ATTRIBUTE_KIND = 'DutAttribute'
FLAT_CONFIG_KIND = 'FlatConfig'
DEVICE_STABILITY_KIND = 'DeviceStability'


def get_ufs_project(env):
  """Return project name based on env argument."""
  if env == 'dev':
    return UFS_DEV_PROJECT
  if env == 'prod':
    return UFS_PROD_PROJECT
  raise RuntimeError('get_ufs_project: environment %s not supported' % env)


def generate_config_bundle_id(bundle):
  """Generate ConfigBundleEntity id as ${program_id}-${design_id}.

  It is possible the ConfigBundle has an empty design_list (e.g. because it is
  from a program repo). In this case, return None.
  """
  if not bundle.design_list:
    return None

  return (bundle.design_list[0].program_id.value + '-' +
          bundle.design_list[0].id.value).lower()


def handle_config_bundle_list(cb_list_path, client):
  """Take a path to a ConfigBundleList, iterate through the list and store into
  UFS datastore based on env.
  """
  cb_list = io_utils.read_json_proto(
      protodb.GetSymbol(CB_INPUT_TYPE)(), cb_list_path)

  for config_bundle in cb_list.values:
    update_config(config_bundle, client, flat=False)


def generate_flat_config_id(bundle):
  """Generate FlatConfigEntity id as ${program_id}-${design_id}-${design_config_id}
  if design_config_id is available. Else ${program_id}-${design_id}."""
  if bundle.hw_design_config.id.value:
    return (bundle.hw_design.program_id.value + '-' \
      + bundle.hw_design.id.value \
      + '-' + bundle.hw_design_config.id.value).lower()
  return (bundle.hw_design.program_id.value + '-' \
    + bundle.hw_design.id.value).lower()


def handle_flat_config_list(fc_list_path, client):
  """Take a path to a FlatConfigList, iterate through the list and store into
  UFS datastore based on env.
  """
  fc_list = io_utils.read_json_proto(
      protodb.GetSymbol(FC_INPUT_TYPE)(), fc_list_path)

  for flat_config in fc_list.values:
    update_config(flat_config, client, flat=True)


def update_config(config, client, flat=False):
  """Take a ConfigBundle or FlatConfig and store it an an entity in the UFS datastore."""
  if flat:
    kind = FLAT_CONFIG_KIND
    eid = generate_flat_config_id(config)
  else:
    kind = CONFIG_BUNDLE_KIND
    eid = generate_config_bundle_id(config)

    if not eid:
      logging.info('no eid for config %s, skipping', config)
      return

  logging.info('update_config: handling %s', eid)

  key = client.key(kind, eid)
  entity = datastore.Entity(
      key=key,
      exclude_from_indexes=['ConfigData'],
  )
  entity['ConfigData'] = config.SerializeToString()
  entity['Updated'] = datetime.datetime.now()

  logging.info('update_config: putting entity into datastore for %s', eid)
  client.put(entity)


def handle_dut_attribute_list(dut_attr_list_path, client):
  """Take a path to a DutAttributeList, iterate through the list and store into
  UFS datastore based on env.
  """
  dut_attr_list = io_utils.read_json_proto(
      protodb.GetSymbol(DA_INPUT_TYPE)(), dut_attr_list_path)

  for dut_attribute in dut_attr_list.dut_attributes:
    update_dut_attribute(dut_attribute, client)


def update_dut_attribute(attr, client):
  """Take a DutAttribute and store it in the UFS datastore as a DutAttributeEntity."""
  eid = attr.id.value
  logging.info('update_dut_attribute: handling %s', eid)

  key = client.key(DUT_ATTRIBUTE_KIND, eid)
  entity = datastore.Entity(
      key=key,
      exclude_from_indexes=['AttributeData'],
  )
  entity['AttributeData'] = attr.SerializeToString()
  entity['Updated'] = datetime.datetime.now()

  logging.info('update_dut_attribute: putting entity into datastore for %s',
               eid)
  client.put(entity)


def handle_device_stability_list(dev_stab_list_path, client):
  """Take a path to a DeviceStabilityList, iterate through the list and store
  into UFS datastore based on env.
  """
  dev_stab_list = io_utils.read_json_proto(
      protodb.GetSymbol(DEV_STAB_INPUT_TYPE)(), dev_stab_list_path)

  all_boards = {}
  for dev_stab in dev_stab_list.values:
    update_device_stability(dev_stab, client)
    for eid in dev_stab.dut_criteria[0].values:
      all_boards[eid] = True

  clean_up_device_stability(all_boards, client)


def clean_up_device_stability(all_boards, client):
  """Clean up boards/models that are deleted from device stability config files"""
  query = client.query(kind=DEVICE_STABILITY_KIND)
  query.keys_only()
  to_delete_keys = []
  for record in query.fetch():
    if record.key.name not in all_boards:
      to_delete_keys.append(record.key)

  client.delete_multi(to_delete_keys)


def update_device_stability(dev_stab, client):
  """Take a DeviceStability and store it in the UFS datastore as a
  DeviceStabilityEntity.
  """
  # TODO (justinsuen): May need to change this eventually. This assumes the use
  # of a single model (DutAttribute ID design_id) when defining a
  # DeviceStability entry.
  # http://cs/chromeos_internal/infra/config/testingconfig/target_test_requirements_config_helper.star?l=155
  for eid in dev_stab.dut_criteria[0].values:
    logging.info('update_device_stability: handling %s', eid)

    key = client.key(DEVICE_STABILITY_KIND, eid)
    entity = datastore.Entity(
        key=key,
        exclude_from_indexes=['StabilityData'],
    )
    entity['StabilityData'] = dev_stab.SerializeToString()
    entity['Updated'] = datetime.datetime.now()

    logging.info(
        'update_device_stability: putting entity into datastore for %s', eid)
    client.put(entity)


if __name__ == '__main__':
  logging.basicConfig(level=logging.INFO)
  parser = argparse.ArgumentParser(
      description=__doc__,
      formatter_class=argparse.RawDescriptionHelpFormatter,
  )

  parser.add_argument(
      '--env',
      type=str,
      default='dev',
      help='environment flag for UFS service',
  )

  # load database of protobuffer name -> Type
  protodb = proto_utils.create_symbol_db()
  options = parser.parse_args()
  ufs_ds_client = datastore.Client(
      project=get_ufs_project(options.env),
      namespace="os",
  )
  script_dir = os.path.dirname(os.path.realpath(__file__))

  handle_config_bundle_list("hw_design/generated/configs.jsonproto",
                            ufs_ds_client)
  handle_dut_attribute_list(
      os.path.realpath(
          os.path.join(
              script_dir,
              "../generated/dut_attributes.jsonproto",
          )), ufs_ds_client)
  handle_flat_config_list("hw_design/generated/flattened.jsonproto",
                          ufs_ds_client)
  handle_device_stability_list(
      os.path.realpath(
          os.path.join(
              script_dir,
              "../../config-internal/board_config/generated/device_stability.cfg"
          )), ufs_ds_client)