// Copyright 2021 The Bazel Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Binary rbe_configs_upload uploads the artifacts generated by rbe_configs_gen to GCS. This tool
// is meant for internal use by the owners of this repository only.
// This tool will upload the given configs tarball & manifest to the following paths on GCS:
// - gs://rbe-bazel-toolchains/configs/latest
// - - rbe_default.tar (The configs tarball)
// - - manifest.json (The JSON manifest)
// - gs://rbe-bazel-toolchains/configs/bazel_<version>/latest
// - - rbe_default.tar (The configs tarball)
// - - manifest.json (The JSON manifest)
// This tool will upload the above files even if the config tarball hasn't changed. This can happen
// if there's been no new Bazel release or toolchain container release since the last time this tool
// was run. Thus, the above GCS artifacts are unstable in the sense that their contents can change
// if either a new Bazel or toolchain container is released. This is to avoid users depending on
// these GCS artifacts in production. Instead, users should copy the artifacts into a GCS bucket
// or other remote location under their control.
package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"time"

	"cloud.google.com/go/storage"
	"github.com/bazelbuild/bazel-toolchains/pkg/monitoring"
	"github.com/bazelbuild/bazel-toolchains/pkg/rbeconfigsgen"
)

var (
	configsTarball        = flag.String("configs_tarball", "", "Path to the configs tarball generated by rbe_configs_gen to be uploaded to GCS.")
	configsManifest       = flag.String("configs_manifest", "", "Path to the JSON manifest generated by rbe_configs_gen.")
	enableMonitoring      = flag.Bool("enable_monitoring", false, "(Optional) Enables reporting reporting results to Google Cloud Monitoring. Defaults to false.")
	monitoringProjectID   = flag.String("monitoring_project_id", "", "GCP Project ID where monitoring results will be reported. Required if --enable_monitoring is true.")
	monitoringDockerImage = flag.String("monitoring_docker_image", "", "Name of the toolchain docker image to be reported as a string label to monitoring. Required if --enable_monitoring is true.")
)

// manifest is the metadata about the configs that'll be uploaded to GCS.
type manifest struct {
	// Wrap around the manifest produced by rbe_configs_gen.
	rbeconfigsgen.Manifest
	// UploadTime is the time this manifest was uploaded. For information only.
	UploadTime time.Time `json:"upload_time"`
}

// manifestFromFile loads the JSON manifest (in the format produced by rbe_configs_gen) from the
// given file and injects the current time in the returned manifest object.
func manifestFromFile(filePath string) (*manifest, error) {
	blob, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, fmt.Errorf("unable to read manifest JSON file %q: %w", filePath, err)
	}
	m := &manifest{}
	if err := json.Unmarshal(blob, m); err != nil {
		return nil, fmt.Errorf("error parsing contents of manifest file %q as JSON: %w", filePath, err)
	}
	// Ensure the mandatory fields used by this binary were specified.
	if len(m.BazelVersion) == 0 {
		return nil, fmt.Errorf("manifest %q did not specify bazel version", filePath)
	}
	m.UploadTime = time.Now()
	return m, nil
}

// storageClient represents the GCS client.
type storageClient struct {
	client *storage.Client
	// bucketName is the GCS bucket all artifacts will be uploaded to.
	bucketName string
}

func newStorage(ctx context.Context) (*storageClient, error) {
	c, err := storage.NewClient(ctx)
	if err != nil {
		return nil, err
	}
	return &storageClient{
		client:     c,
		bucketName: "rbe-toolchain",
	}, nil
}

// upload uploads the bytes represented by the given reader as the given GCS object name.
func (s *storageClient) upload(ctx context.Context, r io.Reader, objectName string) error {
	w := s.client.Bucket(s.bucketName).Object(objectName).NewWriter(ctx)
	if _, err := io.Copy(w, r); err != nil {
		return fmt.Errorf("error while uploading to GCS object %q: %w", objectName, err)
	}
	// The actual upload might happen after Close is called so we need to capture any errors.
	if err := w.Close(); err != nil {
		return fmt.Errorf("error finishing upload to GCS object %q: %w", objectName, err)
	}
	return nil
}

// uploadArtifacts uploads the given blob of bytes representing a JSON manifest and the configs
// tarball at the given path to the given GCS directory.
func (s *storageClient) uploadArtifacts(ctx context.Context, manifest []byte, tarballPath, remoteDir string) error {
	f, err := os.Open(tarballPath)
	if err != nil {
		return fmt.Errorf("unable to open configs tarball file %q: %w", tarballPath, err)
	}
	defer f.Close()

	if err := s.upload(ctx, bytes.NewBuffer(manifest), fmt.Sprintf("%s/manifest.json", remoteDir)); err != nil {
		return fmt.Errorf("error uploading manifest to GCS: %w", err)
	}

	if err := s.upload(ctx, f, fmt.Sprintf("%s/rbe_default.tar", remoteDir)); err != nil {
		return fmt.Errorf("error uploading configs tarball to GCS: %w", err)
	}
	return nil
}

// printFlag prints flag values with the intent of allowing easy copy paste of flags to rerun this
// binary.
func printFlags() {
	log.Println("rbe_configs_upload.go \\")
	log.Printf("--configs_tarball=%q \\", *configsTarball)
	log.Printf("--configs_manifest=%q \\", *configsManifest)
	log.Printf("--enable_monitoring=%v \\", *enableMonitoring)
	log.Printf("--monitoring_project_id=%q \\", *monitoringProjectID)
	log.Printf("--monitoring_docker_image=%q", *monitoringDockerImage)
}

// uploadConfigs is the core config upload logic allowing the caller a convenient wrapper to
// report results to monitoring before triggering a fatal exit.
// containerImage is the name of the toolchain container that will be used to name the directory
// on GCS configs are uploaded to.
func uploadConfigs(ctx context.Context, containerImage string) error {
	sc, err := newStorage(ctx)
	if err != nil {
		return fmt.Errorf("failed to initialize the GCS client: %v", err)
	}

	m, err := manifestFromFile(*configsManifest)
	if err != nil {
		return fmt.Errorf("error reading config manifest: %v", err)
	}
	manifestBlob, err := json.MarshalIndent(m, "", " ")
	if err != nil {
		return fmt.Errorf("error converting manifest into JSON: %v", err)
	}

	uploadDirs := []string{
		fmt.Sprintf("bazel-configs/%s/latest", containerImage),
		fmt.Sprintf("bazel-configs/bazel_%s/%s/latest", m.BazelVersion, containerImage),
	}
	for _, u := range uploadDirs {
		if err := sc.uploadArtifacts(ctx, manifestBlob, *configsTarball, u); err != nil {
			return fmt.Errorf("error uploading configs to GCS bucket %s, directory %s: %v", sc.bucketName, u, err)
		}
		log.Printf("Configs published to GCS bucket %s, directory %s.", sc.bucketName, u)
	}
	return nil
}

func initMonitoringClient(ctx context.Context) (*monitoring.Client, error) {
	if !(*enableMonitoring) {
		return nil, nil
	}
	if len(*monitoringProjectID) == 0 {
		return nil, fmt.Errorf("--monitoring_project_id is required because --enable_monitoring is true")
	}
	if len(*monitoringDockerImage) == 0 {
		return nil, fmt.Errorf("--monitoring_docker_image is required because --enable_monitoring is true")
	}
	c, err := monitoring.NewClient(ctx, *monitoringProjectID)
	if err != nil {
		return nil, fmt.Errorf("unable to initialize the monitoring client: %w", err)
	}
	return c, nil
}

func main() {
	flag.Parse()
	printFlags()

	if len(*configsTarball) == 0 {
		log.Fatalf("--configs_tarball was not specified.")
	}
	if len(*configsManifest) == 0 {
		log.Fatalf("--configs_manifest was not specified.")
	}

	ctx := context.Background()
	mc, err := initMonitoringClient(ctx)
	if err != nil {
		log.Fatalf("Failed to initialize monitoring: %v", err)
	}

	result := true
	if err := uploadConfigs(ctx, *monitoringDockerImage); err != nil {
		log.Printf("Configs upload failed: %v", err)
		result = false
	} else {
		log.Printf("Configs uploaded successfully.")
	}

	// Monitoring is optional and used for internal alerting by the owners of this repo only.
	if mc != nil {
		if err := mc.ReportToolchainConfigsUpload(ctx, *monitoringDockerImage, result); err != nil {
			log.Fatalf("Failed to report results to monitoring: %v", err)
		}
	}
	if !result {
		os.Exit(1)
	}
}
