package orchestrator

import (
	"encoding/json"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	workflowrunstore "github.com/tidbcloud/AgentTeam/internal/store/workflow/run"
	planlang "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"
)

func TestExternalizeNestedDependencyValueUsesAPathBoundArtifact(t *testing.T) {
	large := strings.Repeat("x", runlimit.MaxRunPayloadBytes+1024)
	values := runkernel.Values{"dependencies": map[string]any{"join": map[string]any{"value": large}}}
	var drafts []runartifact.Draft
	compacted, changed, err := externalizeValuesAtKind(runartifact.KindJoinAggregate, "wsp_nested", "wri_nested", "join", "act_join", values, runlimit.MaxRunPayloadBytes, func(draft runartifact.Draft) {
		drafts = append(drafts, draft)
	})
	require.NoError(t, err)
	require.True(t, changed)
	require.Len(t, drafts, 1)
	require.LessOrEqual(t, len(mustCanonicalJSON(t, compacted)), runlimit.MaxRunPayloadBytes)
	dependencies := compacted["dependencies"].(map[string]any)
	join := dependencies["join"].(map[string]any)
	ref, marker, err := runartifact.ParseMarker(join["value"])
	require.NoError(t, err)
	require.True(t, marker)
	require.Equal(t, drafts[0].Ref, ref)
	expectedContent, err := json.Marshal(large)
	require.NoError(t, err)
	require.Equal(t, expectedContent, drafts[0].Content)
}

func mustCanonicalJSON(t *testing.T, value any) []byte {
	t.Helper()
	encoded, err := json.Marshal(value)
	require.NoError(t, err)
	return encoded
}

func TestProjectTurnExternalizesOversizedJoinAggregateAndDownstreamCopies(t *testing.T) {
	artifact, diagnostics, err := planlang.Compile([]byte(largeJoinPlan))
	require.NoError(t, err)
	require.False(t, planlang.HasErrors(diagnostics), "%#v", diagnostics)
	graph, err := runkernel.IndexArtifact(artifact)
	require.NoError(t, err)
	kernel, err := runkernel.NewKernel(graph)
	require.NoError(t, err)
	part := strings.Repeat("x", 2300<<10)
	state, err := kernel.Start(runkernel.Values{"seed": part})
	require.NoError(t, err)
	fork, err := kernel.Step(state, time.Unix(1, 0).UTC())
	require.NoError(t, err)
	join, err := kernel.Step(fork.State, time.Unix(2, 0).UTC())
	require.NoError(t, err)
	require.Greater(t, len(join.State.Activations), 2)

	run := workflowrunstore.Run{
		ID: "wri_large_join", WorkspaceID: "wsp_large_join", Revision: 2,
		Status: workflowrunstore.RunStatusRunning, StartedAt: ptrTime(time.Unix(1, 0).UTC()),
	}
	params, err := ProjectTurn(ProjectTurnInput{
		Run: run, Fence: workflowrunstore.TurnFence{WorkspaceID: run.WorkspaceID, RunID: run.ID, LeaseToken: "lease"},
		Graph: graph, Transition: join, SourceKind: "manual", Now: time.Unix(2, 0).UTC(),
	})
	require.NoError(t, err)
	require.NotEmpty(t, params.Artifacts)
	var draft runartifact.Draft
	oversized := 0
	for _, candidate := range params.Artifacts {
		if candidate.Ref.CanonicalByteLength > int64(runlimit.MaxRunPayloadBytes) {
			draft = candidate
			oversized++
		}
	}
	require.Equal(t, 1, oversized, "the closing Join aggregate is persisted exactly once")
	require.Greater(t, draft.Ref.CanonicalByteLength, int64(runlimit.MaxRunPayloadBytes))
	require.Less(t, len(params.StateJSON), runlimit.MaxRunStateBytes)
	require.Contains(t, string(params.StateJSON), runartifact.MarkerKey)
	for _, activation := range params.Activations {
		for _, payload := range [][]byte{activation.Value.InputJSON, activation.Value.OutputJSON} {
			if len(payload) != 0 {
				require.LessOrEqual(t, len(payload), workflowrunstore.MaxRunPayloadBytes)
			}
		}
	}
	var marker map[string]any
	require.NoError(t, json.Unmarshal(params.StateJSON, &marker))
	require.NotNil(t, marker)
}

const largeJoinPlan = `apiVersion: loop.plan/v1
kind: PlanVersion
input:
  seed: {type: string, required: true}
output:
  left: {type: string, required: true}
  right: {type: string, required: true}
graph:
  entry:
    nodeId: fanout
    port: in
    input:
      seed: {ref: plan.input.seed}
  nodes:
  - id: fanout
    type: fork
    input:
      seed: {type: string, required: true}
    fork:
      branches: [left, right]
  - id: gather
    type: join
    input:
      value: {type: string, required: true}
    join:
      forkId: fanout
      mode: all
      result:
        left: {type: string, required: true}
        right: {type: string, required: true}
      merge:
        left: {ref: branches.left.value}
        right: {ref: branches.right.value}
  - id: success
    type: end
    end: {outcome: succeeded}
  - id: failure
    type: end
    end: {outcome: failed, code: join_failed}
  edges:
  - id: left
    source: {nodeId: fanout, port: left}
    target: {nodeId: gather, port: left}
    bindings:
      value: {ref: input.seed}
  - id: right
    source: {nodeId: fanout, port: right}
    target: {nodeId: gather, port: right}
    bindings:
      value: {ref: input.seed}
  - id: joined
    source: {nodeId: gather, port: completed}
    target: {nodeId: success, port: in}
    bindings:
      left: {ref: output.left}
      right: {ref: output.right}
  - id: join_failed
    source: {nodeId: gather, port: failed}
    target: {nodeId: failure, port: in}
`