// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package cmd

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	uri "path"
	"path/filepath"
	"sort"
	"strings"

	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/proto"
	"github.com/maruel/subcommands"
	tnProto "go.chromium.org/chromiumos/config/go/api/test/harness/tnull/v1"
	rtd "go.chromium.org/chromiumos/config/go/api/test/rtd/v1"
	"go.chromium.org/luci/common/cli"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
)

type gen struct {
	subcommands.CommandRunBase
	destBinaryProto string
	alsoOutputJson  bool
	testList        string
}

func GenerateFullRequest() *subcommands.Command {
	return &subcommands.Command{
		UsageLine: "generate-full-request -output_dir /path/to/output-dir/",
		ShortDesc: "create Invocations requesting each TNull test, as binaryproto files.",
		CommandRun: func() subcommands.CommandRun {
			g := &gen{}
			g.Flags.StringVar(&g.testList, "input_json", "",
				`Optional: path that contains a JSON list of test names to include in invocation.
				Each name should look like 'dummy-pass' or look
				like "remoteTestDrivers/tnull/tests/dummy-pass"`)
			g.Flags.StringVar(&g.destBinaryProto, "output", "", "File to contain binaryproto rtd.Invocation objects")
			g.Flags.BoolVar(&g.alsoOutputJson, "also_output_json", true, "Whether to create output JSON rtd.Invocation files in addition to the binaryprotos, in the same directory")
			return g
		},
	}
}

func (g *gen) Run(a subcommands.Application, args []string, env subcommands.Env) int {
	if err := g.innerRun(cli.GetContext(a, g, env)); err != nil {
		fmt.Fprintf(a.GetErr(), "%s\n", err)
		return 1
	}
	return 0
}

func (g *gen) innerRun(ctx context.Context) error {
	lookup, err := extractTestsFromSpecification()
	if err != nil {
		return err
	}
	names, err := getNames(g.testList)
	if err != nil {
		return err
	}
	nameStructs, err := parseNames(names)
	if err != nil {
		return err
	}
	validMap, err := filterTestMap(lookup.Lookup, nameStructs)
	if err != nil {
		return err
	}

	err = WriteInvocation(ctx, validMap, g.destBinaryProto, g.alsoOutputJson)
	if err != nil {
		return err
	}
	return nil
}

func getNames(path string) ([]string, error) {
	var names []string
	if path == "" {
		return names, nil
	}
	bs, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, errors.Annotate(err, "read in JSON list").Err()
	}
	err = json.Unmarshal(bs, &names)
	if err != nil {
		return nil, errors.Annotate(err, "read in JSON list").Err()
	}
	return names, nil
}

// WriteInvocation writes a binary-encoded Invocation proto to destFile.
func WriteInvocation(ctx context.Context, tests map[string]*tnProto.Steps, destFile string, alsoOutputJson bool) error {
	var reqs []*rtd.Request
	for _, name := range mapKeysToSortedList(tests) {
		req := rtd.Request{
			Name: "request_" + uri.Base(name),
			Test: name,
		}
		reqs = append(reqs, &req)
	}
	Inv := rtd.Invocation{
		Requests: reqs,
	}
	dir := filepath.Dir(destFile)
	// Create the directory if it doesn't exist.
	if err := os.MkdirAll(dir, 0755); err != nil {
		return errors.Annotate(err, "write invocation pb").Err()
	}

	{
		protoBytes, err := proto.Marshal(&Inv)
		if err != nil {
			return errors.Annotate(err, "write Invocation binaryproto").Err()
		}
		if err = ioutil.WriteFile(destFile, protoBytes, 0644); err != nil {
			return errors.Annotate(err, "write Invocation binaryproto").Err()
		}
		logging.Infof(ctx, "Wrote binaryproto Invocation to %v", destFile)
	}

	if alsoOutputJson {
		jsonFile := strings.TrimSuffix(destFile, ".binaryproto") + ".json"
		w, err := os.Create(jsonFile)
		if err != nil {
			return errors.Annotate(err, "write Invocation pb").Err()
		}
		defer w.Close()

		marshaler := jsonpb.Marshaler{Indent: "  "}
		if err := marshaler.Marshal(w, &Inv); err != nil {
			return errors.Annotate(err, "write JSON pb").Err()
		}
		logging.Infof(ctx, "Wrote JSON Invocation to %v", jsonFile)
	}
	return nil
}

func mapKeysToSortedList(m map[string]*tnProto.Steps) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	return keys
}

// We have two canonical name formats; short, e.g. "dummy-pass", and full, e.g.
// "remoteTestDrivers/tnull/tests/dummy-pass". Requiring either is a serious hassle
// for users, so convert to a single format. This accepts either and translates them
// to simple structs which keep the original value for the purposes of debugging
func parseNames(names []string) ([]NameStruct, error) {
	snames := make([]NameStruct, 0, len(names))
	errs := errors.MultiError{}
	for _, name := range names {
		split := strings.SplitN(name, "/", 4)
		if len(split) != 4 {
			snames = append(snames, NameStruct{givenName: name, shortName: name})
			continue
		}
		if uri.Join(split[0:2]...) != "remoteTestDrivers/tnull/tests" {
			errs = append(errs, errors.Reason("test name %s is malformed", name).Err())
			continue
		}
		snames = append(snames, NameStruct{givenName: name, shortName: split[3]})
	}
	if len(errs) > 0 {
		return nil, errs
	}
	return snames, nil
}

func filterTestMap(spec map[string]*tnProto.Steps, names []NameStruct) (map[string]*tnProto.Steps, error) {
	if len(names) == 0 {
		return spec, nil
	}
	errs := errors.MultiError{}
	newLookup := make(map[string]*tnProto.Steps)
	for _, ns := range names {
		name := ns.inferredName()
		if t, present := spec[name]; present {
			newLookup[name] = t
		} else {
			e := errors.Reason("Could not find test %s specified by name %s", name, ns.givenName).Err()
			errs = append(errs, e)
		}
	}
	if len(errs) > 0 {
		return nil, errs
	}
	return newLookup, nil
}

type NameStruct struct {
	givenName string
	shortName string
}

func (n *NameStruct) inferredName() string {
	return uri.Join(TNullName, "tests", n.shortName)
}