package engine

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"path/filepath"
	"sync"

	"tangled.org/core/notifier"
	"tangled.org/core/spindle/config"
	"tangled.org/core/spindle/db"
	"tangled.org/core/spindle/models"
	"tangled.org/core/spindle/secrets"
	"tangled.org/core/spindle/storage"
)

var (
	ErrTimedOut         = errors.New("timed out")
	ErrWorkflowFailed   = errors.New("workflow failed")
	ErrWorkflowCanceled = errors.New("workflow canceled")
)

var (
	activeMu      sync.Mutex
	activeCancels = make(map[models.WorkflowId]context.CancelCauseFunc)
)

func CancelWorkflow(wid models.WorkflowId) {
	activeMu.Lock()
	cancel, ok := activeCancels[wid]
	activeMu.Unlock()
	if ok {
		cancel(ErrWorkflowCanceled)
	}
}

// user cancel, timeout is DeadlineExceeded
func isCanceled(wfCtx context.Context) bool {
	return errors.Is(context.Cause(wfCtx), ErrWorkflowCanceled)
}

// for when recording early wf cancellations
func writeWfError(db *db.DB, n *notifier.Notifier, l *slog.Logger, wfCtx context.Context, wid models.WorkflowId, phase string, err error) {
	l = l.With("wid", wid, "phase", phase)
	switch {
	case isCanceled(wfCtx):
		l.Info("workflow canceled")
		if dbErr := db.StatusCancelled(wid, "User canceled the workflow", -1, n); dbErr != nil {
			l.Error("failed to set workflow status to cancelled", "err", dbErr)
		}
	case errors.Is(err, ErrTimedOut) || errors.Is(wfCtx.Err(), context.DeadlineExceeded):
		l.Info("workflow timed out")
		if dbErr := db.StatusTimeout(wid, n); dbErr != nil {
			l.Error("failed to set workflow status to timeout", "err", dbErr)
		}
	default:
		l.Error("workflow failed", "err", err)
		if dbErr := db.StatusFailed(wid, err.Error(), -1, n); dbErr != nil {
			l.Error("failed to set workflow status to failed", "err", dbErr)
		}
	}
}

type workflowFinalizer interface {
	FinalizeWorkflow(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, wfLogger models.WorkflowLogger) error
}

func StartWorkflows(l *slog.Logger, vault secrets.Manager, cfg *config.Config, db *db.DB, n *notifier.Notifier, cacheStore storage.Storage, ctx context.Context, pipeline *models.Pipeline, pipelineId models.PipelineId) {
	l.Info("starting all workflows in parallel", "pipeline", pipelineId)

	isTrustedRepo := pipeline.TrustedSource && pipeline.RepoDid != ""
	var allSecrets []secrets.UnlockedSecret
	// never pass secrets to pipelines that run untrusted (e.g. fork) code
	if isTrustedRepo {
		if res, err := vault.GetSecretsUnlocked(ctx, secrets.RepoIdentifier(pipeline.RepoDid.String())); err == nil {
			allSecrets = res
		}
	} else if !pipeline.TrustedSource {
		l.Info("skipping secrets for untrusted pipeline source", "pipeline", pipelineId)
	}
	// untrusted runs cant read or write shared caches
	cacheOwnerDID := ""
	cacheEnabled := cacheStore != nil && isTrustedRepo
	if cacheEnabled {
		repo, err := db.GetRepoByDid(pipeline.RepoDid)
		if err != nil {
			l.Warn("cache owner lookup failed; caching disabled", "repo", pipeline.RepoDid, "err", err)
			cacheEnabled = false
		} else {
			cacheOwnerDID = repo.Owner.String()
		}
	} else if cacheStore != nil && !pipeline.TrustedSource {
		l.Info("skipping caches for untrusted pipeline source", "pipeline", pipelineId)
	}

	// hash the checked commit, not the working tree
	cacheRepoPath, cacheRev := "", ""
	if tm := pipeline.TriggerMetadata; tm != nil {
		if rev, err := models.ExtractCommitSHA(*tm); err == nil {
			did := pipeline.RepoDid.String()
			if tm.SourceRepo != nil && *tm.SourceRepo != "" {
				did = *tm.SourceRepo
			}
			cacheRepoPath, cacheRev = filepath.Join(cfg.Server.RepoDir, did), rev
		} else {
			l.Warn("cannot resolve pipeline commit; cache hashing disabled", "err", err)
		}
	}

	secretValues := make([]string, len(allSecrets))
	for i, s := range allSecrets {
		secretValues[i] = s.Value
	}

	s3, err := NewS3(cfg.S3.LogBucket)
	if err != nil {
		l.Error("error creating s3 client", "err", err)
	}

	// wid.String() is lossy so two different names can map to the same key
	// eg. "foo bar" and "foo-bar"...
	wfCounts := make(map[string]int)
	for _, wfs := range pipeline.Workflows {
		for _, w := range wfs {
			wid := models.WorkflowId{
				PipelineId: pipelineId,
				Name:       w.Name,
			}
			wfCounts[wid.String()]++
		}
	}

	var wg sync.WaitGroup
	for eng, wfs := range pipeline.Workflows {
		workflowTimeout := eng.WorkflowTimeout()
		cacheRunner, cachesSupported := eng.(CacheRunner)
		l.Info("using workflow timeout", "timeout", workflowTimeout)

		for _, w := range wfs {
			w := w
			wid := models.WorkflowId{
				PipelineId: pipelineId,
				Name:       w.Name,
			}

			if wfCounts[wid.String()] > 1 {
				l.Warn("skipping workflow due to name collision", "wid", wid, "key", wid.String())
				dbErr := db.StatusFailed(wid, fmt.Sprintf("colliding workflow name: %s; rename to something else", wid.String()), -1, n)
				if dbErr != nil {
					l.Error("failed to set workflow status to failed", "wid", wid, "err", dbErr)
				}
				continue
			}

			wg.Go(func() {
				if st, err := db.GetStatus(wid); err == nil && models.StatusKind(st.Status).IsFinish() {
					l.Info("skipping finished workflow", "wid", wid, "status", st.Status)
					return
				}
				defer func() {
					if s3 != nil {
						logFile := filepath.Join(cfg.Server.LogDir, fmt.Sprintf("%s.log", wid.String()))
						if err := s3.WriteFile(ctx, logFile); err != nil {
							l.Error("error uploading logs", "err", err)
						}
					}
				}()

				wfLogger, err := models.NewFileWorkflowLogger(cfg.Server.LogDir, wid, secretValues)
				if err != nil {
					l.Warn("failed to setup step logger; logs will not be persisted", "error", err)
					wfLogger = models.NullLogger{}
				} else {
					l.Info("setup step logger; logs will be persisted", "logDir", cfg.Server.LogDir, "wid", wid)
					defer wfLogger.Close()
				}

				timeoutCtx, timeoutCancel := context.WithTimeout(ctx, workflowTimeout)
				defer timeoutCancel()

				wfCtx, userCancel := context.WithCancelCause(timeoutCtx)
				defer userCancel(nil)

				// allow wf context to be cancelled properly by manual cancel
				activeMu.Lock()
				activeCancels[wid] = userCancel
				activeMu.Unlock()
				defer func() {
					activeMu.Lock()
					delete(activeCancels, wid)
					activeMu.Unlock()
				}()

				l.Info("waiting for slot", "wid", wid)
				slot := WorkflowSlot(NoopSlot{})
				if s, ok := eng.(WorkflowSlotter); ok {
					slot, err = s.AcquireWorkflowSlot(wfCtx, wid, &w)
					if err != nil {
						writeWfError(db, n, l, wfCtx, wid, "waiting for slot", err)
						return
					}
				}
				defer slot.Release()

				err = db.StatusRunning(wid, n)
				if err != nil {
					l.Error("failed to set workflow status to running", "wid", wid, "err", err)
					return
				}

				err = eng.SetupWorkflow(wfCtx, wid, &w, wfLogger)
				if err != nil {
					if !isCanceled(wfCtx) {
						if destroyErr := eng.DestroyWorkflow(ctx, wid); destroyErr != nil {
							l.Error("failed to destroy workflow after setup failure", "error", destroyErr)
						}
					}
					writeWfError(db, n, l, wfCtx, wid, "setting up workflow", err)
					return
				}
				defer eng.DestroyWorkflow(ctx, wid)

				ctx, cancel := context.WithTimeout(ctx, workflowTimeout)
				defer cancel()

				var resolvedCaches []ResolvedCache
				if cacheEnabled && len(w.Caches) > 0 {
					if cachesSupported {
						resolvedCaches = ResolveCaches(ctx, l, db, pipeline.RepoDid.String(), w.Engine, cacheRepoPath, cacheRev, w.Caches)
						wfLogger.ControlWriter(CacheRestoreStepIdx, CacheRestoreStep, models.StepStatusStart).Write([]byte{0})
						// caches are an optimization, never a reason to fail the workflow
						restoreStore := cacheStoreForRestore(cacheStore, db, l, resolvedCaches)
						if err := cacheRunner.RestoreCache(ctx, wid, &w, restoreStore, resolvedCaches, wfLogger); err != nil {
							l.Warn("cache restore failed", "wid", wid, "err", err)
						}
						wfLogger.ControlWriter(CacheRestoreStepIdx, CacheRestoreStep, models.StepStatusEnd).Write([]byte{0})
					} else {
						l.Warn("engine does not support caches, skipping restore", "wid", wid)
					}
				}

				// dont save on timeouts, their context is already dead
				saveCaches := func(failed bool) {
					toSave := resolvedCaches[:0]
					for _, rc := range resolvedCaches {
						if rc.saveOn(failed) {
							toSave = append(toSave, rc)
						}
					}
					if len(toSave) == 0 {
						return
					}
					wfLogger.ControlWriter(CacheSaveStepIdx, CacheSaveStep, models.StepStatusStart).Write([]byte{0})
					saveStore, err := prepareCacheSaves(ctx, cacheStore, db, l, cacheOwnerDID, pipeline.RepoDid.String(), w.Engine, toSave)
					if err != nil {
						l.Warn("cache metadata setup failed", "wid", wid, "err", err)
					} else {
						if err := cacheRunner.SaveCache(ctx, wid, &w, saveStore, toSave, wfLogger); err != nil {
							l.Warn("cache save failed", "wid", wid, "err", err)
						}
						saveStore.cleanup(context.WithoutCancel(ctx))
					}
					wfLogger.ControlWriter(CacheSaveStepIdx, CacheSaveStep, models.StepStatusEnd).Write([]byte{0})
				}
				for stepIdx, step := range w.Steps {
					if wfLogger != nil {
						wfLogger.
							ControlWriter(stepIdx, step, models.StepStatusStart).
							Write([]byte{0})
					}

					err = eng.RunStep(wfCtx, wid, &w, stepIdx, allSecrets, wfLogger)

					if wfLogger != nil {
						wfLogger.
							ControlWriter(stepIdx, step, models.StepStatusEnd).
							Write([]byte{0})
					}

					if err != nil {
						if !errors.Is(err, ErrTimedOut) && !errors.Is(wfCtx.Err(), context.DeadlineExceeded) && !isCanceled(wfCtx) {
							saveCaches(true)
						}
						writeWfError(db, n, l, wfCtx, wid, "running step", err)
						return
					}
				}

				saveCaches(false)

				if finalizer, ok := eng.(workflowFinalizer); ok {
					if err := finalizer.FinalizeWorkflow(wfCtx, wid, &w, wfLogger); err != nil {
						writeWfError(db, n, l, wfCtx, wid, "finalizing", err)
						return
					}
				}

				if isCanceled(wfCtx) {
					writeWfError(db, n, l, wfCtx, wid, "before success", nil)
					return
				}

				err = db.StatusSuccess(wid, n)
				if err != nil {
					l.Error("failed to set workflow status to success", "wid", wid, "err", err)
				}
			})
		}
	}

	wg.Wait()
	l.Info("all workflows completed")
}