package.go

package check

import (
	"fmt"
	"go/ast"
	"go/token"
	"go/types"
	"path/filepath"
	"text/template/parse"

	"golang.org/x/tools/go/packages"

	"github.com/typelate/check/internal/asteval"
	"github.com/typelate/check/internal/astgen"
)

type pendingCall struct {
	call         *ast.CallExpr
	receiverObj  types.Object
	templateName string
	dataType     types.Type
}

type resolvedTemplate struct {
	templates   asteval.Template
	functions   asteval.TemplateFunctions
	metadata    *asteval.TemplateMetadata
	definitions definitionSet

	// leftDelim and rightDelim are the delimiters the construction
	// expression left in effect. A later call on the same variable
	// continues from them, and nothing else records them: the template
	// keeps its delimiters to itself, so evaluating that later call
	// without these would parse it as though Delims had never been
	// called.
	leftDelim  string
	rightDelim string
}

// definitionsFor locates every template definition in the sources that
// were parsed while resolving a template variable, in parse order.
func definitionsFor(fset *token.FileSet, sources []asteval.ParsedText) []Definition {
	var ordered []Definition
	for _, src := range sources {
		if src.Literal != nil {
			ordered = append(ordered, definitionsInLiteral(fset, src.Literal, src.RootName, src.LeftDelim, src.RightDelim)...)
			continue
		}
		ordered = append(ordered, definitionsIn(src.RootName, src.Filename, src.Text, src.LeftDelim, src.RightDelim)...)
	}
	return ordered
}

// ExecuteTemplateNodeInspectorFunc is called for each ExecuteTemplate
// call that resolves to a known template.
//
// The tree is the invoked template's own tree, the same one def.Tree
// holds. The type is the data argument's type, and def locates where the
// template was defined.
type ExecuteTemplateNodeInspectorFunc func(node *ast.CallExpr, t *parse.Tree, tp types.Type, def Definition)

// Package discovers all .ExecuteTemplate calls in the given package,
// resolves receiver variables to their template construction chains,
// and type-checks each call.
//
// ExecuteTemplate must be called with a string literal for the second parameter.
//
// A non-nil result is a *Error tree; see Execute for how to walk it.
func Package(pkg *packages.Package, inspectCall ExecuteTemplateNodeInspectorFunc, inspectTemplate TemplateNodeInspectorFunc) error {
	pending, receivers := findExecuteCalls(pkg)
	resolved, resolveErrs := resolveTemplates(pkg, receivers)
	callErr := checkCalls(pkg, pending, resolved, inspectCall, inspectTemplate)
	return joinErrors(nil, nil, append(resolveErrs, callErr)...)
}

// findExecuteCalls walks the package syntax looking for ExecuteTemplate calls
// and returns the pending calls along with the set of receiver objects that
// need template resolution.
func findExecuteCalls(pkg *packages.Package) ([]pendingCall, map[types.Object]struct{}) {
	var pending []pendingCall
	receiverSet := make(map[types.Object]struct{})

	for _, file := range pkg.Syntax {
		ast.Inspect(file, func(node ast.Node) bool {
			call, ok := node.(*ast.CallExpr)
			if !ok || len(call.Args) != 3 {
				return true
			}
			sel, ok := call.Fun.(*ast.SelectorExpr)
			if !ok || sel.Sel.Name != "ExecuteTemplate" {
				return true
			}
			// Verify the method belongs to html/template or text/template.
			if !asteval.IsTemplateMethod(pkg.TypesInfo, sel) {
				return true
			}
			receiverIdent, ok := sel.X.(*ast.Ident)
			if !ok {
				return true
			}
			obj := pkg.TypesInfo.Uses[receiverIdent]
			if obj == nil {
				return true
			}
			templateName, ok := asteval.BasicLiteralString(call.Args[1])
			if !ok {
				return true
			}
			dataType := pkg.TypesInfo.TypeOf(call.Args[2])
			pending = append(pending, pendingCall{
				call:         call,
				receiverObj:  obj,
				templateName: templateName,
				dataType:     dataType,
			})
			receiverSet[obj] = struct{}{}
			return true
		})
	}

	return pending, receiverSet
}

// resolveTemplates resolves each unique receiver object to its template
// construction chain, including additional ParseFS/Parse modifications.
func resolveTemplates(pkg *packages.Package, receivers map[types.Object]struct{}) (map[types.Object]*resolvedTemplate, []error) {
	resolved := make(map[types.Object]*resolvedTemplate)

	workingDirectory := packageDirectory(pkg)
	embeddedPaths, err := asteval.RelativeFilePaths(workingDirectory, pkg.EmbedFiles...)
	if err != nil {
		return nil, []error{fmt.Errorf("failed to calculate relative path for embedded files: %w", err)}
	}

	var resolveErrs []error

	resolveExpr := func(obj types.Object, name string, expr ast.Expr) {
		if _, needed := receivers[obj]; !needed {
			return
		}
		funcTypeMap := asteval.DefaultFunctions(pkg.Types)
		meta := &asteval.TemplateMetadata{}
		ts, left, right, err := asteval.EvaluateTemplateSelector(nil, pkg.Types, pkg.TypesInfo, expr, workingDirectory, name, "", "", pkg.Fset, pkg.Syntax, embeddedPaths, funcTypeMap, make(map[string]any), meta)
		if err != nil {
			resolveErrs = append(resolveErrs, err)
			return
		}
		resolved[obj] = &resolvedTemplate{
			templates:  ts,
			functions:  funcTypeMap,
			metadata:   meta,
			leftDelim:  left,
			rightDelim: right,
		}
	}

	// Resolve top-level var declarations.
	for _, tv := range astgen.IterateValueSpecs(pkg.Syntax) {
		for i, ident := range tv.Names {
			if i >= len(tv.Values) {
				continue
			}
			obj := pkg.TypesInfo.Defs[ident]
			if obj == nil {
				continue
			}
			resolveExpr(obj, ident.Name, tv.Values[i])
		}
	}

	// Resolve function-local var declarations and short variable declarations.
	for _, file := range pkg.Syntax {
		ast.Inspect(file, func(node ast.Node) bool {
			switch n := node.(type) {
			case *ast.DeclStmt:
				gd, ok := n.Decl.(*ast.GenDecl)
				if !ok || gd.Tok != token.VAR {
					return true
				}
				for _, spec := range gd.Specs {
					vs, ok := spec.(*ast.ValueSpec)
					if !ok {
						continue
					}
					for i, ident := range vs.Names {
						if i >= len(vs.Values) {
							continue
						}
						obj := pkg.TypesInfo.Defs[ident]
						if obj == nil {
							continue
						}
						resolveExpr(obj, ident.Name, vs.Values[i])
					}
				}
			case *ast.AssignStmt:
				if n.Tok != token.DEFINE {
					return true
				}
				for i, lhs := range n.Lhs {
					if i >= len(n.Rhs) {
						continue
					}
					ident, ok := lhs.(*ast.Ident)
					if !ok {
						continue
					}
					obj := pkg.TypesInfo.Defs[ident]
					if obj == nil {
						continue
					}
					resolveExpr(obj, ident.Name, n.Rhs[i])
				}
			}
			return true
		})
	}

	// Find additional ParseFS/Parse calls on resolved template variables.
	for _, file := range pkg.Syntax {
		ast.Inspect(file, func(node ast.Node) bool {
			call, ok := node.(*ast.CallExpr)
			if !ok {
				return true
			}
			obj := asteval.FindModificationReceiver(call, pkg.TypesInfo)
			if obj == nil {
				return true
			}
			rt, ok := resolved[obj]
			if !ok {
				return true
			}
			meta := &asteval.TemplateMetadata{}
			ts, left, right, err := asteval.EvaluateTemplateSelector(rt.templates, pkg.Types, pkg.TypesInfo, call, workingDirectory, "", rt.leftDelim, rt.rightDelim, pkg.Fset, pkg.Syntax, embeddedPaths, rt.functions, make(map[string]any), meta)
			if err != nil {
				return true
			}
			rt.templates = ts
			rt.leftDelim, rt.rightDelim = left, right
			rt.metadata.EmbedFilePaths = append(rt.metadata.EmbedFilePaths, meta.EmbedFilePaths...)
			rt.metadata.ParseCalls = append(rt.metadata.ParseCalls, meta.ParseCalls...)
			rt.metadata.Sources = append(rt.metadata.Sources, meta.Sources...)
			return true
		})
	}

	return resolved, resolveErrs
}

// checkCalls type-checks each pending ExecuteTemplate call against its
// resolved template.
func checkCalls(pkg *packages.Package, pending []pendingCall, resolved map[types.Object]*resolvedTemplate, inspectCall ExecuteTemplateNodeInspectorFunc, inspectTemplate TemplateNodeInspectorFunc) error {
	mergedFunctions := make(Functions)
	if pkg.Types != nil {
		mergedFunctions = DefaultFunctions(pkg.Types)
	}
	for _, rt := range resolved {
		for name, sig := range rt.functions {
			mergedFunctions[name] = sig
		}
		rt.definitions = newDefinitionSet(definitionsFor(pkg.Fset, rt.metadata.Sources))
		rt.definitions.adoptTrees(rt.templates)
	}

	var errs []error
	for _, p := range pending {
		rt, ok := resolved[p.receiverObj]
		if !ok {
			continue
		}
		looked := rt.templates.Lookup(p.templateName)
		if looked == nil {
			continue
		}
		// A name can be given to a template that is never parsed into,
		// which template.New does on its own, and the lookup finds that
		// entry with its tree still nil. There is nothing to walk.
		// Executing such a template fails at run time as an incomplete
		// template, a mistake this does not report yet, but it must not
		// be walked as though it had a body.
		tree := looked.Tree()
		if tree == nil {
			continue
		}
		global := NewGlobal(pkg.Types, pkg.Fset, rt.templates, mergedFunctions)
		global.InspectTemplateNode = inspectTemplate
		global.Definitions = rt.definitions
		if inspectCall != nil {
			inspectCall(p.call, tree, p.dataType, global.definition(p.templateName))
		}
		if err := Execute(global, tree, p.dataType); err != nil {
			errs = append(errs, err)
		}
	}

	return joinErrors(nil, nil, errs...)
}

func packageDirectory(pkg *packages.Package) string {
	if len(pkg.GoFiles) > 0 {
		return filepath.Dir(pkg.GoFiles[0])
	}
	// A package loaded without NeedFiles still positions its syntax in the
	// file set.
	if len(pkg.Syntax) > 0 && pkg.Fset != nil {
		if pos := pkg.Fset.Position(pkg.Syntax[0].Pos()); pos.Filename != "" {
			return filepath.Dir(pos.Filename)
		}
	}
	return "."
}