package cq

import (
	"context"
	"io"
	"io/ioutil"
	osexec "os/exec"
	"path/filepath"
	"regexp"

	"github.com/bazelbuild/buildtools/build"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/util"
)

// WithUpdateCQConfig reads the given Starlark config file, runs the given
// function to modify it, then writes it back to disk and runs lucicfg to
// generate the proto config files in the given directory. Expects lucicfg to
// be in PATH and "lucicfg auth-login" to have already been run. generatedDir
// must be a descendant of the directory which contains starlarkConfigFile.
func WithUpdateCQConfig(ctx context.Context, starlarkConfigFile, generatedDir string, fn func(*build.File) error) error {
	// Make the generatedDir relative to the directory containing
	// starlarkConfigFile.
	parentDir := filepath.Dir(starlarkConfigFile)
	relConfigFile := filepath.Base(starlarkConfigFile)
	relGenDir, err := filepath.Rel(parentDir, generatedDir)
	if err != nil {
		return skerr.Wrapf(err, "could not make %s relative to %s", generatedDir, parentDir)
	}

	// Check presence and auth status of lucicfg.
	if lucicfg, err := osexec.LookPath("lucicfg"); err != nil || lucicfg == "" {
		return skerr.Fmt("unable to find lucicfg in PATH; do you have depot tools installed?")
	}
	if _, err := exec.RunCwd(ctx, ".", "lucicfg", "auth-info"); err != nil {
		return skerr.Wrapf(err, "please run `lucicfg auth-login`")
	}

	// Read the config file.
	oldCfgBytes, err := ioutil.ReadFile(starlarkConfigFile)
	if err != nil {
		return skerr.Wrapf(err, "failed to read %s", starlarkConfigFile)
	}

	// Update the config.
	newCfgBytes, err := WithUpdateCQConfigBytes(starlarkConfigFile, oldCfgBytes, fn)
	if err != nil {
		return skerr.Wrap(err)
	}

	// Write the new config.
	if err := util.WithWriteFile(starlarkConfigFile, func(w io.Writer) error {
		_, err := w.Write(newCfgBytes)
		return err
	}); err != nil {
		return skerr.Wrapf(err, "failed to write config file")
	}

	// Run lucicfg to generate the proto config files.
	cmd := []string{"lucicfg", "generate", "-validate", "-config-dir", relGenDir, relConfigFile}
	if _, err := exec.RunCwd(ctx, parentDir, cmd...); err != nil {
		return skerr.Wrap(err)
	}
	return nil
}

// WithUpdateCQConfigBytes parses the given bytes as a Config, runs the given
// function to modify the Config, then returns the updated bytes.
func WithUpdateCQConfigBytes(filename string, oldCfgBytes []byte, fn func(*build.File) error) ([]byte, error) {
	// Parse the Config.
	f, err := build.ParseDefault(filename, oldCfgBytes)
	if err != nil {
		return nil, skerr.Wrapf(err, "failed to parse config file")
	}

	// Run the passed-in func.
	if err := fn(f); err != nil {
		return nil, skerr.Wrapf(err, "config modification failed")
	}

	// Write the new config bytes.
	return build.Format(f), nil
}

// FindAssignExpr finds the AssignExpr with the given identifier.
func FindAssignExpr(callExpr *build.CallExpr, identifier string) (int, *build.AssignExpr, error) {
	for idx, expr := range callExpr.List {
		if assignExpr, ok := expr.(*build.AssignExpr); ok {
			if ident, ok := assignExpr.LHS.(*build.Ident); ok {
				if ident.Name == identifier {
					return idx, assignExpr, nil
				}
			}
		}
	}
	return 0, nil, skerr.Fmt("no AssignExpr found for %q", identifier)
}

// FindExprForBranch finds the CallExpr for the given branch.
func FindExprForBranch(f *build.File, branch string) (int, *build.CallExpr, error) {
	for idx, expr := range f.Stmt {
		if callExpr, ok := expr.(*build.CallExpr); ok {
			_, assignExpr, err := FindAssignExpr(callExpr, "name")
			if err != nil {
				continue
			}
			if stringExpr, ok := assignExpr.RHS.(*build.StringExpr); ok {
				if stringExpr.Value == branch {
					return idx, callExpr, nil
				}
			}
		}
	}
	return 0, nil, skerr.Fmt("no config group found for %q", branch)
}

// CopyExprSlice returns a shallow copy of the []Expr.
func CopyExprSlice(slice []build.Expr) []build.Expr {
	cp := make([]build.Expr, 0, len(slice))
	for _, expr := range slice {
		cp = append(cp, expr)
	}
	return cp
}

// CloneBranch updates the given CQ config to create a config for a new
// branch based on a given existing branch. Optionally, include experimental
// tryjobs, include the tree-is-open check, and exclude trybots matching regular
// expressions.
func CloneBranch(f *build.File, oldBranch, newBranch string, includeExperimental, includeTreeCheck bool, excludeTrybotRegexp []*regexp.Regexp) error {
	// Find the CQ config for the old branch.
	_, oldBranchExpr, err := FindExprForBranch(f, oldBranch)
	if err != nil {
		return err
	}

	// Copy the old branch and modify it for the new branch.
	newBranchExpr, ok := oldBranchExpr.Copy().(*build.CallExpr)
	if !ok {
		return skerr.Fmt("expected CallExpr for cq_group")
	}
	newBranchExpr.List = CopyExprSlice(oldBranchExpr.List)

	// CQ group name.
	nameExprIndex, nameExpr, err := FindAssignExpr(newBranchExpr, "name")
	if err != nil {
		return skerr.Wrap(err)
	}
	nameExprCopy := nameExpr.Copy().(*build.AssignExpr)
	nameStr, ok := nameExprCopy.RHS.Copy().(*build.StringExpr)
	if !ok {
		return skerr.Fmt("expected StringExpr for name")
	}
	nameStr.Value = newBranch
	nameExprCopy.RHS = nameStr
	newBranchExpr.List[nameExprIndex] = nameExprCopy

	// Ref match.
	refsetExprIndex, refsetExpr, err := FindAssignExpr(newBranchExpr, "watch")
	if err != nil {
		return skerr.Wrap(err)
	}
	refsetExprCopy := refsetExpr.Copy().(*build.AssignExpr)
	newBranchExpr.List[refsetExprIndex] = refsetExprCopy

	refsetCallExpr, ok := refsetExprCopy.RHS.(*build.CallExpr)
	if !ok {
		return skerr.Fmt("expected CallExpr for refset")
	}
	refsetCallExprCopy := refsetCallExpr.Copy().(*build.CallExpr)
	refsetExprCopy.RHS = refsetCallExprCopy
	refsetCallExprCopy.List = CopyExprSlice(refsetCallExprCopy.List)

	refsExprIndex, refsExpr, err := FindAssignExpr(refsetCallExprCopy, "refs")
	if err != nil {
		return skerr.Wrap(err)
	}
	refsExprCopy := refsExpr.Copy().(*build.AssignExpr)
	refsetCallExprCopy.List[refsExprIndex] = refsExprCopy

	refsListExpr, ok := refsExprCopy.RHS.(*build.ListExpr)
	if !ok {
		return skerr.Fmt("expected ListExpr for refs")
	}
	refsListExprCopy := refsListExpr.Copy().(*build.ListExpr)
	refsExprCopy.RHS = refsListExprCopy

	if len(refsListExprCopy.List) != 1 {
		return skerr.Fmt("expected a single ref but got %d", len(refsListExprCopy.List))
	}
	refExpr, ok := refsListExprCopy.List[0].(*build.StringExpr)
	if !ok {
		return skerr.Fmt("expected StringExpr for ref")
	}
	refExprCopy := refExpr.Copy().(*build.StringExpr)
	refExprCopy.Value = "refs/heads/" + newBranch
	refsListExprCopy.List = []build.Expr{refExprCopy}

	// Tryjobs.
	verifiersExprIndex, verifiersExpr, err := FindAssignExpr(newBranchExpr, "verifiers")
	if err != nil {
		return skerr.Wrap(err)
	}
	verifiersExprCopy := verifiersExpr.Copy().(*build.AssignExpr)
	newBranchExpr.List[verifiersExprIndex] = verifiersExprCopy

	verifiersListExpr, ok := verifiersExprCopy.RHS.(*build.ListExpr)
	if !ok {
		return skerr.Fmt("expected ListExpr for verifiers")
	}
	verifiersListExprCopy := verifiersListExpr.Copy().(*build.ListExpr)
	verifiersExprCopy.RHS = verifiersListExprCopy

	verifiersListExprCopy.List = make([]build.Expr, 0, len(verifiersListExpr.List))
	for _, expr := range verifiersListExpr.List {
		verifierCallExpr, ok := expr.(*build.CallExpr)
		if !ok {
			return skerr.Fmt("expected CallExpr for verifier")
		}
		// Include experimental builders?
		_, _, err := FindAssignExpr(verifierCallExpr, "experiment_percentage")
		if err == nil && !includeExperimental {
			continue
		}

		// Is this builder excluded based on a regex?
		_, builder, err := FindAssignExpr(verifierCallExpr, "builder")
		if err != nil {
			return skerr.Wrap(err)
		}
		builderStringExpr, ok := builder.RHS.(*build.StringExpr)
		if !ok {
			return skerr.Fmt("expected StringExpr for builder name")
		}
		include := true
		for _, regex := range excludeTrybotRegexp {
			if regex.MatchString(builderStringExpr.Value) {
				include = false
				break
			}
		}
		if include {
			// No need to copy, since we're not modifying the verifier
			// expression itself.
			verifiersListExprCopy.List = append(verifiersListExprCopy.List, verifierCallExpr)
		}
	}

	// Tree status.
	if !includeTreeCheck {
		treeCheckIndex, _, err := FindAssignExpr(newBranchExpr, "tree_status_host")
		if err == nil {
			cp := make([]build.Expr, 0, len(newBranchExpr.List))
			for idx, expr := range newBranchExpr.List {
				if idx != treeCheckIndex {
					cp = append(cp, expr)
				}
			}
			newBranchExpr.List = cp
		}
	}

	// Add the new branch config.
	f.Stmt = append(f.Stmt, newBranchExpr)
	return nil
}

// DeleteBranch updates the given CQ config to delete the config matching the
// given branch.
func DeleteBranch(f *build.File, branch string) error {
	branchExprIndex, _, err := FindExprForBranch(f, branch)
	if err != nil {
		return skerr.Wrap(err)
	}
	newStmt := make([]build.Expr, 0, len(f.Stmt)-1)
	for idx, expr := range f.Stmt {
		if idx != branchExprIndex {
			newStmt = append(newStmt, expr)
		}
	}
	f.Stmt = newStmt
	return nil
}