package main

import (
	"bytes"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"io/fs"
	"path/filepath"
	"strings"
	"testing"
)

// String fields cannot encode the region vocabulary in their Go types. This
// lint checks scalar comparisons, call arguments and map keys at recognizable
// record/engine boundaries. Runtime witnesses cover local variable aliases.
func plan13Expr(e ast.Node) string {
	var b bytes.Buffer
	_ = printer.Fprint(&b, token.NewFileSet(), e)
	return b.String()
}
func plan13Placement(s string) bool {
	if strings.Contains(s, "PlacementRegion(") {
		return true
	}
	return s == "to" || s == "from" || s == "rec.Region" || s == "recRegion" || s == "mv.To" || s == "mv.From" || s == "m.To" || s == "m.From" || strings.HasSuffix(s, ".Move.To") || strings.HasSuffix(s, ".Move.From")
}
func plan13Engine(s string) bool {
	if strings.Contains(s, ".Intent[") && (strings.HasSuffix(s, `["to"]`) || strings.HasSuffix(s, `["region"]`)) {
		return true
	}
	if strings.HasSuffix(s, ".Region()") || strings.HasSuffix(s, ".Touched.Regions") || strings.HasSuffix(s, ".QuietAfter") {
		return true
	}
	if strings.HasPrefix(s, `strings.TrimPrefix(`) && strings.HasSuffix(s, `.Owner, "region:")`) {
		return plan13Engine(strings.TrimSuffix(strings.TrimPrefix(s, `strings.TrimPrefix(`), `.Owner, "region:")`) + ".Region")
	}
	for _, field := range []string{".Region", ".PredecessorRegion"} {
		if strings.HasSuffix(s, field) {
			switch strings.TrimSuffix(s, field) {
			case "ob", "o", "a", "predecessor", "entry", "record", "box", "row", "rr.row":
				return true
			}
		}
	}
	switch s {
	case "actor.Region", "src.Region", "dst.Region", "b.Region", "engine.Region", "e.Region", "r.Holder.Region", "in.SourceRegion", "in.TargetRegion", "s.Region", "a.Region", "o.Region", "m.Region", "t.Region", "fence.Region", "verified.Region", "regions":
		return true
	}
	return false
}
func plan13Operand(e ast.Expr, classify func(string) bool) bool {
	if _, ok := e.(*ast.FuncLit); ok {
		return false
	}
	if classify(plan13Expr(e)) {
		return true
	}
	switch x := e.(type) {
	case *ast.ParenExpr:
		return plan13Operand(x.X, classify)
	case *ast.CallExpr:
		switch plan13Expr(x.Fun) {
		case "park.SameRegion", "regionname.SameRegion", "SameRegion", "park.PlacementMatches", "PlacementMatches", "park.QuietUntil", "QuietUntil":
			return false
		}
		for _, arg := range x.Args {
			if plan13Operand(arg, classify) {
				return true
			}
		}
	}
	return false
}
func plan13MixedRegion(x *ast.BinaryExpr) bool { return plan13MixedRegionWith(x, plan13Engine) }
func plan13MixedRegionWith(x *ast.BinaryExpr, engine func(string) bool) bool {
	if x.Op != token.EQL && x.Op != token.NEQ {
		return false
	}
	return plan13Operand(x.X, plan13Placement) && plan13Operand(x.Y, engine) || plan13Operand(x.Y, plan13Placement) && plan13Operand(x.X, engine)
}
func plan13MixedRegionNode(n ast.Node) bool { return plan13MixedRegionNodeWith(n, plan13Engine) }
func plan13MixedRegionNodeWith(n ast.Node, engine func(string) bool) bool {
	switch x := n.(type) {
	case *ast.BinaryExpr:
		return plan13MixedRegionWith(x, engine)
	case *ast.AssignStmt:
		if len(x.Lhs) == len(x.Rhs) {
			for i := range x.Lhs {
				if plan13Operand(x.Lhs[i], engine) && plan13Operand(x.Rhs[i], plan13Placement) && !plan13EngineConversion(x.Rhs[i]) {
					return true
				}
			}
		}
	case *ast.IndexExpr:
		return plan13Operand(x.X, engine) && plan13Operand(x.Index, plan13Placement)
	case *ast.CallExpr:
		name := plan13Expr(x.Fun)
		switch name {
		case "park.SameRegion", "regionname.SameRegion", "SameRegion", "park.PlacementMatches", "PlacementMatches", "park.QuietUntil", "QuietUntil":
			return false
		}
		placement, hasEngine := false, false
		for _, arg := range x.Args {
			placement = placement || plan13Operand(arg, plan13Placement)
			hasEngine = hasEngine || plan13Operand(arg, engine)
		}
		return placement && hasEngine
	}
	return false
}

// Mapping a placement to an engine is valid at a store boundary, but cannot
// replace alias-aware equality between engine identities.
func plan13EngineConversion(e ast.Expr) bool {
	if p, ok := e.(*ast.ParenExpr); ok {
		return plan13EngineConversion(p.X)
	}
	c, ok := e.(*ast.CallExpr)
	if !ok {
		return false
	}
	switch plan13Expr(c.Fun) {
	case "park.PlacementRegion", "regionname.PlacementRegion", "PlacementRegion":
		return true
	}
	return false
}

// A scope maps scalar regions and records from engine-region collections.
func plan13Collection(e ast.Expr) string {
	s := plan13Expr(e)
	if strings.HasSuffix(s, ".Touched.Regions") {
		return "scalar"
	}
	if strings.HasSuffix(s, ".Admissions") || strings.HasSuffix(s, ".Obligations") {
		return "record"
	}
	return ""
}
func plan13Scope(scope map[string]string) map[string]string {
	next := map[string]string{}
	for name, kind := range scope {
		next[name] = kind
	}
	return next
}

type plan13Parent struct {
	switchTag  ast.Expr
	recordType string
}

func plan13Inspect(root ast.Node, scope map[string]string, reject func(ast.Node), parents ...plan13Parent) {
	var parent plan13Parent
	if len(parents) > 0 {
		parent = parents[0]
	}
	engine := func(s string) bool {
		if kind, ok := scope[s]; ok {
			return kind == "scalar"
		}
		for _, field := range []string{".Region", ".PredecessorRegion"} {
			if strings.HasSuffix(s, field) {
				if kind, ok := scope[strings.TrimSuffix(s, field)]; ok {
					return kind == "record"
				}
			}
		}
		return plan13Engine(s)
	}
	ast.Inspect(root, func(n ast.Node) bool {
		if n == nil {
			return false
		}
		if plan13MixedRegionNodeWith(n, engine) {
			reject(n)
		}
		switch x := n.(type) {
		case *ast.RangeStmt:
			plan13Inspect(x.X, scope, reject)
			next := plan13Scope(scope)
			if name, ok := x.Key.(*ast.Ident); ok {
				next[name.Name] = ""
			}
			if name, ok := x.Value.(*ast.Ident); ok {
				next[name.Name] = plan13Collection(x.X)
			}
			plan13Inspect(x.Body, next, reject)
			return false
		case *ast.CallExpr:
			for _, arg := range x.Args {
				if fn, ok := arg.(*ast.FuncLit); ok {
					next := plan13Scope(scope)
					for _, field := range fn.Type.Params.List {
						for _, name := range field.Names {
							next[name.Name] = ""
						}
					}
					if len(x.Args) > 0 && len(fn.Type.Params.List) > 0 {
						for _, name := range fn.Type.Params.List[0].Names {
							next[name.Name] = plan13Collection(x.Args[0])
						}
					}
					plan13Inspect(fn.Body, next, reject)
				} else {
					plan13Inspect(arg, scope, reject)
				}
			}
			plan13Inspect(x.Fun, scope, reject)
			return false
		case *ast.FuncLit:
			next := plan13Scope(scope)
			for _, field := range x.Type.Params.List {
				for _, name := range field.Names {
					next[name.Name] = ""
				}
			}
			plan13Inspect(x.Body, next, reject)
			return false
		case *ast.SwitchStmt:
			if x.Init != nil {
				plan13Inspect(x.Init, scope, reject)
			}
			if x.Tag != nil {
				plan13Inspect(x.Tag, scope, reject)
			}
			for _, stmt := range x.Body.List {
				plan13Inspect(stmt, scope, reject, plan13Parent{switchTag: x.Tag})
			}
			return false
		case *ast.CaseClause:
			for _, value := range x.List {
				if parent.switchTag != nil && plan13MixedRegionWith(&ast.BinaryExpr{X: parent.switchTag, Op: token.EQL, Y: value}, engine) {
					reject(x)
				}
			}
		case *ast.CompositeLit:
			typ := plan13Expr(x.Type)
			for _, element := range x.Elts {
				plan13Inspect(element, scope, reject, plan13Parent{recordType: typ[strings.LastIndex(typ, ".")+1:]})
			}
			return false
		case *ast.KeyValueExpr:
			switch parent.recordType {
			case "Obligation", "Admission", "AdmissionRecord", "Selection":
				if (plan13Expr(x.Key) == "Region" || plan13Expr(x.Key) == "PredecessorRegion") && plan13Operand(x.Value, plan13Placement) && !plan13EngineConversion(x.Value) {
					reject(x)
				}
			}
		}
		return true
	})
}

func TestPlan13NoDirectPlacementEngineComparison(t *testing.T) {
	err := filepath.WalkDir(".", func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if d.IsDir() {
			if path != "." && strings.HasPrefix(d.Name(), ".") {
				return filepath.SkipDir
			}
			return nil
		}
		if !strings.HasSuffix(path, ".go") || strings.HasSuffix(path, "_test.go") {
			return nil
		}
		fset := token.NewFileSet()
		f, err := parser.ParseFile(fset, path, nil, 0)
		if err != nil {
			return err
		}
		plan13Inspect(f, nil, func(n ast.Node) {
			t.Errorf("%s: placement and engine regions need the shared rule: %s", fset.Position(n.Pos()), plan13Expr(n))
		})
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
}

func TestPlan13RegionComparisonLintWitness(t *testing.T) {
	for _, expr := range []string{`rec.Region != src.Region`, `rec.Move.To == c.Op.Intent["to"]`, `in.SourceRegion == mv.From`, `park.PlacementRegion(rec.Region) != src.Region`, `rec.Move.To == s.Region`, `m.From == m.Region`, `rec.Region == fence.Region()`, `to == c.Op.Intent["to"]`, `canonical(rec.Region) == canonical(actor.Region)`} {
		e, err := parser.ParseExpr(expr)
		if err != nil {
			t.Fatal(err)
		}
		if !plan13MixedRegion(e.(*ast.BinaryExpr)) {
			t.Fatal("missed", expr)
		}
	}
	for _, expr := range []string{`rec.Region != cfg.Region`, `rec.Move.To == rec.Region`, `receipt.Region != box.Region`, `canonical(rec.Region) != canonical(cfg.Region)`} {
		e, err := parser.ParseExpr(expr)
		if err != nil {
			t.Fatal(err)
		}
		if plan13MixedRegion(e.(*ast.BinaryExpr)) {
			t.Fatal("false positive", expr)
		}
	}
}

func TestPlan13RegionCallAndMapLintWitness(t *testing.T) {
	for _, expr := range []string{
		`containsString(c.Touched.Regions, rec.Move.To)`,
		`slices.Contains(c.Touched.Regions, rec.Move.To)`,
		`strings.EqualFold(rec.Region, actor.Region)`,
		`same(rec.Move.From, verified.Region())`,
		`same(canonical(rec.Move.From), canonical(verified.Region()))`,
		`c.QuietAfter[rec.Move.To]`,
		`c.QuietAfter[canonical(rec.Move.To)]`,
		`regions[rec.Region]`,
		`slices.Contains(c.Touched.Regions, park.PlacementRegion(rec.Move.To))`,
		`c.QuietAfter[park.PlacementRegion(rec.Move.To)]`,
	} {
		e, err := parser.ParseExpr(expr)
		if err != nil {
			t.Fatal(err)
		}
		if !plan13MixedRegionNode(e) {
			t.Fatal("missed", expr)
		}
	}
	for _, expr := range []string{
		`containsString(c.Touched.Boxes, box)`,
		`strings.EqualFold(rec.Region, cfg.Region)`,
		`park.PlacementMatches(home, rec.Region, actor.Region)`,
		`park.SameRegion(home, park.PlacementRegion(rec.Region), actor.Region)`,
		`park.QuietUntil(c.QuietAfter, home, park.PlacementRegion(rec.Move.To))`,
		`fleet.targets[rec.Region]`,
	} {
		e, err := parser.ParseExpr(expr)
		if err != nil {
			t.Fatal(err)
		}
		if plan13MixedRegionNode(e) {
			t.Fatal("false positive", expr)
		}
	}
}

func TestPlan13RegionFlowLintWitness(t *testing.T) {
	for _, body := range []string{
		`if park.PlacementRegion(rec.Region) == box.Region {}`,
		`if park.PlacementRegion(rec.Region) != row.Region {}`,
		`if park.PlacementRegion(rec.Region) == rr.row.Region {}`,
		`if ob.Region == park.PlacementRegion(rec.Region) {}`,
		`if predecessor.Region == park.PlacementRegion(rec.Region) {}`,
		`if entry.PredecessorRegion == park.PlacementRegion(rec.Region) {}`,
		`if strings.TrimPrefix(ob.Owner, "region:") == park.PlacementRegion(rec.Region) {}`,
		`slices.ContainsFunc(c.Touched.Regions, func(r string) bool { return r == rec.Move.To })`,
		`for _, r := range c.Touched.Regions { if r == park.PlacementRegion(rec.Move.To) {} }`,
		`switch park.PlacementRegion(rec.Region) { case s.Region: }`,
		`_ = hostclaim.Obligation{Region: rec.Move.To}`,
		`s.Region = rec.Region`,
		`if park.PlacementRegion(rec.Region) == record.Region {}`,
		`for _, item := range c.Admissions { if item.PredecessorRegion == rec.Move.From {} }`,
		`for _, item := range c.Obligations { if item.Region == rec.Region {} }`,
		`slices.ContainsFunc(c.Admissions, func(item hostclaim.AdmissionRecord) bool { return item.Region == rec.Region })`,
		`_ = hostclaim.Admission{Region: rec.Region}`,
		`_ = hostbox.Selection{Region: rec.Region}`,
		`s.Region, other = rec.Region, x`,
		`f := func() { s.Region = rec.Region }; f()`,
	} {
		t.Run(body, func(t *testing.T) {
			f, err := parser.ParseFile(token.NewFileSet(), "witness.go", "package p; func f(){"+body+"}", 0)
			if err != nil {
				t.Fatal(err)
			}
			found := false
			plan13Inspect(f, nil, func(ast.Node) { found = true })
			if !found {
				t.Fatal("missed", body)
			}
		})
	}
}

func TestPlan13RegionFlowLintAcceptsBoundaries(t *testing.T) {
	for _, body := range []string{
		`slices.ContainsFunc(c.Touched.Regions,func(r string) bool { return park.PlacementMatches(home, rec.Move.To,r) })`,
		`for _,r:=range c.Touched.Regions { if park.SameRegion(home,r,park.PlacementRegion(rec.Region)) {} }`,
		`for _,item:=range c.Admissions { if park.PlacementMatches(home,rec.Region,item.Region) {} }`,
		`s.Region = park.PlacementRegion(rec.Region)`,
		`_ = hostclaim.Obligation{Region: park.PlacementRegion(rec.Region)}`,
		`_ = provision.Record{Region: rec.Region}`,
		`switch rec.Region { case cfg.Region: }`,
		`for _,r:=range c.Touched.Regions { use(r) }; if r == rec.Region {}`,
		`slices.ContainsFunc(c.Touched.Regions,func(r string) bool { return true }); if r == rec.Region {}`,
		`slices.ContainsFunc(c.Touched.Regions,func(r string) bool { return func(r string) bool { return r == rec.Region }("placement") })`,
	} {
		t.Run(body, func(t *testing.T) {
			f, err := parser.ParseFile(token.NewFileSet(), "witness.go", "package p; func f(){"+body+"}", 0)
			if err != nil {
				t.Fatal(err)
			}
			plan13Inspect(f, nil, func(n ast.Node) { t.Errorf("false positive: %s", plan13Expr(n)) })
		})
	}
}