package orchestrator

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

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

func validateMaterial(run runstore.Run, material Material) (*runkernel.Graph, error) {
	if material.WorkspaceID != run.WorkspaceID || material.PlanID != run.PlanID || material.PlanVersionID != run.PlanVersionID {
		return nil, fmt.Errorf("exact PlanVersion material does not match Run")
	}
	artifact, err := plan.DecodeArtifact(material.ArtifactJSON)
	if err != nil {
		return nil, err
	}
	if artifact.CompilerVersion != strings.TrimSpace(material.CompilerVersion) || artifact.Digest != strings.TrimSpace(material.Digest) {
		return nil, fmt.Errorf("exact PlanVersion artifact metadata does not match")
	}
	return runkernel.IndexArtifact(artifact)
}

func decodeRunState(payload []byte) (*runkernel.State, bool, error) {
	trimmed := bytes.TrimSpace(payload)
	if len(trimmed) == 0 || bytes.Equal(trimmed, []byte(`{}`)) {
		return &runkernel.State{}, false, nil
	}
	decoder := json.NewDecoder(bytes.NewReader(trimmed))
	decoder.DisallowUnknownFields()
	decoder.UseNumber()
	var state runkernel.State
	if err := decoder.Decode(&state); err != nil {
		return nil, false, err
	}
	if err := requireJSONEOF(decoder); err != nil {
		return nil, false, fmt.Errorf("workflow Run state: %w", err)
	}
	if state.Status == "" || state.Activations == nil || state.Scopes == nil || state.Input == nil || state.Variables == nil {
		return nil, false, fmt.Errorf("workflow Run state is incomplete")
	}
	return &state, true, nil
}

func decodeValues(payload []byte, target *runkernel.Values) error {
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.UseNumber()
	if err := decoder.Decode(target); err != nil {
		return err
	}
	if *target == nil {
		return fmt.Errorf("workflow Run input must be an object")
	}
	if err := requireJSONEOF(decoder); err != nil {
		return fmt.Errorf("workflow Run input: %w", err)
	}
	return nil
}

func validateRestoredState(state *runkernel.State, graph *runkernel.Graph) error {
	if len(state.Activations) > plan.MaxRunActivations || state.NextActivation > uint64(plan.MaxRunActivations) {
		return fmt.Errorf("workflow Run state exceeds the %d activation limit", plan.MaxRunActivations)
	}
	seen := make(map[string]struct{}, len(state.Activations))
	activationSequence := make(map[string]uint64, len(state.Activations))
	var maxActivation uint64
	for _, activation := range state.Activations {
		if activation.ID == "" || activation.Sequence == 0 || activation.NodeID == "" {
			return fmt.Errorf("workflow Run activation identity is incomplete")
		}
		if _, duplicate := seen[activation.ID]; duplicate {
			return fmt.Errorf("workflow Run activation %q is duplicated", activation.ID)
		}
		seen[activation.ID] = struct{}{}
		activationSequence[activation.ID] = activation.Sequence
		if _, ok := graph.Node(activation.NodeID); !ok {
			return fmt.Errorf("workflow Run activation %q references unknown node", activation.ID)
		}
		if activation.Sequence > maxActivation {
			maxActivation = activation.Sequence
		}
	}
	if state.NextActivation < maxActivation {
		return fmt.Errorf("workflow Run activation cursor is behind persisted state")
	}
	for _, activation := range state.Activations {
		if activation.ParentActivationID == "" {
			continue
		}
		parentSequence, found := activationSequence[activation.ParentActivationID]
		if !found || parentSequence >= activation.Sequence {
			return fmt.Errorf("workflow Run activation %q has an invalid parent", activation.ID)
		}
	}
	seenScopes := make(map[string]struct{}, len(state.Scopes))
	var maxScope uint64
	for _, scope := range state.Scopes {
		if scope.ID == "" || scope.Sequence == 0 || scope.OwnerID == "" {
			return fmt.Errorf("workflow Run scope identity is incomplete")
		}
		if _, duplicate := seenScopes[scope.ID]; duplicate {
			return fmt.Errorf("workflow Run scope %q is duplicated", scope.ID)
		}
		seenScopes[scope.ID] = struct{}{}
		if scope.Sequence > maxScope {
			maxScope = scope.Sequence
		}
	}
	if state.NextScope < maxScope {
		return fmt.Errorf("workflow Run scope cursor is behind persisted state")
	}
	return nil
}

func stableEffects(values []runstore.Effect) []runstore.Effect {
	result := append([]runstore.Effect(nil), values...)
	sort.SliceStable(result, func(i, j int) bool {
		left, right := result[i].CompletedAt, result[j].CompletedAt
		if left != nil && right != nil && !left.Equal(*right) {
			return left.Before(*right)
		}
		if left != nil && right == nil {
			return true
		}
		if left == nil && right != nil {
			return false
		}
		return result[i].ID < result[j].ID
	})
	return result
}

// availableTaskWatchResponderObservation exposes the selected resolved Agent
// from a committed callback message receipt without treating the callback as
// complete. Task Watch was already changed by the normal message transaction;
// the Run must freeze the same identity before another projected callback can
// advance the causal frontier.
func availableTaskWatchResponderObservation(
	state *runkernel.State,
	effect runstore.Effect,
) (activationID string, responderAgentID string, available bool, err error) {
	if state == nil || (effect.Status != runstore.EffectStatusDelivered &&
		effect.Status != runstore.EffectStatusAwaiting &&
		effect.Status != runstore.EffectStatusCompleted &&
		effect.Status != runstore.EffectStatusFailed &&
		effect.Status != runstore.EffectStatusCanceled) {
		return "", "", false, nil
	}
	request, err := DecodeEffectRequest(effect.RequestJSON)
	if err != nil {
		return "", "", false, err
	}
	if request.Intent.Kind != runkernel.IntentSendMessage || !request.Intent.Callback ||
		request.Intent.TaskWatchHandoff == nil {
		return "", "", false, nil
	}
	var activation *runkernel.Activation
	for index := range state.Activations {
		if state.Activations[index].ID == request.Intent.ActivationID {
			activation = &state.Activations[index]
			break
		}
	}
	if activation == nil || activation.CallbackTaskWatch == nil {
		return "", "", false, fmt.Errorf(
			"Task Watch callback effect references ineligible activation %q",
			request.Intent.ActivationID,
		)
	}
	if (effect.Status == runstore.EffectStatusCompleted ||
		effect.Status == runstore.EffectStatusFailed ||
		effect.Status == runstore.EffectStatusCanceled) &&
		len(bytes.TrimSpace(effect.ResultJSON)) == 0 {
		return "", "", false, nil
	}
	result, err := DecodeEffectResult(effect.ResultJSON)
	if err != nil {
		return "", "", false, err
	}
	if strings.TrimSpace(result.MessageID) == "" {
		if effect.Status == runstore.EffectStatusCompleted ||
			effect.Status == runstore.EffectStatusFailed ||
			effect.Status == runstore.EffectStatusCanceled {
			// A callback may reach a definitive pre-delivery timeout/failure
			// without ever committing a Task message. That terminal result has
			// no responder to observe; the completion reducer restores the
			// provisional predecessor frontier instead.
			return "", "", false, nil
		}
		return "", "", false, fmt.Errorf("Task Watch callback receipt is missing messageId")
	}
	if result.DeliveryScope == nil || *result.DeliveryScope != request.Intent.TaskWatchHandoff.TaskScope {
		return "", "", false, fmt.Errorf("Task Watch callback receipt deliveryScope does not match its handoff")
	}
	responderAgentID, err = ResolveCallbackResponderAuthority(
		true,
		result.CallbackResponderAgentID,
		result.WakeAgentIDs,
	)
	if err != nil {
		return "", "", false, err
	}
	if existing := strings.TrimSpace(activation.CallbackTaskWatch.ResponderAgentID); existing != "" {
		if existing != responderAgentID {
			return "", "", false, fmt.Errorf(
				"Task Watch callback activation %q receipt responder changed from %q to %q",
				activation.ID,
				existing,
				responderAgentID,
			)
		}
		return "", "", false, nil
	}
	return activation.ID, responderAgentID, true, nil
}

// availableTaskWatchPreDeliveryRestoration recognizes only durable evidence
// that no Task message transaction committed. Ambiguous or crossed dispatches
// remain behind hasUnobservedTaskWatchDispatch until they converge to a
// receipt-bearing result or a definitive no-receipt terminal result. Authored
// callback deadlines no longer produce a definitive no-receipt state, so a
// pending first attempt is never restored on time alone.
func availableTaskWatchPreDeliveryRestoration(
	state *runkernel.State,
	effect runstore.Effect,
) (activationID string, available bool, err error) {
	if state == nil {
		return "", false, fmt.Errorf("Task Watch restoration requires Run state")
	}
	switch effect.Status {
	case runstore.EffectStatusCompleted,
		runstore.EffectStatusFailed:
	default:
		return "", false, nil
	}
	request, err := DecodeEffectRequest(effect.RequestJSON)
	if err != nil {
		return "", false, err
	}
	if request.Intent.Kind != runkernel.IntentSendMessage || !request.Intent.Callback ||
		request.Intent.TaskWatchHandoff == nil {
		return "", false, nil
	}
	var activation *runkernel.Activation
	for index := range state.Activations {
		if state.Activations[index].ID == request.Intent.ActivationID {
			activation = &state.Activations[index]
			break
		}
	}
	if activation == nil || activation.CallbackTaskWatch == nil {
		return "", false, fmt.Errorf(
			"Task Watch callback effect references ineligible activation %q",
			request.Intent.ActivationID,
		)
	}
	if activation.Status != runkernel.ActivationWaiting ||
		strings.TrimSpace(activation.CallbackTaskWatch.ResponderAgentID) != "" {
		return "", false, nil
	}
	projectionOutstanding := false
	for _, frontierActivationID := range activation.TaskWatchFrontier {
		if frontierActivationID == activation.ID {
			projectionOutstanding = true
			break
		}
	}
	if !projectionOutstanding {
		return "", false, nil
	}

	if len(bytes.TrimSpace(effect.ResultJSON)) == 0 {
		return activation.ID, true, nil
	}
	result, err := DecodeEffectResult(effect.ResultJSON)
	if err != nil {
		return "", false, err
	}
	if strings.TrimSpace(result.MessageID) != "" {
		return "", false, fmt.Errorf(
			"receipt-bearing Task Watch callback %q was not observed before restoration",
			activation.ID,
		)
	}
	return activation.ID, true, nil
}

// hasUnobservedTaskWatchDispatch fences only the short delivery window of a
// fixed-Task callback. Once the effect receipt is awaiting and its responder
// has been observed, Agent work and sibling branches continue in parallel.
// Authored callback deadlines no longer release this fence; a pending first
// attempt for a waiting activation stays fenced until it crosses dispatch or
// reaches a definitive terminal result.
func hasUnobservedTaskWatchDispatch(
	state *runkernel.State,
	effects []runstore.Effect,
) (bool, error) {
	if state == nil {
		return false, fmt.Errorf("Task Watch dispatch barrier requires Run state")
	}
	for _, effect := range effects {
		switch effect.Status {
		case runstore.EffectStatusPending,
			runstore.EffectStatusDispatching,
			runstore.EffectStatusDelivered,
			runstore.EffectStatusAwaiting:
		default:
			continue
		}
		request, err := DecodeEffectRequest(effect.RequestJSON)
		if err != nil {
			return false, err
		}
		if request.Intent.Kind != runkernel.IntentSendMessage || !request.Intent.Callback ||
			request.Intent.TaskWatchHandoff == nil {
			continue
		}
		found := false
		for index := range state.Activations {
			activation := &state.Activations[index]
			if activation.ID != request.Intent.ActivationID {
				continue
			}
			found = true
			if activation.CallbackTaskWatch == nil {
				return false, fmt.Errorf(
					"Task Watch callback effect references ineligible activation %q",
					activation.ID,
				)
			}
			if strings.TrimSpace(activation.CallbackTaskWatch.ResponderAgentID) == "" {
				switch effect.Status {
				case runstore.EffectStatusDispatching,
					runstore.EffectStatusDelivered,
					runstore.EffectStatusAwaiting:
					// Dispatch has crossed, or may already have crossed, the
					// external Task transaction. Even a canceled loser must wait
					// for its authoritative receipt/failure before a join successor
					// can compute a complete handoff.
					return true, nil
				case runstore.EffectStatusPending:
					// A waiting first attempt still needs delivery. A retried
					// pending effect is conservatively treated as crossed because
					// the prior writer error may have followed an external commit.
					if activation.Status == runkernel.ActivationWaiting || effect.AttemptCount > 0 {
						return true, nil
					}
				}
			}
			break
		}
		if !found {
			return false, fmt.Errorf(
				"Task Watch callback effect references missing activation %q",
				request.Intent.ActivationID,
			)
		}
	}
	return false, nil
}

func transitionHasTaskWatchDispatch(transition runkernel.Transition) bool {
	for _, intent := range transition.Intents {
		if intent.Kind == runkernel.IntentSendMessage && intent.Callback && intent.TaskWatchHandoff != nil {
			return true
		}
	}
	return false
}

func availableEffectCompletion(state *runkernel.State, effect runstore.Effect) (runkernel.Completion, bool, error) {
	request, err := DecodeEffectRequest(effect.RequestJSON)
	if err != nil {
		return runkernel.Completion{}, false, err
	}
	waiting := false
	for _, activation := range state.Activations {
		if activation.ID == request.Intent.ActivationID {
			waiting = activation.Status == runkernel.ActivationWaiting
			break
		}
	}
	if !waiting {
		return runkernel.Completion{}, false, nil
	}
	switch effect.Status {
	case runstore.EffectStatusCompleted:
		result, err := DecodeEffectResult(effect.ResultJSON)
		if err != nil {
			return runkernel.Completion{}, false, err
		}
		if request.Intent.Callback && strings.TrimSpace(result.MessageID) != "" {
			result.CallbackResponderAgentID, err = ResolveCallbackResponderAuthority(
				true,
				result.CallbackResponderAgentID,
				result.WakeAgentIDs,
			)
			if err != nil {
				return runkernel.Completion{}, false, err
			}
		}
		completion, err := result.Completion(request.Intent.ActivationID)
		return completion, err == nil, err
	case runstore.EffectStatusFailed:
		result := EffectResult{Port: "failed", Output: runkernel.Values{}, FailureCode: effect.ErrorCode}
		if len(effect.ResultJSON) != 0 {
			result, err = DecodeEffectResult(effect.ResultJSON)
			if err != nil {
				return runkernel.Completion{}, false, err
			}
			if result.Port == "" {
				result.Port = "failed"
			}
		}
		completion, err := result.Completion(request.Intent.ActivationID)
		return completion, err == nil, err
	default:
		// Pending, delivered, and awaiting callback effects are resolved only by
		// a committed receipt or an infrastructure-level delivery failure. The
		// engine no longer enforces an authored callback deadline, so these live
		// statuses never reduce to a timeout completion.
		return runkernel.Completion{}, false, nil
	}
}

func hasReducibleWork(state *runkernel.State, effects []runstore.Effect) (bool, error) {
	if len(runkernel.StableReadyIDs(state)) != 0 {
		return true, nil
	}
	for _, effect := range stableEffects(effects) {
		_, available, err := availableEffectCompletion(state, effect)
		if err != nil {
			return false, err
		}
		if available {
			return true, nil
		}
	}
	return false, nil
}

func mergeTransition(aggregate *runkernel.Transition, next runkernel.Transition) {
	aggregate.State = next.State
	aggregate.Intents = append(aggregate.Intents, next.Intents...)
	aggregate.Events = append(aggregate.Events, next.Events...)
}

func (engine *Engine) listAllActivations(ctx context.Context, run runstore.Run) ([]runstore.Activation, error) {
	values := make([]runstore.Activation, 0)
	after := ""
	for len(values) <= maxRuntimeChildren {
		limit := collectionProbeLimit(len(values))
		page, err := engine.repository.ListActivations(ctx, runstore.ListChildrenParams{
			WorkspaceID: run.WorkspaceID, RunID: run.ID, AfterID: after, Limit: limit,
		})
		if err != nil {
			return nil, err
		}
		values = append(values, page...)
		if len(values) > maxRuntimeChildren {
			return nil, fmt.Errorf("%w: activations exceed %d", errRuntimeCollectionLimit, maxRuntimeChildren)
		}
		if len(page) < limit {
			return values, nil
		}
		after = page[len(page)-1].ID
	}
	return nil, fmt.Errorf("%w: activations exceed %d", errRuntimeCollectionLimit, maxRuntimeChildren)
}

func (engine *Engine) listAllEffects(ctx context.Context, run runstore.Run) ([]runstore.Effect, error) {
	values := make([]runstore.Effect, 0)
	after := ""
	for len(values) <= maxRuntimeChildren {
		limit := collectionProbeLimit(len(values))
		page, err := engine.repository.ListEffects(ctx, runstore.ListChildrenParams{
			WorkspaceID: run.WorkspaceID, RunID: run.ID, AfterID: after, Limit: limit,
		})
		if err != nil {
			return nil, err
		}
		values = append(values, page...)
		if len(values) > maxRuntimeChildren {
			return nil, fmt.Errorf("%w: effects exceed %d", errRuntimeCollectionLimit, maxRuntimeChildren)
		}
		if len(page) < limit {
			return values, nil
		}
		after = page[len(page)-1].ID
	}
	return nil, fmt.Errorf("%w: effects exceed %d", errRuntimeCollectionLimit, maxRuntimeChildren)
}

func collectionProbeLimit(current int) int {
	remaining := maxRuntimeChildren + 1 - current
	if remaining < runtimeCollectionPage {
		return remaining
	}
	return runtimeCollectionPage
}

func (engine *Engine) saveUnrecoverableFailure(ctx context.Context, claim runstore.ClaimedRun, code string, cause error, now time.Time) error {
	payload, errorJSON, err := unrecoverableFailureSnapshot(claim.Run, code, cause)
	if err != nil {
		return err
	}
	_, err = engine.repository.FailClaimedRun(ctx, runstore.FailClaimedRunParams{
		Fence: claim.Fence, ExpectedRevision: claim.Run.Revision,
		StateJSON: payload, ErrorCode: code, ErrorJSON: errorJSON, Now: now,
	})
	return err
}

// unrecoverableFailureSnapshot always returns a complete, decodable reducer
// snapshot. Normal initialized state is preserved for post-mortem graph/UI
// inspection. If that state itself is poisoned, a minimal terminal snapshot
// is built from the immutable Invocation input; relational children are
// canceled by FailClaimedRun in the same fenced transaction.
func unrecoverableFailureSnapshot(run runstore.Run, code string, cause error) (json.RawMessage, json.RawMessage, error) {
	if strings.TrimSpace(code) == "" || cause == nil {
		return nil, nil, fmt.Errorf("workflow Run deterministic failure requires code and cause")
	}
	// The persisted failure carries only the stable code plus bounded,
	// public-safe problem/hint derived from the code; the arbitrary cause never
	// crosses the storage boundary.
	failure := &runkernel.Failure{Code: strings.TrimSpace(code), Stage: "material", Problem: runkernel.FailureProblem(code), Hint: runkernel.FailureHint(code)}
	state, initialized, decodeErr := decodeRunState(run.StateJSON)
	if decodeErr != nil || !initialized {
		state = minimalFailureState(run, failure)
	} else {
		state = state.Clone()
		state.Status = runkernel.StatusFailed
		state.Failure = failure
		state.Output = nil
		for index := range state.Activations {
			if state.Activations[index].Status == runkernel.ActivationReady || state.Activations[index].Status == runkernel.ActivationWaiting {
				state.Activations[index].Status = runkernel.ActivationCanceled
				state.Activations[index].Failure = failure
			}
		}
		for index := range state.Scopes {
			if state.Scopes[index].Status == runkernel.ScopeActive {
				state.Scopes[index].Status = runkernel.ScopeCanceled
			}
		}
	}
	payload, err := json.Marshal(state)
	if err != nil {
		return nil, nil, fmt.Errorf("encode deterministic workflow Run failure state: %w", err)
	}
	payload, err = runstore.DeterministicStateJSON(payload)
	if err != nil {
		state = minimalFailureState(run, failure)
		payload, err = json.Marshal(state)
		if err != nil {
			return nil, nil, fmt.Errorf("encode bounded deterministic workflow Run failure state: %w", err)
		}
		payload, err = runstore.DeterministicStateJSON(payload)
		if err != nil {
			return nil, nil, fmt.Errorf("validate bounded deterministic workflow Run failure state: %w", err)
		}
	}
	errorJSON, err := runFailureFromKernel(failure, "material")
	if err != nil {
		return nil, nil, fmt.Errorf("encode deterministic workflow Run failure: %w", err)
	}
	return payload, errorJSON, nil
}

func minimalFailureState(run runstore.Run, failure *runkernel.Failure) *runkernel.State {
	input := runkernel.Values{}
	if err := decodeValues(run.InputJSON, &input); err != nil {
		input = runkernel.Values{}
	}
	return &runkernel.State{
		Status: runkernel.StatusFailed, Input: input, Variables: runkernel.Values{},
		Activations: []runkernel.Activation{}, Scopes: []runkernel.Scope{}, Failure: failure,
	}
}

func (engine *Engine) saveStateFailure(
	ctx context.Context,
	claim runstore.ClaimedRun,
	graph *runkernel.Graph,
	state *runkernel.State,
	activations []runstore.Activation,
	effects []runstore.Effect,
	sourceKind, code string,
	cause error,
	now time.Time,
) error {
	failed := state.Clone()
	failed.Status = runkernel.StatusFailed
	failure := &runkernel.Failure{
		Code: code, Stage: "reducer", Problem: runkernel.FailureProblem(code), Hint: runkernel.FailureHint(code),
	}
	failed.Failure = failure
	for index := range failed.Activations {
		if failed.Activations[index].Status == runkernel.ActivationReady || failed.Activations[index].Status == runkernel.ActivationWaiting {
			failed.Activations[index].Status = runkernel.ActivationCanceled
		}
	}
	for index := range failed.Scopes {
		if failed.Scopes[index].Status == runkernel.ScopeActive {
			failed.Scopes[index].Status = runkernel.ScopeCanceled
		}
	}
	transition := runkernel.Transition{State: failed, Events: []runkernel.Event{runkernel.FailureEvent("run_failed", failure, "", "", "")}}
	params, err := ProjectTurn(ProjectTurnInput{
		Run: claim.Run, Fence: claim.Fence, Graph: graph, Transition: transition,
		ExistingActivations: activations, ExistingEffects: effects, SourceKind: sourceKind, Now: now,
	})
	if err != nil {
		if limitCode := runtimeProjectionLimitCode(err); limitCode != "" {
			return engine.saveUnrecoverableFailure(ctx, claim, limitCode, err, now)
		}
		return err
	}
	_, err = engine.repository.SaveTurn(ctx, params)
	return err
}

func runtimeProjectionLimitCode(err error) string {
	switch {
	case errors.Is(err, errRuntimeStateLimit):
		return "run.runtime_state_limit_exceeded"
	case errors.Is(err, errRuntimePayloadLimit):
		return "run.runtime_payload_limit_exceeded"
	case errors.Is(err, runartifact.ErrArtifactBudgetExceeded):
		return "run.artifact_budget_exceeded"
	default:
		return ""
	}
}