// Copyright 2024 The Go 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 modernize

import (
	"cmp"
	"fmt"
	"go/ast"
	"go/constant"
	"go/token"
	"go/types"
	"maps"
	"slices"

	"golang.org/x/tools/go/analysis"
	"golang.org/x/tools/go/analysis/passes/inspect"
	"golang.org/x/tools/go/ast/edge"
	"golang.org/x/tools/go/ast/inspector"
	"golang.org/x/tools/internal/analysis/analyzerutil"
	typeindexanalyzer "golang.org/x/tools/internal/analysis/typeindex"
	"golang.org/x/tools/internal/astutil"
	"golang.org/x/tools/internal/refactor"
	"golang.org/x/tools/internal/typesinternal/typeindex"
)

var StringsBuilderAnalyzer = &analysis.Analyzer{
	Name: "stringsbuilder",
	Doc:  analyzerutil.MustExtractDoc(doc, "stringsbuilder"),
	Requires: []*analysis.Analyzer{
		inspect.Analyzer,
		typeindexanalyzer.Analyzer,
	},
	Run: stringsbuilder,
	URL: "https://pkg.go.dev/golang.org/x/tools/go/analysis/passes/modernize#stringbuilder",
}

// stringsbuilder replaces string += string in a loop by strings.Builder.
func stringsbuilder(pass *analysis.Pass) (any, error) {
	// Skip the analyzer in packages where its
	// fixes would create an import cycle.
	if within(pass, "strings", "runtime") {
		return nil, nil
	}

	var (
		inspect = pass.ResultOf[inspect.Analyzer].(*inspector.Inspector)
		index   = pass.ResultOf[typeindexanalyzer.Analyzer].(*typeindex.Index)
		info    = pass.TypesInfo
	)

	// Gather all local string variables that appear on the
	// LHS of some string += string assignment.
	candidates := make(map[*types.Var]bool)
	for curAssign := range inspect.Root().Preorder((*ast.AssignStmt)(nil)) {
		assign := curAssign.Node().(*ast.AssignStmt)
		if assign.Tok == token.ADD_ASSIGN && is[*ast.Ident](assign.Lhs[0]) {
			if v, ok := info.Uses[assign.Lhs[0].(*ast.Ident)].(*types.Var); ok &&
				v.Kind() == types.LocalVar &&
				types.Identical(v.Type(), builtinString.Type()) {
				candidates[v] = true
			}
		}
	}

	lexicalOrder := func(x, y *types.Var) int { return cmp.Compare(x.Pos(), y.Pos()) }

	// File and Pos of last fix edit,
	// for overlapping fix span detection.
	var (
		lastEditFile *ast.File
		lastEditEnd  token.Pos
	)

	// Now check each candidate variable's decl and uses.
nextcand:
	for _, v := range slices.SortedFunc(maps.Keys(candidates), lexicalOrder) {
		var edits, postEdits []analysis.TextEdit // postEdits are emitted last

		// Check declaration of s has one of these forms:
		//
		//    s := expr
		//    var s [string] [= expr]
		//    var ( ...; s [string] [= expr] )			(s is last)
		//
		// and transform to one of:
		//
		//    var   s strings.Builder ; s.WriteString(expr)
		//    var ( s strings.Builder); s.WriteString(expr)
		//
		def, ok := index.Def(v)
		if !ok {
			continue
		}

		// To avoid semantic conflicts, do not offer a fix if its edit
		// range (ignoring import edits) overlaps a previous fix.
		// This fixes #76983 and is an ad-hoc mitigation of #76476.
		file := astutil.EnclosingFile(def)
		if file == lastEditFile && v.Pos() < lastEditEnd {
			continue
		}

		ek := def.ParentEdgeKind()
		if ek == edge.AssignStmt_Lhs &&
			len(def.Parent().Node().(*ast.AssignStmt).Lhs) == 1 {
			// Have: s := expr
			// => var s strings.Builder; s.WriteString(expr)

			assign := def.Parent().Node().(*ast.AssignStmt)

			// Reject "if s := f(); ..." since in that context
			// we can't replace the assign with two statements.
			switch def.Parent().Parent().Node().(type) {
			case *ast.BlockStmt, *ast.CaseClause, *ast.CommClause:
				// OK: these are the parts of syntax that
				// allow unrestricted statement lists.
			default:
				continue
			}

			// Add strings import.
			prefix, importEdits := refactor.AddImport(
				info, astutil.EnclosingFile(def), "strings", "strings", "Builder", v.Pos())
			edits = append(edits, importEdits...)

			if isEmptyString(info, assign.Rhs[0]) {
				// s := ""
				// ---------------------
				// var s strings.Builder
				edits = append(edits, analysis.TextEdit{
					Pos:     assign.Pos(),
					End:     assign.End(),
					NewText: fmt.Appendf(nil, "var %[1]s %[2]sBuilder", v.Name(), prefix),
				})

			} else {
				// s :=                                 expr
				// -------------------------------------    -
				// var s strings.Builder; s.WriteString(expr)
				edits = append(edits, []analysis.TextEdit{
					{
						Pos:     assign.Pos(),
						End:     assign.Rhs[0].Pos(),
						NewText: fmt.Appendf(nil, "var %[1]s %[2]sBuilder; %[1]s.WriteString(", v.Name(), prefix),
					},
					{
						Pos:     assign.End(),
						End:     assign.End(),
						NewText: []byte(")"),
					},
				}...)

			}

		} else if ek == edge.ValueSpec_Names &&
			len(def.Parent().Node().(*ast.ValueSpec).Names) == 1 &&
			first(def.Parent().Parent().LastChild()) == def.Parent() {
			// Have: var   s [string] [= expr]
			//   or: var ( s [string] [= expr] )
			// => var s strings.Builder; s.WriteString(expr)
			//
			// The LastChild check rejects this case:
			//   var ( s [string] [= expr]; others... )
			// =>
			//   var ( s strings.Builder; others... ); s.WriteString(expr)
			// since it moves 'expr' across 'others', requiring
			// reformatting of syntax, which in general is lossy
			// of comments and vertical space.
			// We expect this to be rare.

			// Add strings import.
			prefix, importEdits := refactor.AddImport(
				info, astutil.EnclosingFile(def), "strings", "strings", "Builder", v.Pos())
			edits = append(edits, importEdits...)

			spec := def.Parent().Node().(*ast.ValueSpec)
			decl := def.Parent().Parent().Node().(*ast.GenDecl)

			init := spec.Names[0].End() // start of " = expr"
			if spec.Type != nil {
				init = spec.Type.End()
			}

			// Replace (possibly absent) type:
			//
			// var s [string]
			//      ----------------
			// var s strings.Builder
			edits = append(edits, analysis.TextEdit{
				Pos:     spec.Names[0].End(),
				End:     init,
				NewText: fmt.Appendf(nil, " %sBuilder", prefix),
			})

			if len(spec.Values) > 0 && !isEmptyString(info, spec.Values[0]) {
				if decl.Rparen.IsValid() {
					// var decl with explicit parens:
					//
					// var ( ...  =               expr )
					//           -                     -
					// var ( ... ); s.WriteString(expr)
					edits = append(edits, []analysis.TextEdit{
						{
							Pos:     init,
							End:     init,
							NewText: []byte(")"),
						},
						{
							Pos: spec.Values[0].End(),
							End: decl.End(),
						},
					}...)
				}

				// =               expr
				// ----------------    -
				// ; s.WriteString(expr)
				edits = append(edits, []analysis.TextEdit{
					{
						Pos:     init,
						End:     spec.Values[0].Pos(),
						NewText: fmt.Appendf(nil, "; %s.WriteString(", v.Name()),
					},
					{
						Pos:     spec.Values[0].End(),
						End:     spec.Values[0].End(),
						NewText: []byte(")"),
					},
				}...)
			} else {
				// delete "= expr"
				edits = append(edits, analysis.TextEdit{
					Pos: init,
					End: spec.End(),
				})
			}

		} else {
			continue
		}

		// Check uses of s.
		//
		// - All uses of s except the final one must be of the form
		//
		//    s += expr
		//
		//   Each of these will become s.WriteString(expr).
		//   At least one of them must be in an intervening loop
		//   w.r.t. the declaration of s:
		//
		//    var s string
		//    for ... { s += expr }
		//
		// - All uses of s after the last += must be rvalue uses (e.g. use(s), not &s).
		//   Each of these will become s.String().
		//
		//   Perhaps surprisingly, it is fine for there to be an
		//   intervening loop or lambda w.r.t. the declaration of s:
		//
		//    var s strings.Builder
		//    for range kSmall { s.WriteString(expr) }
		//    for range kLarge { use(s.String()) } // called repeatedly
		//
		//   Even though that might cause the s.String() operation to be
		//   executed repeatedly, this is not a deoptimization because,
		//   by design, (*strings.Builder).String does not allocate.
		var (
			numLoopAssigns int             // number of += assignments within a loop
			loopAssign     *ast.AssignStmt // first += assignment within a loop
			seenRvalueUse  bool            // => we've seen at least one rvalue use of s
		)
		for curUse := range index.Uses(v) {
			// Strip enclosing parens around Ident.
			ek := curUse.ParentEdgeKind()
			for ek == edge.ParenExpr_X {
				curUse = curUse.Parent()
				ek = curUse.ParentEdgeKind()
			}

			// intervening reports whether cur has an ancestor of
			// one of the given types that is within the scope of v.
			intervening := func(types ...ast.Node) bool {
				for cur := range curUse.Enclosing(types...) {
					if v.Pos() <= cur.Node().Pos() { // in scope of v
						return true
					}
				}
				return false
			}

			if ek == edge.AssignStmt_Lhs {
				// After an rvalue use, no more assignments are allowed.
				if seenRvalueUse {
					continue nextcand
				}

				assign := curUse.Parent().Node().(*ast.AssignStmt)
				if assign.Tok != token.ADD_ASSIGN {
					continue nextcand
				}
				// Have: s += expr

				// At least one of the += operations
				// must appear within a loop.
				// relative to the declaration of s.
				if intervening((*ast.ForStmt)(nil), (*ast.RangeStmt)(nil)) {
					numLoopAssigns++
					if loopAssign == nil {
						loopAssign = assign
					}
				}

				// s +=          expr
				//  -------------    -
				// s.WriteString(expr)
				edits = append(edits, analysis.TextEdit{
					// replace " += " with ".WriteString("
					Pos:     assign.Lhs[0].End(),
					End:     assign.Rhs[0].Pos(),
					NewText: []byte(".WriteString("),
				})

				// Delay inserting the closing parenthesis, in case it overlaps with a
				// .String() edit, since it would need to come after.
				postEdits = append(postEdits, analysis.TextEdit{
					// insert ")"
					Pos:     assign.End(),
					End:     assign.End(),
					NewText: []byte(")"),
				})

			} else if ek == edge.UnaryExpr_X &&
				curUse.Parent().Node().(*ast.UnaryExpr).Op == token.AND {
				// Have: use(&s)
				continue nextcand // s is used as an lvalue; reject

			} else {
				// The only possible l-value uses of a string variable
				// are assignments (s=expr, s+=expr, etc) and &s.
				// (For strings, we can ignore method calls s.m().)
				// All other uses are r-values.
				seenRvalueUse = true

				edits = append(edits, analysis.TextEdit{
					// insert ".String()"
					Pos:     curUse.Node().End(),
					End:     curUse.Node().End(),
					NewText: []byte(".String()"),
				})
			}
		}
		if !seenRvalueUse {
			continue nextcand // no rvalue use; reject
		}
		if numLoopAssigns == 0 {
			continue nextcand // no += in a loop; reject
		}

		edits = append(edits, postEdits...)

		lastEditFile = file
		lastEditEnd = edits[len(edits)-1].End

		pass.Report(analysis.Diagnostic{
			Pos:     loopAssign.Pos(),
			End:     loopAssign.End(),
			Message: "using string += string in a loop is inefficient",
			SuggestedFixes: []analysis.SuggestedFix{{
				Message:   "Replace string += string with strings.Builder",
				TextEdits: edits,
			}},
		})
	}

	return nil, nil
}

// isEmptyString reports whether e (a string-typed expression) has constant value "".
func isEmptyString(info *types.Info, e ast.Expr) bool {
	tv, ok := info.Types[e]
	return ok && tv.Value != nil && constant.StringVal(tv.Value) == ""
}