package knowledgeqa

import (
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"
)

func runtimeActivityCount(value int) *int {
	return &value
}

func TestProjectRuntimeActivityPayloadProjectsKnownStages(t *testing.T) {
	cases := []struct {
		name           string
		activity       protocol.RuntimeActivity
		raw            string
		wantStage      string
		wantDiagnostic string
		wantModelCalls int
		wantCounts     *protocol.RuntimeActivityCounts
	}{
		{
			name:      "scope authorized",
			activity:  protocol.RuntimeActivity{ToolName: "scope_check"},
			raw:       `{"stage":"scope_authorized","authorizedKnowledgeBaseCount":2,"authorizedSourceCount":4,"missingSources":0}`,
			wantStage: "scope_authorized",
			wantCounts: &protocol.RuntimeActivityCounts{
				AuthorizedKnowledgeBases: runtimeActivityCount(2),
				AuthorizedSources:        runtimeActivityCount(4),
				MissingSources:           runtimeActivityCount(0),
			},
		},
		{
			name:      "history read",
			activity:  protocol.RuntimeActivity{ToolName: "history_read"},
			raw:       `{"stage":"history_read","tailTurnCount":3}`,
			wantStage: "history_read",
			wantCounts: &protocol.RuntimeActivityCounts{
				ContextTurns: runtimeActivityCount(3),
			},
		},
		{
			name:           "exploration progress",
			activity:       protocol.RuntimeActivity{ToolName: "model"},
			raw:            `{"stage":"exploration_progress","modelCalls":4}`,
			wantStage:      "exploration_progress",
			wantModelCalls: 4,
			wantCounts: &protocol.RuntimeActivityCounts{
				ModelCalls: runtimeActivityCount(4),
			},
		},
		{
			name:      "retrieval complete",
			activity:  protocol.RuntimeActivity{ToolName: "grep"},
			raw:       `{"stage":"retrieval_complete","grepCalls":2,"listCalls":1,"readCalls":3,"refCalls":1,"modelCalls":5,"searchedRoots":4}`,
			wantStage: "retrieval_complete",
			wantCounts: &protocol.RuntimeActivityCounts{
				SearchCalls:   runtimeActivityCount(2),
				ListCalls:     runtimeActivityCount(1),
				ReadCalls:     runtimeActivityCount(3),
				RefCalls:      runtimeActivityCount(1),
				ModelCalls:    runtimeActivityCount(5),
				SearchedRoots: runtimeActivityCount(4),
			},
		},
		{
			name:           "grounding diagnostic",
			activity:       protocol.RuntimeActivity{ToolName: "knowledge_qa.grounding.claim_empty"},
			raw:            `{"stage":"grounding_diagnostic","code":"knowledge_qa.grounding.marker_mismatch"}`,
			wantStage:      "grounding_diagnostic",
			wantDiagnostic: "knowledge_qa.grounding.marker_mismatch",
		},
	}

	for _, testCase := range cases {
		t.Run(testCase.name, func(t *testing.T) {
			activity := testCase.activity
			projectRuntimeActivityPayload(&activity, []byte(testCase.raw))
			require.Equal(t, testCase.wantStage, activity.Stage)
			require.Equal(t, testCase.wantDiagnostic, activity.DiagnosticCode)
			require.Equal(t, testCase.wantModelCalls, activity.ModelCalls)
			require.Equal(t, testCase.wantCounts, activity.Counts)
		})
	}
}

func TestProjectRuntimeActivityPayloadRejectsNegativeValuesButKeepsExplicitZero(t *testing.T) {
	activity := protocol.RuntimeActivity{ToolName: "model"}
	projectRuntimeActivityPayload(&activity, []byte(`{"stage":"retrieval_complete","grepCalls":0,"listCalls":-1,"readCalls":2,"refCalls":-1,"modelCalls":0,"searchedRoots":-3}`))

	require.Equal(t, "retrieval_complete", activity.Stage)
	require.Equal(t, 0, activity.ModelCalls)
	require.NotNil(t, activity.Counts)
	require.Equal(t, runtimeActivityCount(0), activity.Counts.SearchCalls)
	require.Nil(t, activity.Counts.ListCalls)
	require.Equal(t, runtimeActivityCount(2), activity.Counts.ReadCalls)
	require.Nil(t, activity.Counts.RefCalls)
	require.Equal(t, runtimeActivityCount(0), activity.Counts.ModelCalls)
	require.Nil(t, activity.Counts.SearchedRoots)
}

func TestProjectRuntimeActivityPayloadFailsClosedAndPreservesLegacyFallback(t *testing.T) {
	t.Run("unknown stage keeps legacy model call only", func(t *testing.T) {
		activity := protocol.RuntimeActivity{ToolName: "model"}
		projectRuntimeActivityPayload(&activity, []byte(`{"stage":"future_stage","modelCalls":7,"query":"secret","path":"docs/private.md"}`))

		require.Equal(t, 7, activity.ModelCalls)
		require.Empty(t, activity.Stage)
		require.Empty(t, activity.DiagnosticCode)
		require.Nil(t, activity.Counts)
	})

	t.Run("malformed payload leaves structured projection empty", func(t *testing.T) {
		activity := protocol.RuntimeActivity{ToolName: "history_read"}
		projectRuntimeActivityPayload(&activity, []byte(`{"stage":`))

		require.Empty(t, activity.Stage)
		require.Empty(t, activity.DiagnosticCode)
		require.Zero(t, activity.ModelCalls)
		require.Nil(t, activity.Counts)
	})

	t.Run("unknown grounding code collapses to generic warning", func(t *testing.T) {
		activity := protocol.RuntimeActivity{ToolName: "knowledge_qa.grounding.evidence_unavailable"}
		projectRuntimeActivityPayload(&activity, []byte(`{"stage":"grounding_diagnostic","code":"knowledge_qa.grounding.private_detail","detail":"secret"}`))

		require.Equal(t, "grounding_diagnostic", activity.Stage)
		require.Equal(t, genericGroundingDiagnosticCode, activity.DiagnosticCode)
		require.Nil(t, activity.Counts)
	})
}