# Copyright 2020 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Recipe for computing test owners, for use by Flake Fetcher

Reads a test manifest from GCS that's expected to be a list of JSON objects,
each of which contains a test name and the test's GN label.

Then the recipe uses a Fuchsia checkout to compute test owners and Monorail
components based on the nearest OWNERS file to the GN label. Finally, the
recipe uploads the owners data to another JSON file in GCS that Flake Fetcher
uses to assign bugs for flaky tests to the tests' owners and put the bugs in
the appropriate components.
"""

from recipe_engine.recipe_api import Property

DEPS = [
    "fuchsia/bigquery",
    "fuchsia/bqupload",
    "fuchsia/checkout",
    "fuchsia/gsutil",
    "fuchsia/status_check",
    "recipe_engine/context",
    "recipe_engine/file",
    "recipe_engine/path",
    "recipe_engine/properties",
    "recipe_engine/step",
    "recipe_engine/time",
]

PROPERTIES = {
    "manifest": Property(kind=str, help="Jiri manifest to use"),
    "remote": Property(kind=str, help="Remote manifest repository"),
    "owners_gcs_bucket": Property(kind=str, help="GCS bucket to upload test owners to"),
    "owners_gcs_path": Property(
        kind=str,
        help="Path within `owners_gcs_bucket` to upload the test owners manifest to",
    ),
    "dry_run": Property(kind=bool, help="Don't upload results to GCS", default=False),
}

BIGQUERY_PROJECT = "fuchsia-infra"

# Queries ResultDB to find all unique "test_id" and "gn_label" combinations.
TEST_NAMES_QUERY = """
SELECT DISTINCT
  test_id AS test_name,
  SPLIT(tags.value, ':')[OFFSET(0)] AS gn_label
FROM
  `fuchsia-infra.resultdb.try`,
  UNNEST(tags) AS tags
WHERE
  partition_time >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 7 day)
  AND tags.key = 'gn_label'
  AND tags.value != ''
UNION DISTINCT
SELECT DISTINCT
  test_id AS test_name,
  SPLIT(tags.value, ':')[OFFSET(0)] AS gn_label
FROM
  `fuchsia-infra.resultdb.ci`,
  UNNEST(tags) AS tags
WHERE
  partition_time >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 7 day)
  AND tags.key = 'gn_label'
  AND tags.value != ''
"""


def RunSteps(api, manifest, remote, owners_gcs_bucket, owners_gcs_path, dry_run):
    checkout_dir = api.path["start_dir"].join("checkout")
    api.checkout.with_options(
        path=checkout_dir,
        manifest=manifest,
        remote=remote,
        # No need to fetch CIPD packages, we only need to look at OWNERS
        # files and those are always checked into git.
        fetch_packages=False,
    )

    workdir = api.path.mkdtemp("workdir")
    test_manifest_path = workdir.join("test_manifest.json")
    api.bigquery.project = BIGQUERY_PROJECT
    api.bigquery.query(
        "query test names",
        TEST_NAMES_QUERY,
        test_manifest_path,
    )
    if dry_run:
        api.file.read_text("log test_manifest.json", test_manifest_path)

    owners_path = workdir.join("test_owners.json")
    with api.context(cwd=checkout_dir):
        api.step(
            "find owners",
            [api.resource("find_owners.py"), test_manifest_path, owners_path],
        )

    tests = api.file.read_json(
        "read test_owners.json",
        owners_path,
        test_data=[
            {
                "test_name": "test1",
                "owners": ["x@gmail.com", "y@gmail.com"],
            },
            {
                "test_name": "test2",
                "owners": ["x@gmail.com", "z@gmail.com"],
            },
        ],
    )

    date = str(api.time.utcnow().date())  # E.g. "2021-05-27".
    bq_rows = []
    # The BQ table has a slightly different schema from that of
    # test_owners.json; it has one row per (test, owner) pair.
    for test in tests:
        for owner in test.pop("owners"):
            row = {"owner": owner, "date_created": date}
            row.update(test)
            bq_rows.append(row)

    if not dry_run:
        api.gsutil.upload(
            name="upload to GCS",
            src=owners_path,
            bucket=owners_gcs_bucket,
            dst=owners_gcs_path,
        )
        api.bqupload.insert(
            "upload to bigquery",
            project="fuchsia-infra",
            dataset="test_owners",
            table="test_owners",
            rows=bq_rows,
        )


def GenTests(api):
    def properties(**kwargs):
        props = {
            "manifest": "fuchsia/flower",
            "remote": "https://fuchsia.googlesource.com/integration",
            "owners_gcs_bucket": "owners",
            "owners_gcs_path": "test_owners.json",
        }
        props.update(kwargs)
        return api.properties(**props)

    yield api.status_check.test("basic") + properties()

    yield api.status_check.test("dry_run") + properties(dry_run=True)