package agent

import (
	"context"
	"io"
	"os/exec"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/consts"
)

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

	tests := []struct {
		name      string
		toolName  string
		toolInput any
		want      TokenSource
	}{
		{
			name:     "claude MCP alias",
			toolName: "mcp__agentteam__memory_search",
			want:     TokenSource{Kind: tokenSourceKindTool, Name: "agentteam.memory_search"},
		},
		{
			name:     "opencode MCP alias",
			toolName: "mcp_agentteam_memory_search",
			want:     TokenSource{Kind: tokenSourceKindTool, Name: "agentteam.memory_search"},
		},
		{
			name:     "external MCP alias",
			toolName: "mcp__external__lookup",
			want:     TokenSource{Kind: tokenSourceKindTool, Name: "external.lookup"},
		},
		{
			name:     "MCP server with underscores",
			toolName: "mcp__my__server__lookup",
			want:     TokenSource{Kind: tokenSourceKindTool, Name: "my__server.lookup"},
		},
		{
			name:     "native tool",
			toolName: "Read_File",
			want:     TokenSource{Kind: tokenSourceKindTool, Name: "read_file"},
		},
		{
			name:      "JSON command input",
			toolName:  "shell",
			toolInput: `{"command":"grep -n secret /private/work.txt"}`,
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "grep"},
		},
		{
			name:      "map command input",
			toolName:  "Bash",
			toolInput: map[string]any{"cmd": "git diff --stat"},
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "git"},
		},
		{
			name:      "wrapper skipping",
			toolName:  "Shell",
			toolInput: `env A=B sudo -u root /opt/homebrew/bin/rg needle /private/work`,
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "rg"},
		},
		{
			name:      "command lookup short flag",
			toolName:  "shell",
			toolInput: `command -v python3`,
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "command"},
		},
		{
			name:      "command lookup verbose flag",
			toolName:  "shell",
			toolInput: `command -V python3`,
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "command"},
		},
		{
			name:      "command end options wrapper",
			toolName:  "shell",
			toolInput: `command -- /opt/homebrew/bin/rg needle`,
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "rg"},
		},
		{
			name:      "command default path wrapper",
			toolName:  "shell",
			toolInput: `command -p grep needle`,
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "grep"},
		},
		{
			name:      "command lookup after default path flag",
			toolName:  "shell",
			toolInput: `command -p -v python3`,
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "command"},
		},
		{
			name:      "command combined lookup flags",
			toolName:  "shell",
			toolInput: `command -pv python3`,
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "command"},
		},
		{
			name:      "pipeline keeps first executable",
			toolName:  "exec_command",
			toolInput: map[string]any{"cmd": "cat /private/work.txt | head -20"},
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "cat"},
		},
		{
			name:      "namespaced exec command is shell",
			toolName:  "functions.exec_command",
			toolInput: `{"cmd":"find /private -name token"}`,
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "find"},
		},
		{
			name:      "agentteam exec task action",
			toolName:  "shell",
			toolInput: `{"command":"agentteam exec task get - --secret raw"}`,
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "task:get"},
		},
		{
			name:      "agentteam exec workflow action",
			toolName:  "bash",
			toolInput: `agentteam exec workflow run inspect --run-id run_123`,
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "workflow:run"},
		},
		{
			name:      "agentteam executable path",
			toolName:  "shell",
			toolInput: `/private/bin/agentteam exec review add --path /private/repo`,
			want:      TokenSource{Kind: tokenSourceKindCommand, Name: "review:add"},
		},
		{
			name:     "invalid native name",
			toolName: "../../private/tool",
			want:     unknownTokenSource(),
		},
		{
			name:     "malformed MCP alias",
			toolName: "mcp_missingseparator",
			want:     unknownTokenSource(),
		},
		{
			name:     "overlong native name",
			toolName: strings.Repeat("a", maxTokenSourceNameBytes+1),
			want:     unknownTokenSource(),
		},
		{
			name:      "shell without command",
			toolName:  "shell",
			toolInput: `{"path":"/private/work.txt"}`,
			want:      unknownTokenSource(),
		},
	}

	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			require.Equal(t, tt.want, normalizedTokenSource(tt.toolName, tt.toolInput))
		})
	}
}

type tokenSourceTestDriver struct{}

func (tokenSourceTestDriver) Name() string                    { return "codex" }
func (tokenSourceTestDriver) SupportsStdinNotification() bool { return true }
func (tokenSourceTestDriver) MCPToolPrefix() string           { return "mcp__agentteam__" }
func (tokenSourceTestDriver) Spawn(context.Context, drivercontract.LaunchConfig) (drivercontract.SpawnResult, error) {
	return drivercontract.SpawnResult{Cmd: exec.Command("true"), Stdout: io.NopCloser(strings.NewReader(""))}, nil
}
func (tokenSourceTestDriver) ParseLine(string) []drivercontract.ParsedEvent { return nil }
func (tokenSourceTestDriver) EncodeStdinMessage(string, string) []byte      { return nil }

func newTokenSourceTestProcess() (*Process, *[]TokenSourceObservation) {
	observations := []TokenSourceObservation{}
	ap := &Process{
		AgentID:   "agent-1",
		SessionID: "session-primary",
		processRuntimeBinding: processRuntimeBinding{
			driver: tokenSourceTestDriver{},
		},
		processUsageState: processUsageState{
			model: "gpt-test",
			tokenSourceRecorder: func(observation TokenSourceObservation) {
				observations = append(observations, observation)
			},
		},
	}
	return ap, &observations
}

func newTokenSourceAuditTestProcess() (*Process, *[]TokenSourceAuditObservation) {
	ap, _ := newTokenSourceTestProcess()
	observations := []TokenSourceAuditObservation{}
	ap.tokenSourceAuditRecorder = func(observation TokenSourceAuditObservation) {
		observations = append(observations, observation)
	}
	return ap, &observations
}

func usageEvent(sessionID string) drivercontract.ParsedEvent {
	return drivercontract.ParsedEvent{
		Kind:              consts.ParsedEventKindUsage,
		SessionID:         sessionID,
		Model:             "gpt-test",
		InputTokens:       12,
		CachedInputTokens: 3,
		OutputTokens:      7,
	}
}

func completedTool(ap *Process, sessionID, callID, toolName string, input any) {
	completedToolWithOutput(ap, sessionID, callID, toolName, input, "raw output must not be observed")
}

func completedToolWithOutput(ap *Process, sessionID, callID, toolName string, input any, output string) {
	ap.handleParsedEvent(drivercontract.ParsedEvent{
		Kind:       consts.ParsedEventKindToolCall,
		SessionID:  sessionID,
		ToolCallID: callID,
		ToolName:   toolName,
		ToolInput:  input,
	})
	ap.handleParsedEvent(drivercontract.ParsedEvent{
		Kind:       consts.ParsedEventKindToolOutput,
		SessionID:  sessionID,
		ToolCallID: callID,
		ToolName:   toolName,
		Text:       output,
	})
}

func TestProcessTokenSourceDirect(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()

	ap.handleParsedEvent(usageEvent("session-primary"))

	require.Len(t, *observations, 1)
	require.Equal(t, contextTokenSource(), (*observations)[0].Source)
	require.Equal(t, uint64(12), (*observations)[0].InputTokens)
	require.Equal(t, uint64(3), (*observations)[0].CachedInputTokens)
	require.Equal(t, uint64(7), (*observations)[0].OutputTokens)
}

func TestProcessTokenSourceMatchesMachineUsageWithoutProviderAPIKey(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	require.Empty(t, ap.providerAPIKey)
	completedTool(ap, "session-primary", "call-1", "read_file", nil)

	ap.handleParsedEvent(usageEvent("session-primary"))

	require.Len(t, *observations, 2)
	require.Equal(t, "gpt-test", (*observations)[0].Model)
	require.Equal(t, TokenSource{Kind: tokenSourceKindTool, Name: "read_file"}, (*observations)[0].Source)
	require.Equal(t, contextTokenSource(), (*observations)[1].Source)
}

func TestProcessTokenSourceKeepsUsageWithUnknownModelBucket(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	completedTool(ap, "session-child", "call-1", "read_file", nil)

	missingModel := usageEvent("session-child")
	missingModel.Model = ""
	ap.handleParsedEvent(missingModel)

	require.Len(t, *observations, 2)
	require.Equal(t, "unknown/codex", (*observations)[0].Model)
	require.Equal(t, TokenSource{Kind: tokenSourceKindTool, Name: "read_file"}, (*observations)[0].Source)
	require.Equal(t, "unknown/codex", (*observations)[1].Model)
	require.Equal(t, contextTokenSource(), (*observations)[1].Source)

	withModel := usageEvent("session-child")
	ap.handleParsedEvent(withModel)

	require.Len(t, *observations, 3)
	require.Equal(t, unknownTokenSource(), (*observations)[2].Source, "the first usage must consume its source")
}

func TestProcessTokenSourceSessionInitKeepsPreInitToolState(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	ap.SessionID = ""
	completedTool(ap, "", "call-before-init", "read_file", nil)

	ap.handleParsedEvent(drivercontract.ParsedEvent{
		Kind:      consts.ParsedEventKindSessionInit,
		SessionID: "session-primary",
	})
	ap.handleParsedEvent(usageEvent("session-primary"))

	require.Len(t, *observations, 2)
	require.Equal(t, TokenSource{Kind: tokenSourceKindTool, Name: "read_file"}, (*observations)[0].Source)
}

func TestProcessTokenSourceSessionInitKeepsNamedSessionToolState(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	ap.SessionID = "session-old"
	completedTool(ap, "session-next", "call-before-init", "read_file", nil)

	ap.handleParsedEvent(drivercontract.ParsedEvent{
		Kind:      consts.ParsedEventKindSessionInit,
		SessionID: "session-next",
	})
	ap.handleParsedEvent(usageEvent("session-next"))

	require.Len(t, *observations, 2)
	require.Equal(t, TokenSource{Kind: tokenSourceKindTool, Name: "read_file"}, (*observations)[0].Source)
}

func TestProcessTokenSourceSingleAssociated(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	completedToolWithOutput(
		ap,
		"session-primary",
		"call-1",
		"mcp__agentteam__memory_search",
		`{"query":"raw secret"}`,
		"one two three four five",
	)

	ap.handleParsedEvent(usageEvent("session-primary"))

	require.Len(t, *observations, 2)
	require.Equal(t, TokenSource{Kind: tokenSourceKindTool, Name: "agentteam.memory_search"}, (*observations)[0].Source)
	require.Equal(t, uint64(5), (*observations)[0].InputTokens)
	require.Zero(t, (*observations)[0].CachedInputTokens)
	require.Zero(t, (*observations)[0].OutputTokens)
	require.Equal(t, contextTokenSource(), (*observations)[1].Source)
	require.Equal(t, uint64(7), (*observations)[1].InputTokens)
	require.Equal(t, uint64(3), (*observations)[1].CachedInputTokens)
	require.Equal(t, uint64(7), (*observations)[1].OutputTokens)
}

func TestProcessTokenSourceMultipleAllocatesInputAndContextConservatively(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	completedToolWithOutput(ap, "session-primary", "call-1", "mcp__agentteam__memory_search", `{"query":"raw secret"}`, "one two")
	completedToolWithOutput(ap, "session-primary", "call-2", "shell", `{"command":"grep raw /private/path"}`, "one two three four")

	ap.handleParsedEvent(usageEvent("session-primary"))

	require.Len(t, *observations, 3)
	require.Equal(t, TokenSource{Kind: tokenSourceKindTool, Name: "agentteam.memory_search"}, (*observations)[0].Source)
	require.Equal(t, uint64(2), (*observations)[0].InputTokens)
	require.Equal(t, TokenSource{Kind: tokenSourceKindCommand, Name: "grep"}, (*observations)[1].Source)
	require.Equal(t, uint64(4), (*observations)[1].InputTokens)
	require.Equal(t, contextTokenSource(), (*observations)[2].Source)
	require.Equal(t, uint64(6), (*observations)[2].InputTokens)
	require.Equal(t, uint64(3), (*observations)[2].CachedInputTokens)
	require.Equal(t, uint64(7), (*observations)[2].OutputTokens)
	require.Equal(t, uint64(12), sumTokenSourceInput(*observations))
	require.Equal(t, uint64(3), sumTokenSourceCached(*observations))
	require.Equal(t, uint64(7), sumTokenSourceOutput(*observations))
}

func TestProcessTokenSourceAuditRecordsShellDecision(t *testing.T) {
	ap, audits := newTokenSourceAuditTestProcess()

	completedTool(ap, "session-primary", "call-1", "shell", `{"command":"command -v python3"}`)

	require.Len(t, *audits, 1)
	require.Equal(t, "shell_source", (*audits)[0].Event)
	require.Equal(t, "codex", (*audits)[0].Runtime)
	require.Equal(t, "shell", (*audits)[0].ToolName)
	require.Equal(t, "command -v python3", (*audits)[0].RawCommand)
	require.Equal(t, TokenSource{Kind: tokenSourceKindCommand, Name: "command"}, (*audits)[0].Source)
	require.Equal(t, "command_lookup", (*audits)[0].Decision)
	require.Empty(t, (*audits)[0].Candidates)
}

func TestProcessTokenSourceAuditRecordsMultipleCandidates(t *testing.T) {
	ap, audits := newTokenSourceAuditTestProcess()
	completedTool(ap, "session-primary", "call-1", "mcp__agentteam__memory_search", `{"query":"raw secret"}`)
	completedTool(ap, "session-primary", "call-2", "shell", `{"command":"grep raw /private/path"}`)

	ap.handleParsedEvent(usageEvent("session-primary"))

	require.Len(t, *audits, 2, "shell classification plus usage consumption")
	usageAudit := (*audits)[1]
	require.Equal(t, "usage_source", usageAudit.Event)
	require.Equal(t, multipleTokenSource(), usageAudit.Source)
	require.Equal(t, "aggregate", usageAudit.Decision)
	require.Len(t, usageAudit.Candidates, 2)
	require.Equal(t, TokenSource{Kind: tokenSourceKindTool, Name: "agentteam.memory_search"}, usageAudit.Candidates[0].source)
	require.Equal(t, TokenSource{Kind: tokenSourceKindCommand, Name: "grep"}, usageAudit.Candidates[1].source)
	require.Positive(t, usageAudit.Candidates[0].inputTokens)
	require.Positive(t, usageAudit.Candidates[0].outputTokens)
	require.Positive(t, usageAudit.Candidates[1].inputTokens)
	require.Positive(t, usageAudit.Candidates[1].outputTokens)
	require.Len(t, usageAudit.Allocations, 3)
	require.Empty(t, usageAudit.RawCommand, "usage audit candidates stay sanitized")
}

func TestProcessTokenSourceDuplicateLifecycle(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	for range 2 {
		completedTool(ap, "session-primary", "call-replayed", "mcp_agentteam_message_search", `{"q":"raw"}`)
	}

	ap.handleParsedEvent(usageEvent("session-primary"))

	require.Len(t, *observations, 2)
	require.Equal(t, TokenSource{Kind: tokenSourceKindTool, Name: "agentteam.message_search"}, (*observations)[0].Source)
}

func TestProcessTokenSourceConsumesCompletedSourcesOnce(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	completedTool(ap, "session-primary", "call-1", "mcp__agentteam__memory_search", nil)

	ap.handleParsedEvent(usageEvent("session-primary"))
	ap.handleParsedEvent(usageEvent("session-primary"))

	require.Len(t, *observations, 3)
	require.Equal(t, "agentteam.memory_search", (*observations)[0].Source.Name)
	require.Equal(t, contextTokenSource(), (*observations)[2].Source)
}

func TestProcessTokenSourceSessionIsolation(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	completedTool(ap, "session-child", "child-call", "mcp__agentteam__memory_search", nil)

	ap.handleParsedEvent(usageEvent("session-primary"))
	ap.handleParsedEvent(usageEvent("session-other-child"))
	ap.handleParsedEvent(usageEvent("session-child"))

	require.Len(t, *observations, 4)
	require.Equal(t, contextTokenSource(), (*observations)[0].Source)
	require.Equal(t, unknownTokenSource(), (*observations)[1].Source)
	require.Equal(t, "agentteam.memory_search", (*observations)[2].Source.Name)
}

func TestProcessTokenSourcePendingStateIsBounded(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	for i := 0; i < maxPendingTokenSourcesPerSession+5; i++ {
		callID := "call-" + strings.Repeat("x", i+1)
		completedTool(ap, "session-primary", callID, "read_file", nil)
	}

	require.LessOrEqual(t, ap.pendingTokenSourceCount("session-primary"), maxPendingTokenSourcesPerSession)
	ap.handleParsedEvent(usageEvent("session-primary"))
	require.Len(t, *observations, maxPendingTokenSourcesPerSession+1)
}

func TestProcessTokenSourceSessionStateIsBounded(t *testing.T) {
	ap, _ := newTokenSourceTestProcess()
	for i := 0; i < maxTokenSourceSessions+5; i++ {
		sessionID := "child-" + strings.Repeat("x", i+1)
		completedTool(ap, sessionID, "call-"+sessionID, "read_file", nil)
	}

	require.LessOrEqual(t, ap.tokenSourceSessionCount(), maxTokenSourceSessions)
}

func TestProcessTokenSourceTurnBoundaryClearsIncompleteCalls(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	ap.handleParsedEvent(drivercontract.ParsedEvent{
		Kind:       consts.ParsedEventKindToolCall,
		SessionID:  "session-primary",
		ToolCallID: "incomplete",
		ToolName:   "read_file",
	})
	ap.handleParsedEvent(drivercontract.ParsedEvent{
		Kind:      consts.ParsedEventKindTurnEnd,
		SessionID: "session-primary",
	})
	require.Zero(t, ap.pendingTokenSourceCount("session-primary"))

	ap.handleParsedEvent(usageEvent("session-primary"))

	require.Len(t, *observations, 1)
	require.Equal(t, contextTokenSource(), (*observations)[0].Source)
}

func TestProcessTokenSourceToolOutputWithoutCallIDFallback(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	completedTool(ap, "session-primary", "", "mcp__agentteam__message_history_read", `{"target":"raw"}`)

	ap.handleParsedEvent(usageEvent("session-primary"))

	require.Len(t, *observations, 2)
	require.Equal(t, "agentteam.message_history_read", (*observations)[0].Source.Name)
}

func TestProcessTokenSourceExplicitCorrelationIsExact(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	completedTool(ap, "session-primary", "call-exact", "read_file", nil)
	ev := usageEvent("session-primary")
	ev.ToolCallID = "call-exact"

	ap.handleParsedEvent(ev)

	require.Len(t, *observations, 1)
	require.Equal(t, TokenSource{Kind: tokenSourceKindTool, Name: "read_file"}, (*observations)[0].Source)
}

func TestProcessTokenSourceMultipleUsesLargestRemainderWhenEstimatesExceedProviderInput(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	completedToolWithOutput(ap, "session-primary", "call-1", "read_file", nil, "one two three")
	completedToolWithOutput(ap, "session-primary", "call-2", "shell", `{"command":"grep x"}`, "one two three four five")
	ev := usageEvent("session-primary")
	ev.InputTokens = 5
	ev.CachedInputTokens = 2
	ev.OutputTokens = 4

	ap.handleParsedEvent(ev)

	require.Len(t, *observations, 3)
	require.Equal(t, uint64(1), (*observations)[0].InputTokens)
	require.Equal(t, uint64(2), (*observations)[1].InputTokens)
	require.Equal(t, contextTokenSource(), (*observations)[2].Source)
	require.Equal(t, uint64(2), (*observations)[2].InputTokens)
	require.Equal(t, uint64(2), (*observations)[2].CachedInputTokens)
	require.Equal(t, uint64(4), (*observations)[2].OutputTokens)
	require.Equal(t, uint64(5), sumTokenSourceInput(*observations))
}

func TestLargestRemainderTokenSharesHandlesUint64ProductsWithoutOverflow(t *testing.T) {
	total := ^uint64(0) - 1
	shares := largestRemainderTokenShares(total, []tokenSourceCandidate{
		{outputTokens: ^uint64(0) / 2},
		{outputTokens: ^uint64(0) / 2},
	})

	require.Len(t, shares, 2)
	require.Equal(t, total, shares[0]+shares[1])
	require.LessOrEqual(t, shares[0]-shares[1], uint64(1))
}

func TestProcessTokenSourceZeroResultCandidatesKeepWholeUsageInContext(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	completedToolWithOutput(ap, "session-primary", "call-1", "read_file", nil, "")
	completedToolWithOutput(ap, "session-primary", "call-2", "shell", `{"command":"true"}`, "")

	ap.handleParsedEvent(usageEvent("session-primary"))

	require.Len(t, *observations, 3)
	require.Zero(t, (*observations)[0].InputTokens)
	require.Zero(t, (*observations)[1].InputTokens)
	require.Equal(t, contextTokenSource(), (*observations)[2].Source)
	require.Equal(t, uint64(12), (*observations)[2].InputTokens)
}

func TestProcessTokenSourceExplicitCorrelationKeepsExactWholeUsage(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	completedToolWithOutput(ap, "session-primary", "call-exact", "read_file", nil, "small")
	ev := usageEvent("session-primary")
	ev.ToolCallID = "call-exact"

	ap.handleParsedEvent(ev)

	require.Len(t, *observations, 1)
	require.Equal(t, TokenSource{Kind: tokenSourceKindTool, Name: "read_file"}, (*observations)[0].Source)
	require.Equal(t, uint64(12), (*observations)[0].InputTokens)
	require.Equal(t, uint64(3), (*observations)[0].CachedInputTokens)
	require.Equal(t, uint64(7), (*observations)[0].OutputTokens)
}

func TestProcessTokenSourceCorrelationMissDoesNotConsumeUnrelatedTool(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	completedToolWithOutput(ap, "session-primary", "call-real", "read_file", nil, "result")
	missed := usageEvent("session-primary")
	missed.ToolCallID = "call-missing"

	ap.handleParsedEvent(missed)
	exact := usageEvent("session-primary")
	exact.ToolCallID = "call-real"
	ap.handleParsedEvent(exact)

	require.Len(t, *observations, 2)
	require.Equal(t, unknownTokenSource(), (*observations)[0].Source)
	require.Equal(t, TokenSource{Kind: tokenSourceKindTool, Name: "read_file"}, (*observations)[1].Source)
}

func TestProcessTokenSourceZeroShareCandidatesStillRecordCalls(t *testing.T) {
	ap, observations := newTokenSourceTestProcess()
	completedToolWithOutput(ap, "session-primary", "call-1", "read_file", nil, "")
	completedToolWithOutput(ap, "session-primary", "call-2", "shell", `{"command":"true"}`, "")

	ap.handleParsedEvent(usageEvent("session-primary"))

	require.Len(t, *observations, 3)
	require.Equal(t, uint64(1), (*observations)[0].CallCount)
	require.Equal(t, uint64(1), (*observations)[1].CallCount)
	require.Zero(t, (*observations)[2].CallCount)
}

func sumTokenSourceInput(observations []TokenSourceObservation) uint64 {
	var total uint64
	for _, observation := range observations {
		total += observation.InputTokens
	}
	return total
}

func sumTokenSourceCached(observations []TokenSourceObservation) uint64 {
	var total uint64
	for _, observation := range observations {
		total += observation.CachedInputTokens
	}
	return total
}

func sumTokenSourceOutput(observations []TokenSourceObservation) uint64 {
	var total uint64
	for _, observation := range observations {
		total += observation.OutputTokens
	}
	return total
}