// 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 genparams

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"

	"github.com/dave/dst"
	"github.com/dave/dst/decorator"
	"github.com/dave/dst/dstutil"

	"chromiumos/tast/caller"
)

// envName is the name of the environment variable that instructs Ensure to
// update source files with generated metadata instead of just comparing them.
const envName = "TAST_GENERATE_UPDATE"

// Ensure ensures that parameterized test parameters in a test metadata are
// up-to-date.
//
// file is a file name of a .go file containing a Tast test definition (i.e.
// testing.AddTest call). params is a Go literal expression as a string that
// represents test parameters.
//
// If TAST_GENERATE_UPDATE environment variable is not set, this function checks
// if parameterized test parameters in file matches with params. If any mismatch
// is found, it reports a test error with a message prompting users to
// regenerate parameterized test parameters.
//
// If TAST_GENERATE_UPDATE environment variable is set, this function rewrites
// the testing.Test literal in file with params.
func Ensure(t TestingT, file, params string) {
	t.Helper()

	oldCode, err := ioutil.ReadFile(file)
	if err != nil {
		t.Fatalf("%s: %v", file, err)
	}

	// Construct DST of the source code.
	root, err := decorator.Parse(oldCode)
	if err != nil {
		t.Fatalf("%s: %v", file, err)
	}

	// Ensure that chromiumos/tast/testing is imported without alias.
	for _, im := range root.Imports {
		path, err := strconv.Unquote(im.Path.Value)
		if err != nil {
			continue
		}
		if path == "chromiumos/tast/testing" {
			if im.Name != nil {
				t.Fatalf("chromiumos/tast/testing must be imported without alias")
			}
		}
	}

	// Construct DST of the new metadata.
	_, testPath, _, _ := runtime.Caller(1)
	paramsCode := fmt.Sprintf(`package main
var _ = []testing.Param{
// Parameters generated by %s. DO NOT EDIT.
%s
}`, filepath.Base(testPath), params)
	paramsRoot, err := decorator.Parse(paramsCode)
	if err != nil {
		t.Errorf("%s: %v", file, err)
		return
	}
	paramsExpr := paramsRoot.Decls[0].(*dst.GenDecl).Specs[0].(*dst.ValueSpec).Values[0]

	// Replace Params in testing.Test literals.
	hits := 0
	dstutil.Apply(root, func(cur *dstutil.Cursor) bool {
		// Match with testing.Test composite literal.
		comp, ok := cur.Node().(*dst.CompositeLit)
		if !ok {
			return true
		}
		sel, ok := comp.Type.(*dst.SelectorExpr)
		if !ok || sel.Sel.Name != "Test" {
			return true
		}
		if id, ok := sel.X.(*dst.Ident); !ok || id.Name != "testing" {
			return true
		}

		// Delete Params if it exists.
		elts := append([]dst.Expr(nil), comp.Elts...)
		for i, elt := range elts {
			kv, ok := elt.(*dst.KeyValueExpr)
			if !ok {
				continue
			}
			if id, ok := kv.Key.(*dst.Ident); !ok || id.Name != "Params" {
				continue
			}
			elts = append(elts[:i], elts[i+1:]...)
			break
		}

		// Append Params to the composite literal.
		elts = append(elts, &dst.KeyValueExpr{
			Key:   &dst.Ident{Name: "Params"},
			Value: paramsExpr,
			Decs: dst.KeyValueExprDecorations{
				NodeDecs: dst.NodeDecs{
					After: dst.NewLine,
				},
			},
		})

		repl := &dst.CompositeLit{
			Type: &dst.SelectorExpr{
				X:   &dst.Ident{Name: "testing"},
				Sel: &dst.Ident{Name: "Test"},
			},
			Elts: elts,
		}
		cur.Replace(repl)
		hits++
		return false
	}, nil)

	if hits != 1 {
		t.Fatalf("%s: found %d testing.AddTest calls; want 1", file, hits)
	}

	var b bytes.Buffer
	if err := decorator.Fprint(&b, root); err != nil {
		t.Fatalf("%s: %v", file, err)
	}
	newCode := b.Bytes()

	// If the environment variable is set, update the source code.
	if os.Getenv(envName) != "" {
		if err := ioutil.WriteFile(file, newCode, 0666); err != nil {
			t.Fatalf("%s: failed to save the generated code: %v", file, err)
		}
		return
	}

	if !bytes.Equal(newCode, oldCode) {
		pkg := strings.Split(caller.Get(2), ".")[0]
		t.Errorf(`%s: Params is stale; run the following command to update:
%s=1 ~/trunk/src/platform/tast/tools/go.sh test -count=1 %s`, file, envName, pkg)
	}
}