package runnerpool

import (
	"context"
	"sync"

	"github.com/gruntwork-io/terragrunt/pkg/options"

	"github.com/gruntwork-io/terragrunt/internal/component"
	"github.com/gruntwork-io/terragrunt/internal/multierror"

	"github.com/gruntwork-io/terragrunt/pkg/log"

	"github.com/gruntwork-io/terragrunt/internal/queue"
	"github.com/gruntwork-io/terragrunt/internal/telemetry"

	"github.com/puzpuzpuz/xsync/v4"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
)

// UnitRunner defines a function type that executes a Unit within a given context and returns an error.
type UnitRunner func(ctx context.Context, u *component.Unit) error

// Controller orchestrates concurrent execution over a DAG.
type Controller struct {
	q           *queue.Queue
	runner      UnitRunner
	readyCh     chan struct{}
	unitsMap    map[string]*component.Unit
	concurrency int
}

// ControllerOption is a function that modifies a Controller.
type ControllerOption func(*Controller)

// WithRunner sets the UnitRunner for the Controller.
func WithRunner(runner UnitRunner) ControllerOption {
	return func(dr *Controller) {
		dr.runner = runner
	}
}

// WithMaxConcurrency sets the concurrency for the Controller.
func WithMaxConcurrency(concurrency int) ControllerOption {
	return func(dr *Controller) {
		if concurrency <= 0 {
			concurrency = 1
		}

		dr.concurrency = concurrency
	}
}

// NewController creates a new Controller with the given options and a pre-built queue, which must not be nil.
func NewController(q *queue.Queue, units []*component.Unit, opts ...ControllerOption) *Controller {
	dr := &Controller{
		q:           q,
		readyCh:     make(chan struct{}, 1), // buffered to avoid blocking
		concurrency: options.DefaultParallelism,
	}
	// Map to link runner Units and Queue Entries
	unitsMap := make(map[string]*component.Unit)

	for _, u := range units {
		if u != nil && u.Path() != "" {
			unitsMap[u.Path()] = u
		}
	}

	dr.unitsMap = unitsMap
	for _, opt := range opts {
		opt(dr)
	}

	return dr
}

// Run executes the Queue return error summarizing all entries that failed to run.
func (dr *Controller) Run(ctx context.Context, l log.Logger) error {
	return telemetry.TelemeterFromContext(ctx).
		Collect(ctx, l, "runner_pool_controller", map[string]any{
			"total_tasks":             len(dr.q.Entries),
			"concurrency":             dr.concurrency,
			"fail_fast":               dr.q.FailFast,
			"ignore_dependency_order": dr.q.IgnoreDependencyOrder,
		}, func(childCtx context.Context, l log.Logger) error {
			var (
				wg      sync.WaitGroup
				sem     = make(chan struct{}, dr.concurrency)
				results = xsync.NewMap[string, error]()
			)

			if dr.runner == nil {
				return ErrRunnerNotSet
			}

			l.Debugf("Runner Pool Controller: starting with %d tasks, concurrency %d",
				len(dr.q.Entries), dr.concurrency)

			// Initial signal to start scheduling
			select {
			case dr.readyCh <- struct{}{}:
			default:
			}

			for {
				readyEntries := dr.q.GetReadyWithDependencies(l)
				l.Debugf("Runner Pool Controller: found %d readyEntries tasks", len(readyEntries))

				for _, e := range readyEntries {
					if !dr.q.ClaimForRunning(e) {
						l.Debugf(
							"Runner Pool Controller: skipping %s; fail-fast cancelled before dispatch",
							e.Component.Path(),
						)

						continue
					}

					l.Debugf("Runner Pool Controller: running %s", e.Component.Path())

					sem <- struct{}{}

					wg.Add(1)

					go func(ent *queue.Entry) {
						defer func() {
							<-sem
							wg.Done()

							select {
							case dr.readyCh <- struct{}{}:
							default:
							}
						}()

						unit := dr.unitsMap[ent.Component.Path()]
						if unit == nil {
							err := NewUnitNotDiscoveredError(ent.Component.Path())
							l.Errorf(
								"Runner Pool Controller: unit for path %s not found in discovered units, skipping execution",
								ent.Component.Path(),
							)
							dr.q.FailEntry(ent)
							results.Store(ent.Component.Path(), err)

							return
						}

						err := dr.runner(childCtx, unit)
						results.Store(ent.Component.Path(), err)

						if err != nil {
							l.Debugf("Runner Pool Controller: %s failed", ent.Component.Path())
							dr.q.FailEntry(ent)

							return
						}

						l.Debugf("Runner Pool Controller: %s succeeded", ent.Component.Path())
						dr.q.SetEntryStatus(ent, queue.StatusSucceeded)
					}(e)
				}

				if dr.q.Finished() {
					break
				}

				select {
				case <-dr.readyCh:
				case <-childCtx.Done():
					wg.Wait()
					return nil
				}
			}

			wg.Wait()

			var errCollector []error

			var succeeded, failed, earlyExit int

			for _, entry := range dr.q.Entries {
				switch entry.Status {
				case queue.StatusSucceeded:
					succeeded++
				case queue.StatusFailed:
					failed++
				case queue.StatusEarlyExit:
					earlyExit++
				case queue.StatusPending,
					queue.StatusBlocked,
					queue.StatusUnsorted,
					queue.StatusReady,
					queue.StatusRunning:
					// Non-terminal states are not counted in the summary.
				}

				if err, ok := results.Load(entry.Component.Path()); ok {
					if err == nil {
						continue
					}

					errCollector = append(errCollector, err)

					continue
				}

				if entry.Status == queue.StatusEarlyExit {
					failedDep := findFailedDependency(entry, dr.q)
					errCollector = append(
						errCollector,
						NewUnitEarlyExitError(entry.Component.Path(), failedDep),
					)
				}

				if entry.Status == queue.StatusFailed {
					errCollector = append(errCollector, NewUnitFailedError(entry.Component.Path()))
				}
			}

			if span := trace.SpanFromContext(childCtx); span.IsRecording() {
				span.SetAttributes(
					attribute.Int("tasks_succeeded", succeeded),
					attribute.Int("tasks_failed", failed),
					attribute.Int("tasks_early_exit", earlyExit),
				)
			}

			return multierror.Join(errCollector...)
		})
}
