package message

import (
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/core/domain"
)

func TestValidateWorkflowContextMetadataRejectsEveryReservedKeyWithoutInternalCapability(t *testing.T) {
	deliveryScope := domain.MessageScopeValue{Kind: domain.MessageScopeKindTask, ScopeID: "tsk_callback_work"}
	metadata := BuildWorkflowContextMetadata(WorkflowContextMetadataParams{
		WorkspaceID: "wsp_1", PlanID: "pln_1", PlanVersionID: "plv_1",
		RunID: "wri_1", EffectID: "wfx_1", NodeID: "ask", SourceKind: "manual",
		DeliveryScope: &deliveryScope, ResponderAgentRef: "calculator-a",
	})
	metadata[WorkflowContextMetadataCallbackExpectedKey] = true
	metadata[WorkflowContextMetadataResponderAgentIDKey] = "agt_calculator_a"
	require.Equal(t, map[string]any{
		"workflow_workspace_id":    "wsp_1",
		"workflow_plan_id":         "pln_1",
		"workflow_plan_version_id": "plv_1",
		"workflow_run_id":          "wri_1",
		"workflow_effect_id":       "wfx_1",
		"workflow_node_id":         "ask",
		"workflow_source_kind":     "manual",
		"workflow_delivery_scope": map[string]any{
			"kind": "task", "scopeId": "tsk_callback_work",
		},
		"workflow_callback_expected":   true,
		"workflow_responder_agent_ref": "calculator-a",
		"workflow_responder_agent_id":  "agt_calculator_a",
	}, metadata)
	require.Len(t, metadata, 11)
	for key := range metadata {
		key := key
		t.Run(key, func(t *testing.T) {
			err := ValidateWorkflowContextMetadata(map[string]any{key: nil}, false)
			require.ErrorContains(t, err, key)
		})
	}
	require.NoError(t, ValidateWorkflowContextMetadata(metadata, true))
	require.NoError(t, ValidateWorkflowContextMetadata(map[string]any{"client_key": "allowed"}, false))
}

func TestValidateWorkflowContextMetadataQuarantinesRetiredCallbackExpiryKey(t *testing.T) {
	// A retired workflow key must be rejected as ordinary inbound metadata even
	// with allowReserved=true, so a stale producer cannot smuggle a frozen
	// callback deadline back into the protocol.
	for _, allowReserved := range []bool{false, true} {
		err := ValidateWorkflowContextMetadata(map[string]any{
			WorkflowContextMetadataCallbackExpiresAtKey: "2030-01-02T03:04:05Z",
		}, allowReserved)
		require.ErrorContains(t, err, "retired", "allowReserved=%v", allowReserved)
		require.ErrorContains(t, err, WorkflowContextMetadataCallbackExpiresAtKey, "allowReserved=%v", allowReserved)
	}

	// The builder must never emit a retired key.
	built := BuildWorkflowContextMetadata(WorkflowContextMetadataParams{
		RunID: "wri_1", CallbackExpected: true,
	})
	require.NotContains(t, built, WorkflowContextMetadataCallbackExpiresAtKey)

	// Active workflow keys still pass with allowReserved=true and are rejected
	// with allowReserved=false.
	active := map[string]any{WorkflowContextMetadataRunIDKey: "wri_1"}
	require.NoError(t, ValidateWorkflowContextMetadata(active, true))
	require.ErrorContains(t, ValidateWorkflowContextMetadata(active, false), "reserved")
}

func TestWorkflowContextMetadataForAgentRecipientStripsRetiredKey(t *testing.T) {
	metadata := BuildWorkflowContextMetadata(WorkflowContextMetadataParams{
		RunID: "wri_1", CallbackExpected: true, ResponderAgentRef: "calculator-a",
	})
	metadata[WorkflowContextMetadataResponderAgentIDKey] = "agt_calculator_a"
	metadata[WorkflowContextMetadataCallbackExpiresAtKey] = "2030-01-02T03:04:05Z"
	metadata["custom"] = "visible"

	selected := WorkflowContextMetadataForAgentRecipient(metadata, "agt_calculator_a")
	// The selected responder keeps the trusted callback authority but must never
	// carry a retired callback-expiry key.
	require.Equal(t, "visible", selected["custom"])
	require.True(t, selected[WorkflowContextMetadataCallbackExpectedKey].(bool))
	require.NotContains(t, selected, WorkflowContextMetadataCallbackExpiresAtKey)

	// Ordinary recipients receive no workflow metadata, including the retired key.
	projected := WorkflowContextMetadataForAgentRecipient(metadata, "agt_default")
	require.Equal(t, "visible", projected["custom"])
	for _, key := range reservedWorkflowContextMetadataKeys {
		require.NotContains(t, projected, key)
	}
	require.NotContains(t, projected, WorkflowContextMetadataCallbackExpiresAtKey)
}

func TestValidateWorkflowCallbackResponderRoutingRequiresTrustedCallbackMetadata(t *testing.T) {
	callback := BuildWorkflowContextMetadata(WorkflowContextMetadataParams{
		RunID: "wri_1", CallbackExpected: true,
	})
	require.NoError(t, ValidateWorkflowCallbackResponderRouting(callback, true, true))
	require.NoError(t, ValidateWorkflowCallbackResponderRouting(nil, false, false))
	require.ErrorContains(t, ValidateWorkflowCallbackResponderRouting(callback, false, true), "trusted workflow context")
	require.ErrorContains(t, ValidateWorkflowCallbackResponderRouting(map[string]any{WorkflowContextMetadataCallbackExpectedKey: false}, true, true), "callbackExpected")
	require.ErrorContains(t, ValidateWorkflowCallbackResponderRouting(map[string]any{WorkflowContextMetadataCallbackExpectedKey: "true"}, true, true), "callbackExpected")
}

func TestValidateWorkflowCallbackResponderAgentRefBindsExactMetadata(t *testing.T) {
	metadata := BuildWorkflowContextMetadata(WorkflowContextMetadataParams{
		RunID: "wri_1", CallbackExpected: true, ResponderAgentRef: "agt_exact",
	})
	require.Equal(t, "agt_exact", metadata[WorkflowContextMetadataResponderAgentRefKey])
	require.NoError(t, ValidateWorkflowCallbackResponderAgentRef(metadata, true, "agt_exact"))
	require.ErrorContains(
		t,
		ValidateWorkflowCallbackResponderAgentRef(metadata, true, "agt_other"),
		"does not match workflow metadata",
	)
	require.ErrorContains(
		t,
		ValidateWorkflowCallbackResponderAgentRef(metadata, false, "agt_exact"),
		"callback authority routing",
	)
	require.ErrorContains(
		t,
		ValidateWorkflowCallbackResponderAgentRef(metadata, true, ""),
		"requires a selected responder Agent",
	)
	require.NoError(t, ValidateWorkflowCallbackResponderAgentRef(nil, true, ""))
}

func TestWorkflowCallbackResponderAgentIDIsDerivedAndExcludedFromDigestMetadata(t *testing.T) {
	t.Parallel()

	metadata := map[string]any{
		WorkflowContextMetadataResponderAgentRefKey: "calculator-a",
		WorkflowContextMetadataResponderAgentIDKey:  "agt_calculator_a",
	}
	agentID, err := WorkflowCallbackResponderAgentID(metadata)
	require.NoError(t, err)
	require.Equal(t, "agt_calculator_a", agentID)

	digestMetadata := MetadataWithoutDerivedWorkflowResponder(metadata)
	require.Equal(t, "calculator-a", digestMetadata[WorkflowContextMetadataResponderAgentRefKey])
	require.NotContains(t, digestMetadata, WorkflowContextMetadataResponderAgentIDKey)
	require.Contains(t, metadata, WorkflowContextMetadataResponderAgentIDKey)

	for _, malformed := range []any{"", " agt_calculator_a", 42, true} {
		_, err := WorkflowCallbackResponderAgentID(map[string]any{
			WorkflowContextMetadataResponderAgentIDKey: malformed,
		})
		require.ErrorContains(t, err, "responder Agent ID is invalid")
	}
}

func TestWorkflowContextMetadataForAgentRecipientKeepsAuthorityOnlyForFrozenResponder(t *testing.T) {
	t.Parallel()

	metadata := BuildWorkflowContextMetadata(WorkflowContextMetadataParams{
		WorkspaceID: "wsp_1", PlanID: "pln_1", PlanVersionID: "plv_1",
		RunID: "wri_1", EffectID: "wfx_1", NodeID: "calculate", SourceKind: "manual",
		CallbackExpected:  true,
		ResponderAgentRef: "calculator-a",
	})
	metadata[WorkflowContextMetadataResponderAgentIDKey] = "agt_calculator_a"
	metadata["custom"] = "visible"

	selected := WorkflowContextMetadataForAgentRecipient(metadata, "agt_calculator_a")
	require.Equal(t, metadata, selected)
	require.True(t, selected[WorkflowContextMetadataCallbackExpectedKey].(bool))

	for _, agentID := range []string{"agt_default", "agt_watcher", "agt_second_mention", ""} {
		projected := WorkflowContextMetadataForAgentRecipient(metadata, agentID)
		require.Equal(t, "visible", projected["custom"])
		for _, key := range reservedWorkflowContextMetadataKeys {
			require.NotContains(t, projected, key, "recipient %q must not receive callback metadata", agentID)
		}
	}
	require.Contains(t, metadata, WorkflowContextMetadataResponderAgentIDKey, "projection must not mutate the durable metadata")

	legacy := map[string]any{
		WorkflowContextMetadataRunIDKey:            "wri_legacy",
		WorkflowContextMetadataCallbackExpectedKey: true,
	}
	require.Equal(t, legacy, WorkflowContextMetadataForAgentRecipient(legacy, "agt_legacy"))

	missingDerivedID := CopyMessageMetadata(metadata)
	delete(missingDerivedID, WorkflowContextMetadataResponderAgentIDKey)
	require.NotContains(
		t,
		WorkflowContextMetadataForAgentRecipient(missingDerivedID, "agt_calculator_a"),
		WorkflowContextMetadataCallbackExpectedKey,
		"structured callback without its server-derived authority must fail closed",
	)

	malformed := CopyMessageMetadata(metadata)
	malformed[WorkflowContextMetadataResponderAgentIDKey] = " agt_calculator_a"
	require.NotContains(
		t,
		WorkflowContextMetadataForAgentRecipient(malformed, "agt_calculator_a"),
		WorkflowContextMetadataCallbackExpectedKey,
		"malformed frozen authority must fail closed",
	)
}

func TestWorkflowContextMetadataForAgentRecipientKeepsEditContextOnlyForSelectedAgent(t *testing.T) {
	t.Parallel()

	metadata := map[string]any{
		WorkflowAgentEditMetadataKey: map[string]any{
			"kind":             "request",
			"responderAgentId": "agt_selected",
			"diagnostics":      []any{map[string]any{"code": "warning"}},
		},
		"custom": "visible",
	}
	require.Equal(t, metadata, WorkflowContextMetadataForAgentRecipient(metadata, "agt_selected"))

	for name, edit := range map[string]any{
		"other recipient": metadata[WorkflowAgentEditMetadataKey],
		"malformed value": "request",
		"missing responder": map[string]any{
			"kind": "request",
		},
		"malformed responder": map[string]any{
			"kind":             "request",
			"responderAgentId": " agt_selected",
		},
	} {
		t.Run(name, func(t *testing.T) {
			candidate := map[string]any{
				WorkflowAgentEditMetadataKey: edit,
				"custom":                     "visible",
			}
			recipient := "agt_selected"
			if name == "other recipient" {
				recipient = "agt_other"
			}
			projected := WorkflowContextMetadataForAgentRecipient(candidate, recipient)
			require.Equal(t, "visible", projected["custom"])
			require.NotContains(t, projected, WorkflowAgentEditMetadataKey)
		})
	}
	require.Contains(t, metadata, WorkflowAgentEditMetadataKey, "projection must not mutate durable metadata")
}

func TestNormalizeWorkflowCallbackTaskWatchHandoff(t *testing.T) {
	t.Parallel()

	normalized, err := NormalizeWorkflowCallbackTaskWatchHandoff(
		&WorkflowCallbackTaskWatchHandoff{
			TaskScopeID: " tsk_callback_work ",
			PredecessorAgentIDs: []string{
				"agt_z", "agt_a", "agt_z",
			},
		},
		true,
		true,
		"calculator-a",
	)
	require.NoError(t, err)
	require.Equal(t, &WorkflowCallbackTaskWatchHandoff{
		TaskScopeID: "tsk_callback_work", PredecessorAgentIDs: []string{"agt_a", "agt_z"},
	}, normalized)

	_, err = NormalizeWorkflowCallbackTaskWatchHandoff(
		&WorkflowCallbackTaskWatchHandoff{TaskScopeID: "tsk_callback_work"},
		false,
		true,
		"calculator-a",
	)
	require.ErrorContains(t, err, "trusted selected-responder context")
	_, err = NormalizeWorkflowCallbackTaskWatchHandoff(
		&WorkflowCallbackTaskWatchHandoff{
			TaskScopeID: "tsk_callback_work", PredecessorAgentIDs: []string{" agt_a"},
		},
		true,
		true,
		"calculator-a",
	)
	require.ErrorContains(t, err, "predecessor Agent ID is invalid")
}