package flow

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"sort"
	"sync"
)

// Phase identifies when a hook fires relative to its step's action.
type Phase string

const (
	PhaseBefore  Phase = "before"
	PhaseAfter   Phase = "after"
	PhaseOnError Phase = "onError"
)

// State is the mutable bag shared across a single flow run. Actions and hooks
// read and write Vars to pass data down the chain (e.g. an authorization
// reference produced by one step and consumed by the next). ID is the domain
// entity the flow operates on (an order id, for the order flows).
type State struct {
	ID     uint64
	Vars   map[string]any
	Logger *slog.Logger
}

// NewState returns an empty state for entity id.
func NewState(id uint64, logger *slog.Logger) *State {
	if logger == nil {
		logger = slog.Default()
	}
	return &State{ID: id, Vars: map[string]any{}, Logger: logger}
}

// Action performs a step's work. params is the step's opaque JSON config.
type Action func(ctx context.Context, st *State, params json.RawMessage) error

// Predicate gates a step (the step runs only when it returns true).
type Predicate func(ctx context.Context, st *State) (bool, error)

// HookInfo describes the step and phase a hook is firing for.
type HookInfo struct {
	Step  string
	Phase Phase
	// Err is the action error when Phase == PhaseOnError, else nil.
	Err error
}

// Hook reacts around a step. Hook errors never abort the flow — they are logged
// — so observability/notification hooks can't break the business transaction.
type Hook func(ctx context.Context, st *State, info HookInfo, params json.RawMessage) error

// CapabilityMeta is the editor-facing documentation attached at registration
// time. The server generates `/sagas/capabilities` from these registrations, so
// UIs should consume this payload instead of hardcoding per-tool help.
//
// Register every action/predicate/hook with the matching *WithMeta helper when
// it should be discoverable in editors. ExampleParams is arbitrary JSON shown as
// a starter payload for params-capable tools.
type CapabilityMeta struct {
	Description   string          `json:"description,omitempty"`
	ExampleParams json.RawMessage `json:"exampleParams,omitempty"`
}

// Registry holds the named actions, predicates and hooks a flow can reference.
type Registry struct {
	mu            sync.RWMutex
	actions       map[string]Action
	actionMeta    map[string]CapabilityMeta
	predicates    map[string]Predicate
	predicateMeta map[string]CapabilityMeta
	hooks         map[string]Hook
	hookMeta      map[string]CapabilityMeta
}

// NewRegistry returns an empty registry.
func NewRegistry() *Registry {
	return &Registry{
		actions:       map[string]Action{},
		actionMeta:    map[string]CapabilityMeta{},
		predicates:    map[string]Predicate{},
		predicateMeta: map[string]CapabilityMeta{},
		hooks:         map[string]Hook{},
		hookMeta:      map[string]CapabilityMeta{},
	}
}

// Action registers an action under name (overwrites an existing one).
func (r *Registry) Action(name string, fn Action) {
	r.ActionWithMeta(name, fn, CapabilityMeta{})
}

// ActionWithMeta registers an action plus the metadata editors should display.
func (r *Registry) ActionWithMeta(name string, fn Action, meta CapabilityMeta) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.actions[name] = fn
	r.actionMeta[name] = meta
}

// Predicate registers a predicate under name.
func (r *Registry) Predicate(name string, fn Predicate) {
	r.PredicateWithMeta(name, fn, CapabilityMeta{})
}

// PredicateWithMeta registers a predicate plus the metadata editors should display.
func (r *Registry) PredicateWithMeta(name string, fn Predicate, meta CapabilityMeta) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.predicates[name] = fn
	r.predicateMeta[name] = meta
}

// Hook registers a hook under name.
func (r *Registry) Hook(name string, fn Hook) {
	r.HookWithMeta(name, fn, CapabilityMeta{})
}

// HookWithMeta registers a hook plus the metadata editors should display.
func (r *Registry) HookWithMeta(name string, fn Hook, meta CapabilityMeta) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.hooks[name] = fn
	r.hookMeta[name] = meta
}

func (r *Registry) action(name string) (Action, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	a, ok := r.actions[name]
	return a, ok
}

func (r *Registry) predicate(name string) (Predicate, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	p, ok := r.predicates[name]
	return p, ok
}

func (r *Registry) hook(name string) (Hook, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	h, ok := r.hooks[name]
	return h, ok
}

// Capabilities is the set of registered names, surfaced so an editor UI can
// offer exactly the actions/predicates/hooks a flow may reference.
type Capabilities struct {
	Actions       []string                  `json:"actions"`
	ActionMeta    map[string]CapabilityMeta `json:"actionMeta,omitempty"`
	Predicates    []string                  `json:"predicates"`
	PredicateMeta map[string]CapabilityMeta `json:"predicateMeta,omitempty"`
	Hooks         []string                  `json:"hooks"`
	HookMeta      map[string]CapabilityMeta `json:"hookMeta,omitempty"`
}

// Capabilities returns the sorted registered names.
func (r *Registry) Capabilities() Capabilities {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return Capabilities{
		Actions:       sortedKeys(r.actions),
		ActionMeta:    cloneMeta(r.actionMeta),
		Predicates:    sortedKeys(r.predicates),
		PredicateMeta: cloneMeta(r.predicateMeta),
		Hooks:         sortedKeys(r.hooks),
		HookMeta:      cloneMeta(r.hookMeta),
	}
}

func cloneMeta(in map[string]CapabilityMeta) map[string]CapabilityMeta {
	out := make(map[string]CapabilityMeta, len(in))
	for k, v := range in {
		out[k] = v
	}
	return out
}

func sortedKeys[V any](m map[string]V) []string {
	out := make([]string, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	sort.Strings(out)
	return out
}

// StepResult records the outcome of one step.
type StepResult struct {
	Name    string `json:"name"`
	Skipped bool   `json:"skipped,omitempty"`
	Error   string `json:"error,omitempty"`
}

// Result is the outcome of a whole flow run.
type Result struct {
	Flow        string       `json:"flow"`
	Steps       []StepResult `json:"steps"`
	Compensated []string     `json:"compensated,omitempty"`
	Failed      bool         `json:"failed"`
}

// Engine runs flow definitions against a registry.
type Engine struct {
	reg    *Registry
	logger *slog.Logger
}

// NewEngine returns an engine bound to reg.
func NewEngine(reg *Registry, logger *slog.Logger) *Engine {
	if logger == nil {
		logger = slog.Default()
	}
	return &Engine{reg: reg, logger: logger}
}

// Validate checks that every action, predicate and hook a definition references
// is registered. Call it at load time to fail fast on a bad flow config.
func (e *Engine) Validate(def *Definition) error {
	for _, s := range def.Steps {
		if _, ok := e.reg.action(s.Action); !ok {
			return fmt.Errorf("flow %q step %q: unknown action %q", def.Name, s.Name, s.Action)
		}
		if s.When != "" {
			if _, ok := e.reg.predicate(s.When); !ok {
				return fmt.Errorf("flow %q step %q: unknown predicate %q", def.Name, s.Name, s.When)
			}
		}
		if s.Compensate != nil {
			if _, ok := e.reg.action(s.Compensate.Action); !ok {
				return fmt.Errorf("flow %q step %q: unknown compensate action %q", def.Name, s.Name, s.Compensate.Action)
			}
		}
		for _, hs := range [][]HookRef{s.Hooks.Before, s.Hooks.After, s.Hooks.OnError} {
			for _, h := range hs {
				if _, ok := e.reg.hook(h.Type); !ok {
					return fmt.Errorf("flow %q step %q: unknown hook %q", def.Name, s.Name, h.Type)
				}
				if h.When != "" {
					if _, ok := e.reg.predicate(h.When); !ok {
						return fmt.Errorf("flow %q step %q: hook %q: unknown predicate %q", def.Name, s.Name, h.Type, h.When)
					}
				}
			}
		}
	}
	return nil
}

// Run executes the flow. On a step error (without ContinueOnError) it fires the
// step's on-error hooks, compensates every already-executed step that declared a
// compensating action (in reverse), and returns the error. Hook failures are
// logged, never fatal.
func (e *Engine) Run(ctx context.Context, def *Definition, st *State) (*Result, error) {
	res := &Result{Flow: def.Name}
	var executed []Step // steps that completed, for reverse compensation

	for _, step := range def.Steps {
		if step.When != "" {
			pred, ok := e.reg.predicate(step.When)
			if !ok {
				return res, e.fail(res, step, fmt.Errorf("unknown predicate %q", step.When))
			}
			run, err := pred(ctx, st)
			if err != nil {
				return res, e.abort(ctx, res, &executed, step, err, st)
			}
			if !run {
				res.Steps = append(res.Steps, StepResult{Name: step.Name, Skipped: true})
				st.Logger.Info("flow step skipped", "flow", def.Name, "step", step.Name)
				continue
			}
		}

		action, ok := e.reg.action(step.Action)
		if !ok {
			return res, e.fail(res, step, fmt.Errorf("unknown action %q", step.Action))
		}

		e.fireHooks(ctx, st, step.Hooks.Before, HookInfo{Step: step.Name, Phase: PhaseBefore})

		err := action(ctx, st, step.Params)
		if err != nil {
			e.fireHooks(ctx, st, step.Hooks.OnError, HookInfo{Step: step.Name, Phase: PhaseOnError, Err: err})
			if step.ContinueOnError {
				res.Steps = append(res.Steps, StepResult{Name: step.Name, Error: err.Error()})
				st.Logger.Warn("flow step failed, continuing", "flow", def.Name, "step", step.Name, "err", err)
				executed = append(executed, step)
				continue
			}
			return res, e.abort(ctx, res, &executed, step, err, st)
		}

		e.fireHooks(ctx, st, step.Hooks.After, HookInfo{Step: step.Name, Phase: PhaseAfter})
		res.Steps = append(res.Steps, StepResult{Name: step.Name})
		executed = append(executed, step)
	}
	return res, nil
}

// fail records a step error without compensation (used for config errors found
// before the action ran, so nothing needs undoing for this step).
func (e *Engine) fail(res *Result, step Step, err error) error {
	res.Failed = true
	res.Steps = append(res.Steps, StepResult{Name: step.Name, Error: err.Error()})
	return fmt.Errorf("flow %q step %q: %w", res.Flow, step.Name, err)
}

// abort records the failing step, then compensates executed steps in reverse.
func (e *Engine) abort(ctx context.Context, res *Result, executed *[]Step, step Step, cause error, st *State) error {
	res.Failed = true
	res.Steps = append(res.Steps, StepResult{Name: step.Name, Error: cause.Error()})
	e.compensate(ctx, res, *executed, st)
	return fmt.Errorf("flow %q step %q: %w", res.Flow, step.Name, cause)
}

// compensate runs the compensating action of each executed step in reverse.
func (e *Engine) compensate(ctx context.Context, res *Result, executed []Step, st *State) {
	for i := len(executed) - 1; i >= 0; i-- {
		step := executed[i]
		if step.Compensate == nil {
			continue
		}
		action, ok := e.reg.action(step.Compensate.Action)
		if !ok {
			st.Logger.Error("flow compensation skipped: unknown action",
				"flow", res.Flow, "step", step.Name, "action", step.Compensate.Action)
			continue
		}
		if err := action(ctx, st, step.Compensate.Params); err != nil {
			st.Logger.Error("flow compensation failed", "flow", res.Flow, "step", step.Name, "err", err)
			continue
		}
		res.Compensated = append(res.Compensated, step.Name)
	}
}

// fireHooks runs every hook for a phase, logging (never propagating) failures.
// Hooks with a When predicate are evaluated first; the hook is skipped when the
// predicate returns false.
func (e *Engine) fireHooks(ctx context.Context, st *State, refs []HookRef, info HookInfo) {
	for _, ref := range refs {
		if ref.When != "" {
			pred, ok := e.reg.predicate(ref.When)
			if !ok {
				st.Logger.Error("flow hook skipped: unknown predicate", "hook", ref.Type, "predicate", ref.When, "step", info.Step)
				continue
			}
			run, err := pred(ctx, st)
			if err != nil {
				st.Logger.Error("flow hook predicate error", "hook", ref.Type, "predicate", ref.When, "step", info.Step, "err", err)
				continue
			}
			if !run {
				st.Logger.Debug("flow hook skipped by predicate", "hook", ref.Type, "predicate", ref.When, "step", info.Step)
				continue
			}
		}
		h, ok := e.reg.hook(ref.Type)
		if !ok {
			st.Logger.Error("flow hook skipped: unknown type", "hook", ref.Type, "step", info.Step)
			continue
		}
		if err := h(ctx, st, info, ref.Params); err != nil {
			st.Logger.Warn("flow hook error", "hook", ref.Type, "step", info.Step, "phase", info.Phase, "err", err)
		}
	}
}