// Copyright 2025 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 typeindex_test

import (
	"go/ast"
	"go/parser"
	"go/token"
	"go/types"
	"slices"
	"testing"

	"golang.org/x/tools/go/ast/inspector"
	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/go/types/typeutil"
	"golang.org/x/tools/internal/testenv"
	"golang.org/x/tools/internal/testfiles"
	"golang.org/x/tools/internal/typesinternal/typeindex"
	"golang.org/x/tools/txtar"
)

func TestIndex(t *testing.T) {
	testenv.NeedsGoPackages(t)
	var (
		pkg        = loadNetHTTP(t)
		inspect    = inspector.New(pkg.Syntax)
		index      = typeindex.New(inspect, pkg.Types, pkg.TypesInfo)
		fmtSprintf = index.Object("fmt", "Sprintf")
	)

	// Gather calls and uses of fmt.Sprintf in net/http.
	var (
		wantUses  []*ast.Ident
		wantCalls []*ast.CallExpr
	)
	for n := range inspect.PreorderSeq((*ast.CallExpr)(nil), (*ast.Ident)(nil)) {
		switch n := n.(type) {
		case *ast.CallExpr:
			if typeutil.Callee(pkg.TypesInfo, n) == fmtSprintf {
				wantCalls = append(wantCalls, n)
			}
		case *ast.Ident:
			if pkg.TypesInfo.Uses[n] == fmtSprintf {
				wantUses = append(wantUses, n)
			}
		}
	}
	// sanity check (expect about 60 of each)
	if wantUses == nil || wantCalls == nil {
		t.Fatalf("no calls or uses of fmt.Sprintf in net/http")
	}

	var (
		gotUses  []*ast.Ident
		gotCalls []*ast.CallExpr
	)
	for curId := range index.Uses(fmtSprintf) {
		gotUses = append(gotUses, curId.Node().(*ast.Ident))
	}
	for curCall := range index.Calls(fmtSprintf) {
		gotCalls = append(gotCalls, curCall.Node().(*ast.CallExpr))
	}

	if !slices.Equal(gotUses, wantUses) {
		t.Errorf("index.Uses(fmt.Sprintf) = %v, want %v", gotUses, wantUses)
	}
	if !slices.Equal(gotCalls, wantCalls) {
		t.Errorf("index.Calls(fmt.Sprintf) = %v, want %v", gotCalls, wantCalls)
	}
}

func loadNetHTTP(tb testing.TB) *packages.Package {
	cfg := &packages.Config{Mode: packages.LoadSyntax}
	pkgs, err := packages.Load(cfg, "net/http")
	if err != nil {
		tb.Fatal(err)
	}
	return pkgs[0]
}

func BenchmarkIndex(b *testing.B) {
	testenv.NeedsGoPackages(b)
	// Load net/http, a large package, and find calls to net.Dial.
	//
	// There is currently exactly one, which provides an extreme
	// demonstration of the performance advantage of the Index.
	//
	// Index construction costs approximately 7x the cursor
	// traversal, so it breaks even when it replaces 7 passes.
	// The cost of index lookup is approximately zero.
	pkg := loadNetHTTP(b)

	// Build the Inspector (~2.8ms).
	var inspect *inspector.Inspector
	b.Run("inspector.New", func(b *testing.B) {
		for b.Loop() {
			inspect = inspector.New(pkg.Syntax)
		}
	})

	// Build the Index (~6.6ms).
	var index *typeindex.Index
	b.Run("typeindex.New", func(b *testing.B) {
		b.ReportAllocs() // 2.48MB/op
		for b.Loop() {
			index = typeindex.New(inspect, pkg.Types, pkg.TypesInfo)
		}
	})

	target := index.Object("net", "Dial")

	var countA, countB, countC int

	// unoptimized inspect implementation (~1.6ms, 1x)
	b.Run("inspect", func(b *testing.B) {
		for b.Loop() {
			countA = 0
			for _, file := range pkg.Syntax {
				ast.Inspect(file, func(n ast.Node) bool {
					if call, ok := n.(*ast.CallExpr); ok {
						if typeutil.Callee(pkg.TypesInfo, call) == target {
							countA++
						}
					}
					return true
				})
			}
		}
	})
	if countA == 0 {
		b.Errorf("target %v not found", target)
	}

	// unoptimized cursor implementation (~390us, 4x faster)
	b.Run("cursor", func(b *testing.B) {
		for b.Loop() {
			countB = 0
			for curCall := range inspect.Root().Preorder((*ast.CallExpr)(nil)) {
				call := curCall.Node().(*ast.CallExpr)
				if typeutil.Callee(pkg.TypesInfo, call) == target {
					countB++
				}
			}
		}
	})

	// indexed implementation (~120ns, >10,000x faster)
	b.Run("index", func(b *testing.B) {
		for b.Loop() {
			countC = 0
			for range index.Calls(target) {
				countC++
			}
		}
	})

	if countA != countB || countA != countC {
		b.Fatalf("inconsistent results (inspect=%d, cursor=%d, index=%d)", countA, countB, countC)
	}
}

// TestOrigin checks that index.Uses on field or method of a generic
// type reports references to its instantiations.
// Lexical references (to functions, types, and aliases) always
// resolve to the generic; there is no Origin distinction.
func TestOrigin(t *testing.T) {
	// Each nonblank declared name must have at least one use.
	for _, src := range []string{
		`
package alias
type A[T any] = *T
var _ A[int]`,
		`
package localalias
func _[T any]() {
	type A[U any] = *U
	var _ A[T]
}`,
		`
package named
type Named[_ any] int
var _ Named[int]`,
		`
package localnamed
func _[T any]() {
	type N[U any] struct { f U }
	_ = N[T]{}.f
}`,
		`
package field
type Named[T any] struct { field T }
var _ = Named[int]{}.field`,
		`
package function
func F[T any](T) {}
var _ = F[int]`,
		`
package method
type Named[_ any] int
func (Named[T]) method(T) {}
var _ = Named[int].method`,
		`
package uninstantiatedmethod
type Named[_ any] int
func (Named[_]) f() {}
func (Named[T]) g() { Named[T](0).f() }
func (n Named[_]) _() { n.g() }`,
	} {
		// parse
		fset := token.NewFileSet()
		f, _ := parser.ParseFile(fset, "", src, parser.PackageClauseOnly)
		t.Run(f.Name.Name, func(t *testing.T) {
			f, _ := parser.ParseFile(fset, "", src, 0)
			files := []*ast.File{f}

			// typecheck
			info := &types.Info{
				Defs: make(map[*ast.Ident]types.Object),
				Uses: make(map[*ast.Ident]types.Object),
			}
			pkg, err := new(types.Config).Check(f.Name.Name, fset, files, info)
			if err != nil {
				t.Fatal(err)
			}

			// index
			var (
				inspect = inspector.New(files)
				index   = typeindex.New(inspect, pkg, info)
			)

			// check Def/Uses of all objects
			for id := range inspect.PreorderSeq((*ast.Ident)(nil)) {
				id := id.(*ast.Ident)
				if obj := info.Defs[id]; obj != nil && id.Name != "_" {
					t.Logf("%T %s @ %s", obj, id.Name, fset.Position(id.Pos()))

					// check Def
					if curDef, ok := index.Def(obj); !ok {
						t.Errorf("- index.Def missing")
					} else if curDef.Node() != id {
						t.Errorf("- index.Def wrong: @%s", fset.Position(curDef.Node().Pos()))
					}

					// check Uses is nonempty
					ok := false
					for use := range index.Uses(obj) {
						t.Logf("- index.Use at %s", fset.Position(use.Node().Pos()))
						ok = true
					}
					if !ok {
						t.Errorf("- no index.Uses")
					}
				}
			}
		})
	}
}

func TestCalls(t *testing.T) {
	testenv.NeedsGoPackages(t)
	const file = `-- go.mod --
module foo

go 1.24

-- pkg/pkg.go --
package pkg

type Bar struct {}

func (f Bar) Method() {}

var Foo Bar
-- test.go --
package foo

import "foo/pkg"

func foo() {}

func _() {
	foo()
	(foo)()
	((foo))()
}

func fooGeneric[T any]() {}

func _() {
	fooGeneric[int]()
	(fooGeneric[int])()
	((fooGeneric[int]))()
}

func fooGeneric2[T, T2 any]() {}

func _() {
	fooGeneric2[int, int]()
	(fooGeneric2[int, int])()
	((fooGeneric2[int, int]))()
}

func _() {
	var foo pkg.Bar
	(foo.Method)()
	((foo).Method)()
	(foo).Method()
	pkg.Bar.Method(pkg.Bar{})
	(pkg.Bar).Method(pkg.Bar{})
	((pkg.Bar).Method)(pkg.Bar{})

	var bar struct { field pkg.Bar }
	bar.field.Method()
}

func _() {
	pkg.Foo.Method()
	(pkg.Foo).Method()
	((pkg.Foo).Method)()
	(pkg.Foo.Method)()
}
`
	fs, err := txtar.FS(txtar.Parse([]byte(file)))
	if err != nil {
		t.Fatal(err)
	}
	dir := testfiles.CopyToTmp(t, fs)

	fset := token.NewFileSet()
	cfg := packages.Config{Mode: packages.LoadAllSyntax, Fset: fset, Dir: dir}
	pkgs, err := packages.Load(&cfg, "foo")
	if err != nil {
		t.Fatal(err)
	}
	if packages.PrintErrors(pkgs) > 0 {
		t.Fatal("Load: there were errors")
	}

	var (
		pkg     = pkgs[0]
		inspect = inspector.New(pkg.Syntax)
		index   = typeindex.New(inspect, pkg.Types, pkg.TypesInfo)
	)

	// Emulate index.Calls with a tree walk.
	objToCall := make(map[types.Object][]token.Pos)
	for call := range inspector.All[*ast.CallExpr](inspect) {
		obj := typeutil.Callee(pkg.TypesInfo, call)
		if obj != nil {
			objToCall[obj] = append(objToCall[obj], call.Pos())
		}
	}

	// Assert index.Calls returns the same results, as the slow version (above).
	for obj, wantCalls := range objToCall {
		var gotCalls []token.Pos
		for call := range index.Calls(obj) {
			gotCalls = append(gotCalls, call.Node().Pos())
		}
		if !slices.Equal(gotCalls, wantCalls) {
			t.Errorf("%v: got calls: %v; want: %v", obj.Name(), gotCalls, wantCalls)
		}
	}
}
