// Copyright 2018 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.

// This file defines the LKGS CLI tool which computes the last-known-good
// (-jiri)-snapshot given a set of parameters.

// TODO (nmulcahey): Add tests for this tool.

package main

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path"
	"strings"

	"google.golang.org/genproto/protobuf/field_mask"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/auth/client/authcli"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/hardcoded/chromeinfra"
	"go.chromium.org/luci/logdog/client/coordinator"
	"go.chromium.org/luci/logdog/common/renderer"
	"go.chromium.org/luci/logdog/common/types"
)

// Implementing flag.Value
type builderFlag []string

func (builders *builderFlag) Set(builder string) error {
	*builders = append(*builders, builder)
	return nil
}

func (builders *builderFlag) String() string {
	builderStrings := []string{}
	for _, builder := range *builders {
		builderStrings = append(builderStrings, string(builder))
	}
	return strings.Join(builderStrings, ",")
}

type pollErrorType int

const (
	noBuilds           = 0
	noGreenBuilds      = 1
	buildBucketFailure = 2
)

type buildPollError struct {
	builder string
	kind    pollErrorType
}

func (bpe buildPollError) Error() string {
	switch kind := bpe.kind; kind {
	case noBuilds:
		return fmt.Sprintf("no builds found for builder %s", bpe.builder)
	case noGreenBuilds:
		return fmt.Sprintf("no green build found for builder %s", bpe.builder)
	case buildBucketFailure:
		return fmt.Sprintf("buildbucker error %s", bpe.builder)
	}
	return fmt.Sprintf("Unknown Error: %s", bpe.builder)
}

var (
	host       string
	builderIDs builderFlag
	output     string

	// LUCI flags used to parse command-line authentication options.
	authFlags authcli.Flags
	buildMap  = make(map[string][]*buildbucketpb.Build)
)

func init() {
	flag.StringVar(&host, "host", "cr-buildbucket.appspot.com", "the buildbucket host to use (default is cr-buildbucket.appspot.com)")
	flag.Var(&builderIDs, "builder-id", "[repeatable] name of the builders to use as a reference (e.g. fuchsia/ci/garnet-x64)")
	flag.StringVar(&output, "output-file", "", "name of the file to write snapshot to (default is stdout)")

	authFlags = authcli.Flags{}
	authFlags.Register(flag.CommandLine, chromeinfra.DefaultAuthOptions())
}

// praseBuilderID parses a builder ID of the form "fuchsia/ci/garnet-x64".
func parseBuilderID(builderID string) (*buildbucketpb.BuilderID, error) {
	components := strings.SplitN(builderID, "/", 3)
	if len(components) != 3 {
		return nil, fmt.Errorf("failed to parse builder ID")
	}
	return &buildbucketpb.BuilderID{
		Project: components[0],
		Bucket:  components[1],
		Builder: components[2],
	}, nil
}

// getLastKnownGoodBuild retrieves the last known good build for N builder IDs.
// TODO(nmulcahey): This should use BuildsClient.Batch to get builds from all builders.
func getLastKnownGoodBuild(ctx context.Context, buildsClient buildbucketpb.BuildsClient, builderIDs []*buildbucketpb.BuilderID, previousBuild *buildbucketpb.Build) (*buildbucketpb.Build, error) {
	builderID := builderIDs[0]
	if len(buildMap[builderID.String()]) == 0 {
		res, err := buildsClient.SearchBuilds(ctx, &buildbucketpb.SearchBuildsRequest{
			Predicate: &buildbucketpb.BuildPredicate{
				Builder: builderID,
				Status:  buildbucketpb.Status_SUCCESS,
			},
			// Retrieve only the Infra & Input fields
			// Infra.Logdog is how we retrieve the Jiri snapshot
			// Input.GitilesCommit.Id is what we match builds on
			Fields: &field_mask.FieldMask{
				Paths: []string{"builds.*.infra", "builds.*.input"},
			},
		})
		if err != nil {
			return nil, buildPollError{builderID.String(), buildBucketFailure}
		}
		if len(res.Builds) == 0 {
			return nil, buildPollError{builderID.String(), noBuilds}
		}
		buildMap[builderID.String()] = res.Builds
	}
	// TODO(nmulcahey): Use iteration instead of recursion here.
	for _, currentBuild := range buildMap[builderID.String()] {
		// If the currentBuild has no GitilesCommit, someone manually triggered it and we need to skip it
		if currentBuild.Input.GitilesCommit == nil {
			continue
		}
		// If previousBuild is nil, we are at the top-level and can recurse, otherwise only recurse
		// if we found a matching green build
		if previousBuild == nil || currentBuild.Input.GitilesCommit.Id == previousBuild.Input.GitilesCommit.Id {
			// If processing the last builderID, return when a match is found
			if len(builderIDs) > 1 {
				nextBuild, err := getLastKnownGoodBuild(ctx, buildsClient, builderIDs[1:], currentBuild)
				if err != nil {
					// One of the builders has no builds; short-circuit to the top-level
					// and exit with error
					if obj, ok := err.(buildPollError); ok == true {
						switch obj.kind {
						case noGreenBuilds:
							if previousBuild == nil {
								continue
							}
							return nil, err
						case noBuilds:
							fallthrough
						case buildBucketFailure:
							fallthrough
						default:
							return nil, err
						}
					}
					continue
				}
				return nextBuild, nil
			}
			return currentBuild, nil
		}
	}
	return nil, buildPollError{builderID.String(), noGreenBuilds}
}

// getSnapshot retrieves the jiri snapshot from LogDog related to the build using that build's LogDog details.
func getSnapshot(ctx context.Context, client *http.Client, logdog *buildbucketpb.BuildInfra_LogDog) ([]byte, error) {
	coordClient := coordinator.NewClient(&prpc.Client{
		C:       client,
		Host:    logdog.Hostname,
		Options: prpc.DefaultOptions(),
	})
	logProject := types.ProjectName(logdog.Project)
	// TODO(mknyszek): Consider making these snapshots easier to find. This should be resilient against
	// step name changes, so long as the log itself has the same name, but it is kind of a hack.
	logPath := path.Join(logdog.Prefix, "+", "**", "snapshot_contents", "*")

	// Perform the query, capturing exactly one log stream and erroring otherwise.
	var log *coordinator.LogStream
	err := coordClient.Query(ctx, logProject, logPath, coordinator.QueryOptions{}, func(s *coordinator.LogStream) bool {
		log = s
		return false
	})
	switch {
	case err != nil:
		return nil, err
	case log == nil:
		return nil, fmt.Errorf(
			"unable to find jiri snapshot in project %s at path %s",
			logdog.Project, logPath)
	}

	// Read the source manifest from the log stream.
	var buf bytes.Buffer
	_, err = buf.ReadFrom(&renderer.Renderer{
		Source: coordClient.Stream(logProject, log.Path).Fetcher(ctx, nil),
		Raw:    true,
	})
	return buf.Bytes(), err
}

func main() {
	flag.Parse()

	if builderIDs == nil {
		flag.PrintDefaults()
		return
	}

	ids := []*buildbucketpb.BuilderID{}
	for _, builderID := range builderIDs {
		id, err := parseBuilderID(builderID)
		if err != nil {
			log.Fatalf(err.Error())
		}
		ids = append(ids, id)
	}

	opts, err := authFlags.Options()
	if err != nil {
		log.Fatalf(err.Error())
	}

	ctx := context.Background()
	authenticator := auth.NewAuthenticator(ctx, auth.OptionalLogin, opts)
	client, err := authenticator.Client()
	if err != nil {
		log.Fatalf(err.Error())
	}

	buildsClient := buildbucketpb.NewBuildsPRPCClient(&prpc.Client{
		C:    client,
		Host: host,
	})

	build, err := getLastKnownGoodBuild(ctx, buildsClient, ids, nil)
	if err != nil {
		log.Fatalf(err.Error())
	}

	snapshotBytes, err := getSnapshot(ctx, client, build.Infra.Logdog)
	if err != nil {
		log.Fatalf(err.Error())
	}

	var outputFile *os.File
	if output == "" {
		outputFile = os.Stdout
	} else {
		outputFile, err = os.Create(output)
		if err != nil {
			log.Fatalf(err.Error())
		}
		defer outputFile.Close()
	}
	_, err = outputFile.Write(snapshotBytes)
	if err != nil {
		log.Fatalf("writing output: %v", err)
	}
}