// 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/discovery"
	"github.com/gruntwork-io/terragrunt/internal/errors"
	"github.com/gruntwork-io/terragrunt/internal/experiment"
	"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"
)

// 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, ensuring parent stacks complete before their children.
// Worktrees must be provided by the caller if needed; this function will never create worktrees internally.
func GenerateStacks(
	ctx context.Context,
	l log.Logger,
	opts *options.TerragruntOptions,
	wts *worktrees.Worktrees,
) error {
	foundFiles, err := ListStackFiles(ctx, l, opts, opts.WorkingDir, 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
	}

	generatedFiles := make(map[string]bool)

	stackTrees := BuildStackTopology(l, foundFiles, opts.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
		}

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

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

	return nil
}

// generateLevel handles the concurrent generation of all stack files at a given level.
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

		// Before attempting to generate the stack file, we need to double-check that the file exists.
		// Generation at a higher level might have resulted in this file being removed.
		if !util.FileExists(node.FilePath) {
			continue
		}

		wp.Submit(func() error {
			return config.GenerateStackFile(ctx, l, opts, wp, node.FilePath)
		})
	}

	return wp.Wait()
}

// discoverAndAddNewNodes discovers new stack files and adds them to the dependency graph.
func discoverAndAddNewNodes(
	ctx context.Context,
	l log.Logger,
	opts *options.TerragruntOptions,
	worktrees *worktrees.Worktrees,
	dependencyGraph map[string]*StackNode,
	generatedFiles map[string]bool,
	minLevel int,
) error {
	newFiles, listErr := ListStackFiles(ctx, l, opts, opts.WorkingDir, worktrees)
	if listErr != nil {
		return errors.Errorf("Failed to list stack files after level %d: %w", minLevel-1, listErr)
	}

	addNewNodesToGraph(l, dependencyGraph, newFiles, generatedFiles, opts.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 searches for stack files in the specified directory.
//
// We only want to use the discovery package when the filter flag experiment is enabled, as we need to filter discovery
// results to ensure that we get the right files back for generation.
func ListStackFiles(
	ctx context.Context,
	l log.Logger,
	opts *options.TerragruntOptions,
	dir string,
	worktrees *worktrees.Worktrees,
) ([]string, error) {
	discovery, err := discovery.NewForStackGenerate(discovery.StackGenerateOptions{
		WorkingDir:    opts.WorkingDir,
		FilterQueries: opts.FilterQueries,
		Experiments:   opts.Experiments,
	})
	if err != nil {
		return nil, errors.Errorf("Failed to create discovery for stack generate: %w", err)
	}

	discoveredComponents, err := discovery.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, opts.Experiments)
	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
		}

		foundFiles = append(foundFiles, filepath.Join(c.Path(), config.DefaultStackFile))
	}

	for _, c := range worktreeStacks {
		foundFiles = append(foundFiles, filepath.Join(c.Path(), config.DefaultStackFile))
	}

	return foundFiles, nil
}

// worktreeStacksToGenerate returns a slice of stacks that need to be generated from the worktree stacks.
func worktreeStacksToGenerate(
	ctx context.Context,
	l log.Logger,
	opts *options.TerragruntOptions,
	w *worktrees.Worktrees,
	experiments experiment.Experiments,
) (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)
	}

	// 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.
	//
	// Based on business logic, the from branch here should never be used, but we'll check it anyways for completeness.
	// We only require parsing for reading filters, and those only trigger in expanded Git expressions when
	// the file is modified (which would result in a toFilter being returned).

	fullDiscoveries := map[string]*discovery.Discovery{}

	for _, pair := range w.WorktreePairs {
		fromFilters, toFilters := pair.Expand()

		if _, requiresParse := fromFilters.RequiresParse(); requiresParse {
			disc, err := discovery.NewForStackGenerate(discovery.StackGenerateOptions{
				WorkingDir:    pair.FromWorktree.Path,
				FilterQueries: []string{"type=stack"},
				Experiments:   experiments,
			})
			if err != nil {
				return nil, errors.Errorf("Failed to create discovery for worktree %s: %w", pair.FromWorktree.Ref, err)
			}

			fullDiscoveries[pair.FromWorktree.Ref] = disc
		}

		if _, requiresParse := toFilters.RequiresParse(); requiresParse {
			disc, err := discovery.NewForStackGenerate(discovery.StackGenerateOptions{
				WorkingDir:    pair.ToWorktree.Path,
				FilterQueries: []string{"type=stack"},
				Experiments:   experiments,
			})
			if err != nil {
				return nil, errors.Errorf("Failed to create discovery for worktree %s: %w", pair.ToWorktree.Ref, err)
			}

			fullDiscoveries[pair.ToWorktree.Ref] = disc
		}
	}

	g, ctx := errgroup.WithContext(ctx)
	g.SetLimit(min(runtime.NumCPU(), len(fullDiscoveries)))

	var (
		mu   sync.Mutex
		errs []error
	)

	for ref, disc := range fullDiscoveries {
		// Create per-iteration local copies to avoid closure capture bug
		refCopy := ref
		discCopy := disc

		g.Go(func() error {
			components, err := discCopy.Discover(ctx, l, opts)
			if err != nil {
				mu.Lock()

				errs = append(errs, errors.Errorf("Failed to discover stacks in worktree %s: %w", refCopy, err))

				mu.Unlock()

				return nil
			}

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

			return nil
		})
	}

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

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

	return stacksToGenerate.ToComponents(), nil
}
