package orchestrator

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/store/workflow/run"
	"github.com/tidbcloud/AgentTeam/internal/workflow/plan"
	runkernel "github.com/tidbcloud/AgentTeam/internal/workflow/run"
)

const (
	defaultRunClaimLimit            = 32
	defaultRunTurnLease             = 30 * time.Second
	defaultMaxReductions            = 128
	runtimeCollectionPage           = 500
	maxRuntimeChildren              = plan.MaxRunActivations
	initialRunStartedEvent          = "run_started"
	runtimeCollectionLimitErrorCode = "run.runtime_collection_limit_exceeded"
)

var errRuntimeCollectionLimit = errors.New("workflow Run runtime collection limit exceeded")

// Material is the immutable executable authority for a Run. Implementations
// load the exact PlanVersion pinned on the Invocation; they must never follow
// mutable Plan current/default roles or the Binding active revision.
type Material struct {
	WorkspaceID     string
	PlanID          string
	PlanVersionID   string
	CompilerVersion string
	Digest          string
	ArtifactJSON    json.RawMessage
	SourceKind      string
}

type MaterialLoader interface {
	LoadRunMaterial(context.Context, run.Run) (Material, error)
}

// ArtifactStateMaterializer is optional for test repositories and mandatory
// for the production TiDB store. It expands only fully validated references
// before pure reducer evaluation; ProjectTurn compacts them again atomically.
type ArtifactStateMaterializer interface {
	MaterializeRunState(context.Context, string, string, []byte) ([]byte, error)
}

// Repository is the narrow durable boundary used by the reducer worker.
// Effect dispatch and callback decisions have separate fenced methods and do
// not hold a Run turn lease while performing network I/O.
type Repository interface {
	ClaimRunnableRuns(context.Context, run.ClaimRunnableParams) ([]run.ClaimedRun, error)
	ListActivations(context.Context, run.ListChildrenParams) ([]run.Activation, error)
	ListEffects(context.Context, run.ListChildrenParams) ([]run.Effect, error)
	SaveTurn(context.Context, run.SaveTurnParams) (run.SaveTurnResult, error)
	FailClaimedRun(context.Context, run.FailClaimedRunParams) (run.FailClaimedRunResult, error)
}

type EngineOptions struct {
	WorkerID      string
	ClaimLimit    int
	TurnLeaseTTL  time.Duration
	MaxReductions int
	Now           func() time.Time
	// AfterTurnCommitted is a test-only crash seam. Production leaves it nil;
	// when injected, the callback runs after SaveTurn commits so a fresh engine
	// must recover the exact durable reference rather than replaying a partial
	// reducer turn.
	AfterTurnCommitted func() error
	// AfterAggregateArtifactCommitted fires only after the exact turn that
	// durably finalized an oversized Join aggregate. It is a test-only crash
	// seam; production leaves it nil.
	AfterAggregateArtifactCommitted func() error
}

type Engine struct {
	repository                      Repository
	materials                       MaterialLoader
	workerID                        string
	claimLimit                      int
	turnLeaseTTL                    time.Duration
	maxReductions                   int
	now                             func() time.Time
	afterTurnCommitted              func() error
	afterAggregateArtifactCommitted func() error
}

func NewEngine(repository Repository, materials MaterialLoader, options EngineOptions) (*Engine, error) {
	if repository == nil || materials == nil {
		return nil, fmt.Errorf("workflow Run repository and material loader are required")
	}
	options.WorkerID = strings.TrimSpace(options.WorkerID)
	if options.WorkerID == "" || len(options.WorkerID) > 128 {
		return nil, fmt.Errorf("workflow Run workerId is required and must be bounded")
	}
	if options.ClaimLimit == 0 {
		options.ClaimLimit = defaultRunClaimLimit
	}
	if options.ClaimLimit < 1 || options.ClaimLimit > runtimeCollectionPage {
		return nil, fmt.Errorf("workflow Run claim limit must be between 1 and %d", runtimeCollectionPage)
	}
	if options.TurnLeaseTTL == 0 {
		options.TurnLeaseTTL = defaultRunTurnLease
	}
	if options.TurnLeaseTTL <= 0 {
		return nil, fmt.Errorf("workflow Run turn lease TTL must be positive")
	}
	if options.MaxReductions == 0 {
		options.MaxReductions = defaultMaxReductions
	}
	if options.MaxReductions < 1 || options.MaxReductions > 10_000 {
		return nil, fmt.Errorf("workflow Run max reductions must be between 1 and 10000")
	}
	if options.Now == nil {
		options.Now = time.Now
	}
	return &Engine{
		repository: repository, materials: materials, workerID: options.WorkerID,
		claimLimit: options.ClaimLimit, turnLeaseTTL: options.TurnLeaseTTL,
		maxReductions: options.MaxReductions, now: options.Now, afterTurnCommitted: options.AfterTurnCommitted,
		afterAggregateArtifactCommitted: options.AfterAggregateArtifactCommitted,
	}, nil
}

type RunOnceOutcome struct {
	Claimed   int
	Succeeded int
	Failed    int
}

// RunOnce executes one bounded scan. A bad Run does not prevent the remaining
// independently leased Runs from progressing; errors are joined for logging.
func (engine *Engine) RunOnce(ctx context.Context) (RunOnceOutcome, error) {
	if engine == nil || engine.repository == nil {
		return RunOnceOutcome{}, fmt.Errorf("workflow Run engine is not configured")
	}
	scanAt := engine.now().UTC().Truncate(time.Microsecond)
	claimed, err := engine.repository.ClaimRunnableRuns(ctx, run.ClaimRunnableParams{
		WorkerID: engine.workerID, Now: scanAt, LeaseTTL: engine.turnLeaseTTL, Limit: engine.claimLimit,
	})
	if err != nil {
		return RunOnceOutcome{}, err
	}
	outcome := RunOnceOutcome{Claimed: len(claimed)}
	var failures []error
	for _, claim := range claimed {
		// Loading and reducing a prior Run can consume a meaningful fraction of
		// the batch. Give each independently claimed Run its actual processing
		// time rather than reusing the scan timestamp for deadline evaluation and
		// its fenced SaveTurn.
		processAt := engine.now().UTC().Truncate(time.Microsecond)
		if err := engine.ProcessClaim(ctx, claim, processAt); err != nil {
			outcome.Failed++
			failures = append(failures, fmt.Errorf("process workflow Run %s: %w", claim.Run.ID, err))
			continue
		}
		outcome.Succeeded++
	}
	return outcome, errors.Join(failures...)
}

// ProcessClaim executes only pure reductions and one bounded atomic SaveTurn.
// It never performs a network effect while the Run lease is held.
func (engine *Engine) ProcessClaim(ctx context.Context, claim run.ClaimedRun, now time.Time) error {
	if engine == nil || engine.repository == nil || engine.materials == nil {
		return fmt.Errorf("workflow Run engine is not configured")
	}
	if now.IsZero() {
		now = engine.now()
	}
	now = now.UTC().Truncate(time.Microsecond)
	material, err := engine.materials.LoadRunMaterial(ctx, claim.Run)
	if err != nil {
		if !IsPermanentMaterialError(err) {
			// An unknown loader error is infrastructure failure, not a durable
			// workflow outcome. Leave the claimed snapshot unchanged and let the
			// lease expire so a later turn can retry the exact pinned material.
			return err
		}
		return engine.saveUnrecoverableFailure(ctx, claim, "run.plan_version_unavailable", err, now)
	}
	graph, err := validateMaterial(claim.Run, material)
	if err != nil {
		return engine.saveUnrecoverableFailure(ctx, claim, "run.plan_version_invalid", err, now)
	}
	kernel, err := runkernel.NewKernel(graph)
	if err != nil {
		return engine.saveUnrecoverableFailure(ctx, claim, "run.plan_version_invalid", err, now)
	}
	activations, err := engine.listAllActivations(ctx, claim.Run)
	if err != nil {
		if errors.Is(err, errRuntimeCollectionLimit) {
			return engine.saveUnrecoverableFailure(ctx, claim, runtimeCollectionLimitErrorCode, err, now)
		}
		return err
	}
	effects, err := engine.listAllEffects(ctx, claim.Run)
	if err != nil {
		if errors.Is(err, errRuntimeCollectionLimit) {
			return engine.saveUnrecoverableFailure(ctx, claim, runtimeCollectionLimitErrorCode, err, now)
		}
		return err
	}
	if materializer, ok := engine.repository.(ArtifactStateMaterializer); ok {
		for index := range effects {
			if (effects[index].Status != run.EffectStatusCompleted && effects[index].Status != run.EffectStatusFailed) || len(bytes.TrimSpace(effects[index].ResultJSON)) == 0 {
				continue
			}
			materialized, materializeErr := materializer.MaterializeRunState(ctx, claim.Run.WorkspaceID, claim.Run.ID, effects[index].ResultJSON)
			if materializeErr != nil {
				return engine.saveUnrecoverableFailure(ctx, claim, "run.artifact_unavailable", materializeErr, now)
			}
			effects[index].ResultJSON = materialized
		}
	}
	stateJSON := claim.Run.StateJSON
	if materializer, ok := engine.repository.(ArtifactStateMaterializer); ok {
		stateJSON, err = materializer.MaterializeRunState(ctx, claim.Run.WorkspaceID, claim.Run.ID, stateJSON)
		if err != nil {
			return engine.saveUnrecoverableFailure(ctx, claim, "run.artifact_unavailable", err, now)
		}
	}
	state, initialized, err := decodeRunState(stateJSON)
	if err != nil {
		return engine.saveUnrecoverableFailure(ctx, claim, "run.state_invalid", err, now)
	}
	aggregate := runkernel.Transition{State: state, Intents: []runkernel.Intent{}, Events: []runkernel.Event{}}
	if !initialized {
		var input runkernel.Values
		if err := decodeValues(claim.Run.InputJSON, &input); err != nil {
			return engine.saveUnrecoverableFailure(ctx, claim, "run.input_invalid", err, now)
		}
		state, err = kernel.Start(input)
		if err != nil {
			return engine.saveUnrecoverableFailure(ctx, claim, "run.input_contract_failed", err, now)
		}
		aggregate.State = state
		aggregate.Events = append(aggregate.Events, runkernel.Event{Kind: initialRunStartedEvent})
	} else if err := validateRestoredState(state, graph); err != nil {
		return engine.saveUnrecoverableFailure(ctx, claim, "run.state_invalid", err, now)
	}

	reductions := 0
	// A callback message receipt is authoritative before the Agent replies: the
	// ordinary Task message transaction has already promoted the selected Agent
	// to Task Watch. Freeze that authoritative identity into the pure Run state
	// before any sibling callback can use the projected frontier as coverage.
	for _, effect := range stableEffects(effects) {
		if reductions >= engine.maxReductions {
			break
		}
		activationID, responderAgentID, available, err := availableTaskWatchResponderObservation(
			aggregate.State,
			effect,
		)
		if err != nil {
			return engine.saveStateFailure(ctx, claim, graph, aggregate.State, activations, effects, material.SourceKind, "run.effect_receipt_invalid", err, now)
		}
		if !available {
			continue
		}
		transition, err := kernel.ObserveCallbackResponder(aggregate.State, activationID, responderAgentID)
		if err != nil {
			return engine.saveStateFailure(ctx, claim, graph, aggregate.State, activations, effects, material.SourceKind, "run.effect_receipt_invalid", err, now)
		}
		mergeTransition(&aggregate, transition)
		reductions++
	}
	// A definitive no-receipt result proves that the projected successor never
	// entered Task Watch. Restore only that bookkeeping before any graph
	// completion can cancel/collapse the activation; the normal completion pass
	// below still owns failure routing and timestamp order.
	for _, effect := range stableEffects(effects) {
		if reductions >= engine.maxReductions {
			break
		}
		activationID, available, err := availableTaskWatchPreDeliveryRestoration(
			aggregate.State,
			effect,
		)
		if err != nil {
			return engine.saveStateFailure(ctx, claim, graph, aggregate.State, activations, effects, material.SourceKind, "run.effect_receipt_invalid", err, now)
		}
		if !available {
			continue
		}
		transition, err := kernel.RestoreUndeliveredCallbackTaskWatch(aggregate.State, activationID)
		if err != nil {
			return engine.saveStateFailure(ctx, claim, graph, aggregate.State, activations, effects, material.SourceKind, "run.effect_receipt_invalid", err, now)
		}
		mergeTransition(&aggregate, transition)
		reductions++
	}

	// Do not let an unrelated completion (notably a join_any winner) cancel
	// and collapse a projected Task callback while its message transaction is
	// pending or in flight. The receipt observation above removes this barrier
	// immediately after delivery; a definitive pre-delivery failure is reduced
	// below and restores the predecessor instead. This still serializes only the
	// short delivery-confirmation window, not Agent execution.
	taskWatchDispatchBarrier, err := hasUnobservedTaskWatchDispatch(aggregate.State, effects)
	if err != nil {
		return engine.saveStateFailure(ctx, claim, graph, aggregate.State, activations, effects, material.SourceKind, "run.effect_receipt_invalid", err, now)
	}
	if !taskWatchDispatchBarrier {
		for _, effect := range stableEffects(effects) {
			if reductions >= engine.maxReductions {
				break
			}
			completion, available, err := availableEffectCompletion(aggregate.State, effect)
			if err != nil {
				return engine.saveStateFailure(ctx, claim, graph, aggregate.State, activations, effects, material.SourceKind, "run.effect_receipt_invalid", err, now)
			}
			if !available {
				continue
			}
			transition, err := kernel.Complete(aggregate.State, completion, now)
			if err != nil {
				return engine.saveStateFailure(ctx, claim, graph, aggregate.State, activations, effects, material.SourceKind, "run.effect_completion_invalid", err, now)
			}
			mergeTransition(&aggregate, transition)
			reductions++
		}
	}
	for !taskWatchDispatchBarrier && reductions < engine.maxReductions && len(runkernel.StableReadyIDs(aggregate.State)) != 0 {
		transition, err := kernel.Step(aggregate.State, now)
		if err != nil {
			return engine.saveStateFailure(ctx, claim, graph, aggregate.State, activations, effects, material.SourceKind, "run.reducer_failed", err, now)
		}
		mergeTransition(&aggregate, transition)
		reductions++
		if transitionHasTaskWatchDispatch(transition) {
			// Persist and dispatch one Task callback message at a time until its
			// receipt is observed. Agent work still runs in parallel after each
			// quick delivery transaction; only the unsafe pre-delivery projection
			// window is serialized.
			taskWatchDispatchBarrier = true
		}
	}
	hasMoreWork := false
	if !taskWatchDispatchBarrier {
		hasMoreWork, err = hasReducibleWork(aggregate.State, effects)
		if err != nil {
			return engine.saveStateFailure(ctx, claim, graph, aggregate.State, activations, effects, material.SourceKind, "run.effect_receipt_invalid", err, now)
		}
	}
	if engine.afterAggregateArtifactCommitted != nil {
		_, err = ProjectAndSaveTurnWithArtifactCommitHook(
			ctx, engine.repository, claim, graph, aggregate, activations, effects, material.SourceKind, now,
			hasMoreWork, taskWatchDispatchBarrier, engine.afterAggregateArtifactCommitted,
		)
	} else {
		_, err = ProjectAndSaveTurn(
			ctx, engine.repository, claim, graph, aggregate, activations, effects, material.SourceKind, now,
			hasMoreWork, taskWatchDispatchBarrier,
		)
	}
	if err != nil {
		if code := runtimeProjectionLimitCode(err); code != "" {
			return engine.saveUnrecoverableFailure(ctx, claim, code, err, now)
		}
		return err
	}
	if engine.afterTurnCommitted != nil {
		return engine.afterTurnCommitted()
	}
	return nil
}