// Package generate provides functionality for generating stacks from stack files.
package generate

import (
	"context"
	"path/filepath"
	"runtime"
	"slices"
	"sync"

	"github.com/gruntwork-io/terragrunt/internal/component"
	"github.com/gruntwork-io/terragrunt/internal/configbridge"
	"github.com/gruntwork-io/terragrunt/internal/discovery"
	"github.com/gruntwork-io/terragrunt/internal/errors"
	"github.com/gruntwork-io/terragrunt/internal/filter"
	"github.com/gruntwork-io/terragrunt/internal/util"
	"github.com/gruntwork-io/terragrunt/internal/worker"
	"github.com/gruntwork-io/terragrunt/internal/worktrees"
	"github.com/gruntwork-io/terragrunt/pkg/config"
	"github.com/gruntwork-io/terragrunt/pkg/log"
	"github.com/gruntwork-io/terragrunt/pkg/options"
	"golang.org/x/sync/errgroup"
)

// Generator owns the per-working-directory lock for in-process GenerateStacks calls.
type Generator struct {
	locks *util.KeyLocks
}

// NewGenerator returns a ready-to-use Generator.
func NewGenerator() *Generator {
	return &Generator{locks: util.NewKeyLocks()}
}

// StackNode represents a stack file in the file system.
// The parent is the node that generates the current node,
// and children are the nodes that are generated by the current node.
type StackNode struct {
	Parent   *StackNode
	FilePath string
	Children []*StackNode
	Level    int
}

// NewStackNode creates a new stack node.
func NewStackNode(filePath string) *StackNode {
	return &StackNode{
		FilePath: filePath,
		Level:    -1,
		Children: make([]*StackNode, 0),
	}
}

// GenerateStacks generates the stack files using topological ordering to
// prevent race conditions. Stack files are generated level by level, with
// parent stacks completing before their children. Worktrees must be
// provided by the caller if needed; this function will never create
// worktrees internally.
func (g *Generator) GenerateStacks(
	ctx context.Context,
	l log.Logger,
	opts *options.TerragruntOptions,
	wts *worktrees.Worktrees,
) error {
	workingDir, err := util.CanonicalResolvedPath(opts.WorkingDir, opts.WorkingDir)
	if err != nil {
		return &CanonicalizeWorkingDirError{Path: opts.WorkingDir, Err: err}
	}

	g.locks.Lock(workingDir)
	defer g.locks.Unlock(workingDir)

	foundFiles, err := ListStackFiles(ctx, l, opts, wts)
	if err != nil {
		return errors.Errorf("Failed to list stack files in %s %w", opts.WorkingDir, err)
	}

	if len(foundFiles) == 0 {
		if opts.StackAction == "generate" {
			l.Warnf("No stack files found in %s Nothing to generate.", opts.WorkingDir)
		}

		return nil
	}

	// Dedup keyed by canonical stack-file path; accessed only from the main goroutine.
	generatedFiles := make(map[string]bool)

	// Tracks which parent stack file first claimed each generated stack, so repeat claims can be warned.
	claimedBy := make(map[string]string)

	stackTrees := BuildStackTopology(l, foundFiles, workingDir)

	const maxLevel = 1024
	for level := range maxLevel {
		if level == maxLevel-1 {
			return errors.Errorf("Cycle detected: maximum level (%d) exceeded", maxLevel)
		}

		levelNodes := getNodesAtLevel(stackTrees, level)
		if len(levelNodes) == 0 {
			break
		}

		warnOnRepeatedClaims(l, levelNodes, claimedBy)

		if err := generateLevel(ctx, l, opts, level, levelNodes, generatedFiles); err != nil {
			return err
		}

		if err := discoverAndAddNewNodes(ctx, l, opts, wts, workingDir, stackTrees, generatedFiles, level+1); err != nil {
			return err
		}
	}

	return nil
}

// warnOnRepeatedClaims logs a warning when a stack file is claimed by more than
// one parent in the same invocation. All nodes here are stack files by
// construction — ListStackFiles filters to *component.Stack only.
func warnOnRepeatedClaims(l log.Logger, levelNodes []*StackNode, claimedBy map[string]string) {
	for _, node := range levelNodes {
		parent := "root"
		if node.Parent != nil {
			parent = node.Parent.FilePath
		}

		if prev, seen := claimedBy[node.FilePath]; seen && prev != parent {
			l.Warnf("Stack %q is being generated by %q, but it was already generated by %q. "+
				"You have contention between multiple stacks generating to the same directory. "+
				"This is probably a bug in your stack configurations.",
				node.FilePath, parent, prev)
		}

		claimedBy[node.FilePath] = parent
	}
}

// generateLevel handles the concurrent generation of all stack files at a
// given level. See the "Concurrency invariant" note on generatedFiles in
// GenerateStacks: the map is mutated here from the main goroutine only,
// before any wp.Submit call.
func generateLevel(
	ctx context.Context,
	l log.Logger,
	opts *options.TerragruntOptions,
	level int,
	levelNodes []*StackNode,
	generatedFiles map[string]bool,
) error {
	l.Debugf("Generating stack level %d with %d files", level, len(levelNodes))

	wp := worker.NewWorkerPool(opts.Parallelism)
	defer wp.Stop()

	for _, node := range levelNodes {
		if generatedFiles[node.FilePath] {
			continue
		}

		generatedFiles[node.FilePath] = true

		// Best-effort skip; GenerateStackFile surfaces ENOENT if the file is removed in the TOCTOU window.
		if !util.FileExists(node.FilePath) {
			continue
		}

		wp.Submit(func() error {
			_, pctx := configbridge.NewParsingContext(ctx, l, opts)
			return config.GenerateStackFile(ctx, l, pctx, wp, node.FilePath)
		})
	}

	return wp.Wait()
}

// discoverAndAddNewNodes discovers new stack files and adds them to the
// dependency graph. workingDir must already be canonicalised (cleaned,
// absolute, symlink-resolved) so topology lookups key consistently against
// the canonicalised foundFiles.
func discoverAndAddNewNodes(
	ctx context.Context,
	l log.Logger,
	opts *options.TerragruntOptions,
	worktrees *worktrees.Worktrees,
	workingDir string,
	dependencyGraph map[string]*StackNode,
	generatedFiles map[string]bool,
	minLevel int,
) error {
	newFiles, listErr := ListStackFiles(ctx, l, opts, worktrees)
	if listErr != nil {
		return errors.Errorf("Failed to list stack files after level %d: %w", minLevel-1, listErr)
	}

	addNewNodesToGraph(l, dependencyGraph, newFiles, generatedFiles, workingDir)

	return nil
}

// BuildStackTopology creates a topological tree based on directory hierarchy.
func BuildStackTopology(l log.Logger, stackFiles []string, workingDir string) map[string]*StackNode {
	nodes := make(map[string]*StackNode)

	for _, file := range stackFiles {
		nodes[file] = NewStackNode(file)
	}

	for _, node := range nodes {
		assignNodeLevel(l, node, nodes, workingDir)
	}

	return nodes
}

// assignNodeLevel recursively assigns levels to nodes based on directory depth.
func assignNodeLevel(l log.Logger, node *StackNode, allNodes map[string]*StackNode, workingDir string) int {
	if node.Level != -1 {
		return node.Level
	}

	nodeDir := filepath.Dir(node.FilePath)
	parentPath := findParentStackFile(nodeDir, allNodes, workingDir)

	if parentPath == "" {
		node.Level = 0

		return node.Level
	}

	parent := allNodes[parentPath]
	if parent == nil {
		node.Level = 0

		return node.Level
	}

	parentLevel := assignNodeLevel(l, parent, allNodes, workingDir)
	node.Level = parentLevel + 1
	node.Parent = parent
	parent.Children = append(parent.Children, node)

	l.Debugf("Stack %s (level %d) is child of %s (level %d)", node.FilePath, node.Level, parent.FilePath, parent.Level)

	return node.Level
}

// findParentStackFile finds the parent stack file for a given directory.
func findParentStackFile(childDir string, allNodes map[string]*StackNode, workingDir string) string {
	currentDir := childDir

	for {
		parentDir := filepath.Dir(currentDir)
		if parentDir == currentDir {
			break
		}

		if parentDir == workingDir {
			potentialParent := filepath.Join(workingDir, config.DefaultStackFile)
			if _, exists := allNodes[potentialParent]; exists {
				return potentialParent
			}

			break
		}

		potentialParent := filepath.Join(parentDir, config.DefaultStackFile)
		if _, exists := allNodes[potentialParent]; exists {
			return potentialParent
		}

		currentDir = parentDir
	}

	return ""
}

// getNodesAtLevel returns all nodes at a specific level.
func getNodesAtLevel(nodes map[string]*StackNode, level int) []*StackNode {
	var levelNodes []*StackNode

	for _, node := range nodes {
		if node.Level == level {
			levelNodes = append(levelNodes, node)
		}
	}

	return levelNodes
}

// addNewNodesToGraph adds newly discovered stack files to the dependency graph.
func addNewNodesToGraph(
	l log.Logger,
	existingNodes map[string]*StackNode,
	allFiles []string,
	generatedFiles map[string]bool,
	workingDir string,
) {
	newFiles := make([]string, 0)

	for _, file := range allFiles {
		if _, exists := existingNodes[file]; !exists && !generatedFiles[file] {
			newFiles = append(newFiles, file)
		}
	}

	if len(newFiles) == 0 {
		return
	}

	l.Debugf("Adding %d new stack files to topology graph", len(newFiles))

	for _, file := range newFiles {
		existingNodes[file] = NewStackNode(file)
	}

	for _, file := range newFiles {
		node := existingNodes[file]
		assignNodeLevel(l, node, existingNodes, workingDir)
	}
}

// ListStackFiles returns canonical, symlink-resolved stack-file paths under the absolute opts.WorkingDir.
func ListStackFiles(
	ctx context.Context,
	l log.Logger,
	opts *options.TerragruntOptions,
	worktrees *worktrees.Worktrees,
) ([]string, error) {
	d, err := discovery.NewForStackGenerate(l, discovery.StackGenerateOptions{
		WorkingDir:  opts.WorkingDir,
		Filters:     opts.Filters,
		Experiments: opts.Experiments,
	})
	if err != nil {
		return nil, errors.Errorf("Failed to create discovery for stack generate: %w", err)
	}

	discoveredComponents, err := d.Discover(ctx, l, opts)
	if err != nil {
		return nil, errors.Errorf("Failed to discover stack files: %w", err)
	}

	worktreeStacks, err := worktreeStacksToGenerate(ctx, l, opts, worktrees)
	if err != nil {
		return nil, errors.Errorf("Failed to get worktree stacks to generate: %w", err)
	}

	foundFiles := make([]string, 0, len(discoveredComponents)+len(worktreeStacks))
	for _, c := range discoveredComponents {
		if _, ok := c.(*component.Stack); !ok {
			continue
		}

		canonical, err := util.CanonicalResolvedPath(filepath.Join(c.Path(), config.DefaultStackFile), opts.WorkingDir)
		if err != nil {
			return nil, err
		}

		foundFiles = append(foundFiles, canonical)
	}

	for _, c := range worktreeStacks {
		canonical, err := util.CanonicalResolvedPath(filepath.Join(c.Path(), config.DefaultStackFile), opts.WorkingDir)
		if err != nil {
			return nil, err
		}

		foundFiles = append(foundFiles, canonical)
	}

	return foundFiles, nil
}

// ListStackFilesWithExcludes searches for stack files and also returns the set
// of unit paths that should be excluded from the current tofu/terraform command.
// Both results come from a single discovery walk. Stack-file paths and
// excludedPaths keys are canonical symlink-resolved absolute paths; exclusion
// follows discovery's IsActionListed + If logic using opts.TerraformCommand.
func ListStackFilesWithExcludes(
	ctx context.Context,
	l log.Logger,
	opts *options.TerragruntOptions,
	worktrees *worktrees.Worktrees,
) ([]string, map[string]struct{}, error) {
	d, err := discovery.NewForStackGenerate(l, discovery.StackGenerateOptions{
		WorkingDir:  opts.WorkingDir,
		Filters:     opts.Filters,
		Experiments: opts.Experiments,
	})
	if err != nil {
		return nil, nil, errors.Errorf("Failed to create discovery for stack generate: %w", err)
	}

	d = d.WithParseExclude().WithSuppressParseErrors()

	discoveredComponents, err := d.Discover(ctx, l, opts)
	if err != nil {
		return nil, nil, errors.Errorf("Failed to discover stack files: %w", err)
	}

	worktreeStacks, err := worktreeStacksToGenerate(ctx, l, opts, worktrees)
	if err != nil {
		return nil, nil, errors.Errorf("Failed to get worktree stacks to generate: %w", err)
	}

	foundFiles, excludedPaths, err := collectStackAndExcludedPaths(discoveredComponents, opts.WorkingDir)
	if err != nil {
		return nil, nil, err
	}

	foundFiles, err = appendWorktreeStackPaths(foundFiles, worktreeStacks, opts.WorkingDir)
	if err != nil {
		return nil, nil, err
	}

	return foundFiles, excludedPaths, nil
}

// collectStackAndExcludedPaths splits discovered components into canonical stack-file paths and excluded unit paths.
func collectStackAndExcludedPaths(
	components component.Components,
	workingDir string,
) ([]string, map[string]struct{}, error) {
	foundFiles := make([]string, 0, len(components))
	excludedPaths := make(map[string]struct{})

	for _, c := range components {
		switch v := c.(type) {
		case *component.Stack:
			canonical, err := util.CanonicalResolvedPath(filepath.Join(c.Path(), config.DefaultStackFile), workingDir)
			if err != nil {
				return nil, nil, err
			}

			foundFiles = append(foundFiles, canonical)
		case *component.Unit:
			if !v.Excluded() {
				continue
			}

			canonical, err := util.CanonicalResolvedPath(v.Path(), workingDir)
			if err != nil {
				return nil, nil, err
			}

			excludedPaths[canonical] = struct{}{}
		}
	}

	return foundFiles, excludedPaths, nil
}

// appendWorktreeStackPaths appends canonical stack-file paths for worktree stacks to foundFiles.
func appendWorktreeStackPaths(
	foundFiles []string,
	worktreeStacks component.Components,
	workingDir string,
) ([]string, error) {
	for _, c := range worktreeStacks {
		canonical, err := util.CanonicalResolvedPath(filepath.Join(c.Path(), config.DefaultStackFile), workingDir)
		if err != nil {
			return nil, err
		}

		foundFiles = append(foundFiles, canonical)
	}

	return foundFiles, nil
}

// worktreeStacksToGenerate returns a slice of stacks that need to be generated from the worktree stacks.
// When reading filters are present (changes to files read by a stack), it also
// identifies which stacks are affected and records them on the Worktrees object so that the worktree
// discovery phase can walk them for unit-level changes.
func worktreeStacksToGenerate(
	ctx context.Context,
	l log.Logger,
	opts *options.TerragruntOptions,
	w *worktrees.Worktrees,
) (component.Components, error) {
	// If worktrees is nil, there are no worktrees to process, return empty components.
	if w == nil {
		return component.Components{}, nil
	}

	stacksToGenerate := component.NewThreadSafeComponents(component.Components{})

	// If we edit a stack in a worktree, we need to generate it, at the minimum.
	stackDiff := w.Stacks()

	editedStacks := slices.Concat(
		stackDiff.Added,
		stackDiff.Removed,
	)

	for _, changed := range stackDiff.Changed {
		editedStacks = append(editedStacks, changed.FromStack, changed.ToStack)
	}

	for _, stack := range editedStacks {
		stacksToGenerate.EnsureComponent(stack)
	}

	// Build set of stack dirs already handled by direct changes, so we don't
	// duplicate them in ReadingAffected.
	handledStackDirs := make(map[string]struct{}, len(stackDiff.Changed))

	for _, changed := range stackDiff.Changed {
		if dc := changed.ToStack.DiscoveryContext(); dc != nil {
			if rel, err := filepath.Rel(dc.WorkingDir, changed.ToStack.Path()); err == nil {
				handledStackDirs[filepath.Clean(rel)] = struct{}{}
			}
		}
	}

	// When the expanded filter for a given Git expression requires parsing,
	// we need to generate all the stacks in the given worktree, as units within the generated stack might be affected.
	// We also identify which specific stacks are affected by reading filters so that the worktree discovery phase
	// can walk them for unit-level changes.

	g, ctx := errgroup.WithContext(ctx)
	// Allow up to 2 generation tasks per worktree pair (at least 1), capped by available CPUs.
	g.SetLimit(min(runtime.GOMAXPROCS(0), max(1, len(w.WorktreePairs)*2))) //nolint:mnd

	var (
		mu              sync.Mutex
		errs            []error
		readingAffected []worktrees.StackDiffChangedPair
	)

	for _, pair := range w.WorktreePairs {
		_, toFilters, err := pair.Expand()
		if err != nil {
			return nil, errors.Errorf("failed to expand worktree pair: %w", err)
		}

		parseExpr, requiresParse := toFilters.RequiresParse()
		if !requiresParse {
			continue
		}

		g.Go(func() error {
			// Discover all stacks in both worktrees for generation.
			allFromStacks, err := discoverStacks(ctx, l, opts, pair.FromWorktree, false)
			if err != nil {
				mu.Lock()

				errs = append(errs, err)
				mu.Unlock()

				return nil
			}

			// The "to" discovery uses readFiles to force all stacks through
			// the parse phase (populating Reading()), while still returning every stack
			// (because they all match type=stack).
			allToStacks, err := discoverStacks(ctx, l, opts, pair.ToWorktree, true)
			if err != nil {
				mu.Lock()

				errs = append(errs, err)
				mu.Unlock()

				return nil
			}

			for _, c := range allFromStacks {
				stacksToGenerate.EnsureComponent(c)
			}

			for _, c := range allToStacks {
				stacksToGenerate.EnsureComponent(c)
			}

			// Identify which "to" stacks read the changed files.
			matchedToStacks, err := filter.Evaluate(l, parseExpr, allToStacks)
			if err != nil {
				mu.Lock()

				errs = append(errs, err)
				mu.Unlock()

				return nil
			}

			// Record reading-affected stack pairs on the Worktrees object so that
			// discoverChangesInWorktreeStacks walks them for unit-level diffs.
			for _, toComp := range matchedToStacks {
				toStack, ok := toComp.(*component.Stack)
				if !ok {
					continue
				}

				rel, err := filepath.Rel(pair.ToWorktree.Path, toStack.Path())
				if err != nil {
					continue
				}

				if _, handled := handledStackDirs[filepath.Clean(rel)]; handled {
					continue
				}

				// If the stack only exists in one worktree (e.g. newly added),
				// it will be handled by stackDiff.Added/Removed instead.
				fromStack := findStackByRelPath(allFromStacks, pair.FromWorktree.Path, rel)
				if fromStack == nil {
					continue
				}

				mu.Lock()

				readingAffected = append(readingAffected, worktrees.StackDiffChangedPair{
					FromStack: fromStack,
					ToStack:   toStack,
				})
				mu.Unlock()
			}

			return nil
		})
	}

	if err := g.Wait(); err != nil {
		return nil, err
	}

	w.ReadingAffectedStacks = readingAffected

	if len(errs) > 0 {
		return stacksToGenerate.ToComponents(), errors.Join(errs...)
	}

	return stacksToGenerate.ToComponents(), nil
}

// discoverStacks discovers stacks in a worktree.
// User-provided filters from opts.Filters are included (restricted to stacks) so that
// explicit exclusions like --filter '!./land-mine | type=stack' are respected.
// When readFiles is true, all discovered stacks are parsed to populate their Reading
// attribute (used by reading-affected detection).
func discoverStacks(
	ctx context.Context,
	l log.Logger,
	opts *options.TerragruntOptions,
	wt worktrees.Worktree,
	readFiles bool,
) (component.Components, error) {
	allFilters := slices.Concat(stackTypeFilter(), opts.Filters.RestrictToStacks())

	d := discovery.NewDiscovery(wt.Path).
		WithSuppressParseErrors().
		WithFilters(allFilters)

	if readFiles {
		d = d.WithReadFiles()
	}

	components, err := d.Discover(ctx, l, opts)
	if err != nil {
		return nil, errors.Errorf("failed to discover stacks in worktree %s: %w", wt.Ref, err)
	}

	for _, c := range components {
		dc := c.DiscoveryContext()
		if dc == nil {
			dc = &component.DiscoveryContext{}
		}

		dc.WorkingDir = wt.Path
		dc.Ref = wt.Ref
		c.SetDiscoveryContext(dc)
	}

	return components, nil
}

// findStackByRelPath finds a stack in the given components whose path relative to
// worktreePath matches relPath.
func findStackByRelPath(stacks component.Components, worktreePath string, relPath string) *component.Stack {
	for _, c := range stacks {
		s, ok := c.(*component.Stack)
		if !ok {
			continue
		}

		rel, err := filepath.Rel(worktreePath, s.Path())
		if err != nil {
			continue
		}

		if filepath.Clean(rel) == filepath.Clean(relPath) {
			return s
		}
	}

	return nil
}

// stackTypeFilter returns a filter.Filters that restricts to stack components.
func stackTypeFilter() filter.Filters {
	attrExpr := filter.NewTypeExpression(component.StackKind)

	return filter.Filters{filter.NewFilter(attrExpr, attrExpr.String())}
}
