package orchestrator

import (
	"encoding/json"
	"fmt"
	"strings"

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

// externalizeTransitionArtifacts converts only values that would cross the
// ordinary per-payload boundary into immutable references. The reducer keeps
// its complete in-memory value while preparing the turn; after this function
// returns every persisted copy is compact. Producer identity is frozen from
// the graph activation that created the value, so replay produces the same
// artifact ID and digest.
func externalizeTransitionArtifacts(input *ProjectTurnInput) ([]runartifact.Draft, error) {
	if input == nil || input.Transition.State == nil {
		return nil, nil
	}
	drafts := make([]runartifact.Draft, 0, 4)
	seen := make(map[string]struct{})
	add := func(draft runartifact.Draft) {
		if _, ok := seen[draft.Ref.ArtifactID]; ok {
			return
		}
		seen[draft.Ref.ArtifactID] = struct{}{}
		drafts = append(drafts, draft)
	}
	producer := func(activation runkernel.Activation) (string, string, error) {
		if input.Graph == nil {
			return "", "", fmt.Errorf("workflow artifact producer graph is missing")
		}
		if _, ok := input.Graph.Node(activation.NodeID); !ok {
			return "", "", fmt.Errorf("workflow artifact producer node %q is missing", activation.NodeID)
		}
		return activation.NodeID, activation.ID, nil
	}
	for index := range input.Transition.State.Activations {
		activation := &input.Transition.State.Activations[index]
		if input.Graph == nil {
			return nil, fmt.Errorf("workflow artifact producer graph is missing")
		}
		activationNode, ok := input.Graph.Node(activation.NodeID)
		if !ok {
			return nil, fmt.Errorf("workflow artifact producer node %q is missing", activation.NodeID)
		}
		nodeID, producerActivationID, err := producer(*activation)
		if err != nil {
			return nil, err
		}
		inputFromJoin := false
		if activation.ParentActivationID != "" {
			for parentIndex := range input.Transition.State.Activations {
				parent := input.Transition.State.Activations[parentIndex]
				if parent.ID != activation.ParentActivationID {
					continue
				}
				if parentNode, parentActivation, parentErr := producer(parent); parentErr == nil {
					parentCompiled, _ := input.Graph.Node(parent.NodeID)
					inputFromJoin = parentCompiled != nil && isJoinOpcode(parentCompiled.Opcode)
					// A downstream activation input is a continuation of its parent
					// output. Bind both copies to that single producer identity.
					nodeID, producerActivationID = parentNode, parentActivation
				}
				break
			}
		}
		if inputFromJoin {
			activation.Input, _, err = externalizeValues(input.Run.WorkspaceID, input.Run.ID, nodeID, producerActivationID, activation.Input, add)
			if err != nil {
				return nil, fmt.Errorf("externalize activation %q input: %w", activation.ID, err)
			}
		}
		if isJoinOpcode(activationNode.Opcode) || inputFromJoin {
			outputNodeID, outputActivationID := activation.NodeID, activation.ID
			if inputFromJoin {
				outputNodeID, outputActivationID = nodeID, producerActivationID
			}
			activation.Output, _, err = externalizeValues(input.Run.WorkspaceID, input.Run.ID, outputNodeID, outputActivationID, activation.Output, add)
			if err != nil {
				return nil, fmt.Errorf("externalize activation %q output: %w", activation.ID, err)
			}
		}
	}
	if len(input.Transition.State.Output) != 0 {
		encoded, encodeErr := runartifact.CanonicalJSON(input.Transition.State.Output)
		if encodeErr != nil {
			return nil, fmt.Errorf("inspect Run output artifact: %w", encodeErr)
		}
		if len(encoded) > runlimit.MaxRunPayloadBytes {
			for _, draft := range drafts {
				if string(draft.Content) == string(encoded) {
					input.Transition.State.Output = runkernel.Values{runartifact.MarkerKey: draft.Ref}
					break
				}
			}
		}
	}
	if err := compactStateForStorage(input, add); err != nil {
		return nil, err
	}
	for index := range input.Transition.Intents {
		intent := &input.Transition.Intents[index]
		active := false
		for _, activation := range input.Transition.State.Activations {
			if activation.ID == intent.ActivationID {
				active = activation.Status == runkernel.ActivationWaiting
				break
			}
		}
		if !active || intent.Kind != runkernel.IntentSendMessage || strings.TrimSpace(intent.Text) == "" {
			continue
		}
		encoded, err := json.Marshal(intent.Text)
		if err != nil {
			return nil, fmt.Errorf("encode effect content: %w", err)
		}
		if len(encoded) <= runlimit.MaxRunPayloadBytes {
			continue
		}
		draft, err := runartifact.NewDraftForKind(runartifact.KindEffectText, input.Run.WorkspaceID, input.Run.ID, intent.NodeID, intent.ActivationID, encoded)
		if err != nil {
			return nil, err
		}
		add(draft)
		intent.TextRef = &draft.Ref
		intent.Text = ""
	}
	return drafts, nil
}

func isJoinOpcode(opcode plan.ArtifactOpcode) bool {
	return opcode == plan.OpcodeJoinAll || opcode == plan.OpcodeJoinAny
}

func compactStateForStorage(input *ProjectTurnInput, add func(runartifact.Draft)) error {
	if input == nil || input.Transition.State == nil {
		return nil
	}
	encoded, err := json.Marshal(input.Transition.State)
	if err != nil {
		return fmt.Errorf("measure workflow Run state storage size: %w", err)
	}
	if len(encoded) <= runlimit.MaxRunStateStorageBytes {
		return nil
	}
	state := input.Transition.State
	hasAggregateScope := false
	for _, scope := range state.Scopes {
		if scope.Parallel != nil || scope.FanOut != nil {
			hasAggregateScope = true
			break
		}
	}
	// Preserve the existing deterministic fatal path for an unrelated giant
	// state. Only a graph aggregate that has an explicit Join/Fan-out scope may
	// opt into storage compaction.
	if !hasAggregateScope {
		return nil
	}
	producerNode, producerActivation := "run", "run_input"
	if len(state.Activations) > 0 {
		producerNode, producerActivation = state.Activations[0].NodeID, state.Activations[0].ID
	}
	externalize := func(values runkernel.Values, nodeID, activationID string) (runkernel.Values, error) {
		compacted, _, err := externalizeValuesAtKind(runartifact.KindReducerValue, input.Run.WorkspaceID, input.Run.ID, nodeID, activationID, values, runlimit.MaxRunStateCompactionValueBytes, add)
		return compacted, err
	}
	state.Input, err = externalize(state.Input, producerNode, producerActivation)
	if err != nil {
		return err
	}
	state.Variables, err = externalize(state.Variables, producerNode, producerActivation)
	if err != nil {
		return err
	}
	state.Output, err = externalize(state.Output, producerNode, producerActivation)
	if err != nil {
		return err
	}
	for index := range state.Activations {
		activation := &state.Activations[index]
		activation.Input, err = externalize(activation.Input, activation.NodeID, activation.ID)
		if err != nil {
			return err
		}
		activation.Output, err = externalize(activation.Output, activation.NodeID, activation.ID)
		if err != nil {
			return err
		}
	}
	for scopeIndex := range state.Scopes {
		scope := &state.Scopes[scopeIndex]
		if scope.Parallel != nil {
			for branchIndex := range scope.Parallel.Branches {
				branch := &scope.Parallel.Branches[branchIndex]
				branch.Value, err = externalize(branch.Value, producerNode, producerActivation)
				if err != nil {
					return err
				}
			}
		}
		if scope.Loop != nil {
			scope.Loop.State, err = externalize(scope.Loop.State, producerNode, producerActivation)
			if err != nil {
				return err
			}
		}
		if scope.FanOut != nil {
			for itemIndex := range scope.FanOut.Items {
				item := &scope.FanOut.Items[itemIndex]
				item.Value, err = externalize(item.Value, producerNode, producerActivation)
				if err != nil {
					return err
				}
			}
		}
	}
	encoded, err = json.Marshal(state)
	if err != nil {
		return fmt.Errorf("measure compacted workflow Run state: %w", err)
	}
	if len(encoded) > runlimit.MaxRunStateStorageBytes {
		return fmt.Errorf("workflow Run state remains above storage row budget (%d > %d)", len(encoded), runlimit.MaxRunStateStorageBytes)
	}
	return nil
}

func externalizeValues(workspaceID, runID, nodeID, activationID string, values runkernel.Values, add func(runartifact.Draft)) (runkernel.Values, bool, error) {
	return externalizeValuesAtKind(runartifact.KindJoinAggregate, workspaceID, runID, nodeID, activationID, values, runlimit.MaxRunPayloadBytes, add)
}

func externalizeValuesAtKind(kind, workspaceID, runID, nodeID, activationID string, values runkernel.Values, threshold int, add func(runartifact.Draft)) (runkernel.Values, bool, error) {
	if values == nil {
		return values, false, nil
	}
	if _, marker, err := runartifact.ParseMarker(map[string]any(values)); marker {
		if err != nil {
			return nil, false, err
		}
		return values, false, nil
	}
	encoded, err := runartifact.CanonicalJSON(values)
	if err != nil {
		return nil, false, err
	}
	if len(encoded) <= threshold {
		return values, false, nil
	}
	compacted, changed, err := externalizeNestedValue(kind, workspaceID, runID, nodeID, activationID, values, threshold, add)
	if err != nil {
		return nil, false, err
	}
	if changed {
		switch compactedValues := compacted.(type) {
		case map[string]any:
			return runkernel.Values(compactedValues), true, nil
		case runkernel.Values:
			return compactedValues, true, nil
		}
	}
	draft, err := runartifact.NewDraftForKind(kind, workspaceID, runID, nodeID, activationID, encoded)
	if err != nil {
		return nil, false, err
	}
	add(draft)
	return runkernel.Values{runartifact.MarkerKey: draft.Ref}, true, nil
}

// externalizeNestedValue keeps large dependency/business envelopes compact. A
// scalar or leaf object that exceeds the consumer boundary becomes its own
// artifact; enclosing maps/arrays retain their ordinary shape and markers.
func externalizeNestedValue(kind, workspaceID, runID, nodeID, activationID string, value any, threshold int, add func(runartifact.Draft)) (any, bool, error) {
	encoded, err := runartifact.CanonicalJSON(value)
	if err != nil {
		return nil, false, err
	}
	if len(encoded) <= threshold {
		return value, false, nil
	}
	changed := false
	switch typed := value.(type) {
	case map[string]any:
		for key, child := range typed {
			compacted, childChanged, childErr := externalizeNestedValue(kind, workspaceID, runID, nodeID, activationID, child, threshold, add)
			if childErr != nil {
				return nil, false, childErr
			}
			typed[key], changed = compacted, changed || childChanged
		}
	case []any:
		for index, child := range typed {
			compacted, childChanged, childErr := externalizeNestedValue(kind, workspaceID, runID, nodeID, activationID, child, threshold, add)
			if childErr != nil {
				return nil, false, childErr
			}
			typed[index], changed = compacted, changed || childChanged
		}
	case runkernel.Values:
		for key, child := range typed {
			compacted, childChanged, childErr := externalizeNestedValue(kind, workspaceID, runID, nodeID, activationID, child, threshold, add)
			if childErr != nil {
				return nil, false, childErr
			}
			typed[key], changed = compacted, changed || childChanged
		}
	}
	if changed {
		compacted, compactErr := runartifact.CanonicalJSON(value)
		if compactErr != nil {
			return nil, false, compactErr
		}
		if len(compacted) <= threshold {
			return value, true, nil
		}
		draft, draftErr := runartifact.NewDraftForKind(kind, workspaceID, runID, nodeID, activationID, compacted)
		if draftErr != nil {
			return nil, false, draftErr
		}
		add(draft)
		return runartifact.Marker(draft.Ref), true, nil
	}
	draft, err := runartifact.NewDraftForKind(kind, workspaceID, runID, nodeID, activationID, encoded)
	if err != nil {
		return nil, false, err
	}
	add(draft)
	return runartifact.Marker(draft.Ref), true, nil
}