package runtime

import (
	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"

	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestRuntimeRetriesTransientModelFailureWithinCurrentStep(t *testing.T) {
	now := time.Date(2026, 8, 2, 15, 30, 0, 0, time.UTC)
	handles := &memoryHandleStore{}
	checkpoints := &memoryAccountingCheckpointStore{}
	inference := &fakeInference{
		results: []InferenceResult{
			{Usage: map[string]any{"input_tokens": 7}, SettlementRef: "usage_timeout"},
			accountedToolResult(
				knowledgeNoAnswerFunction, `{}`,
				map[string]any{"input_tokens": 11}, "usage_success",
			),
		},
		errs: []error{&Error{Code: upstreamTimeoutCode, Message: "upstream timed out"}, nil},
	}
	var retryDelays []time.Duration
	runtime := &Runtime{
		Resolver: &Resolver{Catalog: &fakeAuthorizedRootCatalog{candidates: []AuthorizedRootCandidate{{
			KnowledgeBaseID: "kb_1", SourceID: "src_1", RootKind: RootKindPrimary,
			RepoRole: RepoRolePrimary, Revision: "sha-fixed", RepositoryRef: "repo-internal-1",
		}}}, Handles: handles, Clock: func() time.Time { return now }},
		Tools:       &Service{Handles: handles, Repository: &fakeRepository{}, Clock: func() time.Time { return now }},
		Inference:   inference,
		Checkpoints: checkpoints,
		Clock:       func() time.Time { return now },
		modelRetryWait: func(ctx context.Context, delay time.Duration) error {
			require.NoError(t, contextError(ctx))
			retryDelays = append(retryDelays, delay)
			return nil
		},
	}

	result, err := runtime.Execute(context.Background(), ExecuteRequest{
		Principal: testPrincipal(), Question: "server", Model: "agentteam/model", ReasonLevel: "high",
		Capability: "signed-token", IdempotencyKey: "attempt-12345678",
	})

	require.NoError(t, err)
	require.Equal(t, "no_answer", result.Status)
	require.Empty(t, result.ErrorCode)
	require.Equal(t, 2, result.ToolUsage.ModelCalls)
	require.Equal(t, float64(18), result.Usage["input_tokens"])
	require.Equal(t, "usage_success", result.SettlementRef)
	require.Equal(t, []string{"usage_timeout", "usage_success"}, result.SettlementRefs)
	require.Equal(t, []time.Duration{time.Second}, retryDelays)
	require.Equal(t, 1, handles.loadCalls, "a model retry must not restart root authorization or exploration")
	require.Len(t, inference.requests, 2)
	require.Equal(t, "attempt-12345678-step-01", inference.requests[0].IdempotencyKey)
	require.Equal(t, "attempt-12345678-step-02", inference.requests[1].IdempotencyKey)
	require.Equal(t, inference.requests[0].Mode, inference.requests[1].Mode)
	require.Equal(t, inference.requests[0].System, inference.requests[1].System)
	require.Equal(t, inference.requests[0].Messages, inference.requests[1].Messages)
	require.Len(t, checkpoints.writes, 4)
	require.Equal(t, 1, checkpoints.writes[1].Generation)
	require.Equal(t, []string{"usage_timeout"}, checkpoints.writes[1].SettlementRefs)
	require.Equal(t, 2, checkpoints.writes[2].Generation)
	require.Equal(t, []string{"usage_timeout", "usage_success"}, checkpoints.writes[2].SettlementRefs)
}

func TestModelCallAccountingStopsAfterFiveTransientRetries(t *testing.T) {
	now := time.Date(2026, 8, 2, 15, 30, 0, 0, time.UTC)
	inference := &fakeInference{}
	for call := 1; call <= maxTransientModelCallRetries+1; call++ {
		inference.results = append(inference.results, InferenceResult{
			Usage: map[string]any{"input_tokens": call}, SettlementRef: fmt.Sprintf("usage_%d", call),
		})
		inference.errs = append(inference.errs, &Error{Code: upstreamUnavailableCode, Message: "upstream unavailable"})
	}
	checkpoints := &memoryAccountingCheckpointStore{}
	var retryDelays []time.Duration
	runtime := &Runtime{
		Inference: inference, Checkpoints: checkpoints, Clock: func() time.Time { return now },
		modelRetryWait: func(_ context.Context, delay time.Duration) error {
			retryDelays = append(retryDelays, delay)
			return nil
		},
	}
	accounting := &modelCallAccounting{
		runtime: runtime,
		request: ExecuteRequest{
			Principal: testPrincipal(), Model: "agentteam/model", Capability: "signed-token",
			IdempotencyKey: "attempt-12345678",
		},
		expiresAt: now.Add(time.Hour),
	}

	_, err := accounting.generate(
		context.Background(), InferenceModeExploration, "high", "system",
		[]InferenceMessage{{Role: "user", Content: "question"}}, 0,
	)

	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, upstreamUnavailableCode, runtimeErr.Code)
	require.Equal(t, maxTransientModelCallRetries+1, inference.calls)
	require.Equal(t, maxTransientModelCallRetries+1, accounting.generation)
	require.Equal(t, maxTransientModelCallRetries+1, accounting.modelCalls)
	require.Equal(t, float64(21), accounting.usage["input_tokens"])
	require.Equal(t, []string{"usage_1", "usage_2", "usage_3", "usage_4", "usage_5", "usage_6"}, accounting.settlementRefs)
	require.Equal(t, []time.Duration{time.Second, 2 * time.Second, 4 * time.Second, 8 * time.Second, 8 * time.Second}, retryDelays)
	require.Len(t, checkpoints.writes, maxTransientModelCallRetries+1)
	require.Equal(t, maxTransientModelCallRetries+1, checkpoints.writes[len(checkpoints.writes)-1].Generation)
	for index, request := range inference.requests {
		require.Equal(t, fmt.Sprintf("attempt-12345678-step-%02d", index+1), request.IdempotencyKey)
	}
}

func TestModelCallAccountingReservesSettlementLineageBeforeCallingInference(t *testing.T) {
	exactByteBoundary := &modelCallAccounting{
		settlementRefs:          []string{"usage_existing"},
		settlementRefsJSONBytes: protocol.MaxAccountingSettlementLineageJSONBytes - 1 - maxEncodedSettlementReferenceJSONBytes,
	}
	require.True(t, exactByteBoundary.canReserveNextSettlementReference())
	exactByteBoundary.settlementRefsJSONBytes++
	require.False(t, exactByteBoundary.canReserveNextSettlementReference())

	exactCountBoundary := &modelCallAccounting{
		settlementRefs:          make([]string, protocol.MaxAccountingSettlementReferences-1),
		settlementRefsJSONBytes: 2,
	}
	require.True(t, exactCountBoundary.canReserveNextSettlementReference())
	exactCountBoundary.settlementRefs = append(exactCountBoundary.settlementRefs, "usage_last")
	require.False(t, exactCountBoundary.canReserveNextSettlementReference())

	tests := []struct {
		name                   string
		settlementRefs         []string
		settlementRefsJSONSize int
	}{
		{
			name:           "encoded bytes",
			settlementRefs: []string{"usage_existing"},
			settlementRefsJSONSize: protocol.MaxAccountingSettlementLineageJSONBytes -
				maxEncodedSettlementReferenceJSONBytes,
		},
		{
			name:                   "reference count",
			settlementRefs:         make([]string, protocol.MaxAccountingSettlementReferences),
			settlementRefsJSONSize: 2,
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			inference := &fakeInference{}
			accounting := &modelCallAccounting{
				runtime: &Runtime{
					Inference: inference, Checkpoints: &memoryAccountingCheckpointStore{},
				},
				request: ExecuteRequest{
					Principal: testPrincipal(), Model: "agentteam/model", Capability: "signed-token",
					IdempotencyKey: "attempt-12345678",
				},
				settlementRefs:          test.settlementRefs,
				settlementRefsJSONBytes: test.settlementRefsJSONSize,
			}

			_, err := accounting.generate(
				context.Background(), InferenceModeExploration, "high", "system",
				[]InferenceMessage{{Role: "user", Content: "question"}}, 0,
			)

			var runtimeErr *Error
			require.ErrorAs(t, err, &runtimeErr)
			require.Equal(t, CodeResourceExhausted, runtimeErr.Code)
			require.Zero(t, inference.calls, "capacity exhaustion must not create an uncheckpointed billable call")
			require.Zero(t, accounting.generation)
			require.Zero(t, accounting.modelCalls)
		})
	}
}

func TestModelCallAccountingTracksEncodedSettlementLineageBytes(t *testing.T) {
	accounting := &modelCallAccounting{}
	for _, settlementRef := range []string{"usage_\x00", "usage_<second>"} {
		accounting.recordSettlement(settlementRef)
		encoded, err := json.Marshal(accounting.settlementRefs)
		require.NoError(t, err)
		require.Equal(t, len(encoded), accounting.settlementRefsJSONBytes)
	}

	encodedBytes := accounting.settlementRefsJSONBytes
	accounting.recordSettlement("usage_\x00")
	require.Equal(t, encodedBytes, accounting.settlementRefsJSONBytes)
	require.Len(t, accounting.settlementRefs, 2)
}

func TestModelCallAccountingPersistsBoundedInitialProducerAccounting(t *testing.T) {
	now := time.Date(2026, 8, 3, 13, 0, 0, 0, time.UTC)
	rawRef := strings.Repeat("private-usage-ref", 32)
	inference := &fakeInference{results: []InferenceResult{{
		Usage: map[string]any{"input_tokens": 7}, SettlementRef: rawRef,
	}}}
	checkpoints := &memoryAccountingCheckpointStore{}
	accounting := &modelCallAccounting{
		runtime: &Runtime{Inference: inference, Checkpoints: checkpoints, Clock: func() time.Time { return now }},
		request: ExecuteRequest{
			Principal: testPrincipal(), Model: "agentteam/model", Capability: "signed-token",
			IdempotencyKey: "attempt-12345678",
		},
		expiresAt: now.Add(time.Hour),
	}

	_, err := accounting.generate(
		context.Background(), InferenceModeExploration, "high", "system",
		[]InferenceMessage{{Role: "user", Content: "question"}}, 0,
	)

	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, CodeMaaSAccountingRecovery, runtimeErr.Code)
	require.Equal(t, 1, inference.calls)
	require.Equal(t, 1, accounting.generation)
	require.Equal(t, float64(7), accounting.usage["input_tokens"])
	require.Len(t, checkpoints.writes, 1)
	wantRef := accountingRecoveryReference("attempt-12345678-step-01")
	require.Equal(t, []string{wantRef}, checkpoints.writes[0].SettlementRefs)
	require.NotContains(t, checkpoints.writes[0].SettlementRefs[0], rawRef)
}

func TestModelCallAccountingProjectsOversizedInitialUsageBeforeCheckpoint(t *testing.T) {
	now := time.Date(2026, 8, 3, 13, 0, 0, 0, time.UTC)
	rawUsage := map[string]any{
		"input_tokens":             7,
		"provider_private_payload": strings.Repeat("x", protocol.MaxAccountingUsageJSONBytes),
	}
	rawJSON, err := json.Marshal(rawUsage)
	require.NoError(t, err)
	require.Greater(t, len(rawJSON), protocol.MaxAccountingUsageJSONBytes)
	inference := &fakeInference{results: []InferenceResult{{
		Usage: rawUsage, SettlementRef: "usage_first",
	}}}
	checkpoints := &memoryAccountingCheckpointStore{}
	accounting := &modelCallAccounting{
		runtime: &Runtime{Inference: inference, Checkpoints: checkpoints, Clock: func() time.Time { return now }},
		request: ExecuteRequest{
			Principal: testPrincipal(), Model: "agentteam/model", Capability: "signed-token",
			IdempotencyKey: "attempt-12345678",
		},
		expiresAt: now.Add(time.Hour),
	}

	result, err := accounting.generate(
		context.Background(), InferenceModeExploration, "high", "system",
		[]InferenceMessage{{Role: "user", Content: "question"}}, 0,
	)

	require.NoError(t, err)
	require.Equal(t, map[string]any{"input_tokens": float64(7)}, result.Usage)
	require.Len(t, checkpoints.writes, 1)
	require.NotContains(t, checkpoints.writes[0].Usage, "provider_private_payload")
	require.Equal(t, []string{"usage_first"}, checkpoints.writes[0].SettlementRefs)
}

func TestModelCallAccountingDoesNotRetryTransientFailureWithInvalidProducerRef(t *testing.T) {
	now := time.Date(2026, 8, 3, 13, 0, 0, 0, time.UTC)
	rawRef := strings.Repeat("private-transient-ref", 32)
	inference := &fakeInference{
		results: []InferenceResult{
			{Usage: map[string]any{"input_tokens": 7}, SettlementRef: rawRef},
			{Usage: map[string]any{"input_tokens": 11}, SettlementRef: "usage_second"},
		},
		errs: []error{
			&Error{Code: upstreamUnavailableCode, Message: "upstream unavailable"}, nil,
		},
	}
	checkpoints := &memoryAccountingCheckpointStore{}
	var retryDelays []time.Duration
	accounting := &modelCallAccounting{
		runtime: &Runtime{
			Inference: inference, Checkpoints: checkpoints, Clock: func() time.Time { return now },
			modelRetryWait: func(_ context.Context, delay time.Duration) error {
				retryDelays = append(retryDelays, delay)
				return nil
			},
		},
		request: ExecuteRequest{
			Principal: testPrincipal(), Model: "agentteam/model", Capability: "signed-token",
			IdempotencyKey: "attempt-12345678",
		},
		expiresAt: now.Add(time.Hour),
	}

	_, err := accounting.generate(
		context.Background(), InferenceModeExploration, "high", "system",
		[]InferenceMessage{{Role: "user", Content: "question"}}, 0,
	)

	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, CodeMaaSAccountingRecovery, runtimeErr.Code)
	require.Equal(t, 1, inference.calls)
	require.Empty(t, retryDelays)
	require.Len(t, checkpoints.writes, 1)
	require.Equal(t, []string{
		accountingRecoveryReference("attempt-12345678-step-01"),
	}, checkpoints.writes[0].SettlementRefs)
}

func TestModelCallAccountingProjectsOversizedTransientUsageBeforeRetry(t *testing.T) {
	now := time.Date(2026, 8, 3, 13, 0, 0, 0, time.UTC)
	rawUsage := map[string]any{
		"input_tokens":             7,
		"provider_private_payload": strings.Repeat("x", protocol.MaxAccountingUsageJSONBytes),
	}
	rawJSON, err := json.Marshal(rawUsage)
	require.NoError(t, err)
	require.Greater(t, len(rawJSON), protocol.MaxAccountingUsageJSONBytes)
	inference := &fakeInference{
		results: []InferenceResult{
			{Usage: rawUsage, SettlementRef: "usage_first"},
			{Usage: map[string]any{"input_tokens": 11}, SettlementRef: "usage_second"},
		},
		errs: []error{
			&Error{Code: upstreamUnavailableCode, Message: "upstream unavailable"}, nil,
		},
	}
	checkpoints := &memoryAccountingCheckpointStore{}
	accounting := &modelCallAccounting{
		runtime: &Runtime{
			Inference: inference, Checkpoints: checkpoints, Clock: func() time.Time { return now },
			modelRetryWait: func(context.Context, time.Duration) error { return nil },
		},
		request: ExecuteRequest{
			Principal: testPrincipal(), Model: "agentteam/model", Capability: "signed-token",
			IdempotencyKey: "attempt-12345678",
		},
		expiresAt: now.Add(time.Hour),
	}

	result, err := accounting.generate(
		context.Background(), InferenceModeExploration, "high", "system",
		[]InferenceMessage{{Role: "user", Content: "question"}}, 0,
	)

	require.NoError(t, err)
	require.Equal(t, 2, inference.calls)
	require.Equal(t, map[string]any{"input_tokens": float64(11)}, result.Usage)
	require.Equal(t, float64(18), accounting.usage["input_tokens"])
	require.Len(t, checkpoints.writes, 2)
	require.NotContains(t, checkpoints.writes[0].Usage, "provider_private_payload")
	require.Equal(t, []string{"usage_first", "usage_second"}, checkpoints.writes[1].SettlementRefs)
}

func TestModelCallRetryStopsWhenContextIsCanceledDuringBackoff(t *testing.T) {
	now := time.Date(2026, 8, 2, 15, 30, 0, 0, time.UTC)
	ctx, cancel := context.WithCancel(context.Background())
	inference := &fakeInference{
		results: []InferenceResult{{Usage: map[string]any{"input_tokens": 3}, SettlementRef: "usage_rate_limited"}},
		errs:    []error{&Error{Code: upstreamRateLimitCode, Message: "rate limited"}},
	}
	checkpoints := &memoryAccountingCheckpointStore{}
	runtime := &Runtime{
		Inference: inference, Checkpoints: checkpoints, Clock: func() time.Time { return now },
		modelRetryWait: func(ctx context.Context, _ time.Duration) error {
			cancel()
			return contextError(ctx)
		},
	}
	accounting := &modelCallAccounting{
		runtime: runtime,
		request: ExecuteRequest{
			Principal: testPrincipal(), Model: "agentteam/model", Capability: "signed-token",
			IdempotencyKey: "attempt-12345678",
		},
		expiresAt: now.Add(time.Hour),
	}

	_, err := accounting.generate(
		ctx, InferenceModeExploration, "high", "system",
		[]InferenceMessage{{Role: "user", Content: "question"}}, 0,
	)

	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, CodeCanceled, runtimeErr.Code)
	require.Equal(t, 1, inference.calls)
	require.Equal(t, 1, accounting.generation)
	require.Equal(t, []string{"usage_rate_limited"}, accounting.settlementRefs)
	require.Len(t, checkpoints.writes, 1)
}

func TestTransientModelCallErrorClassificationIsExact(t *testing.T) {
	for _, code := range []string{upstreamTimeoutCode, upstreamUnavailableCode, upstreamRateLimitCode} {
		require.True(t, isTransientModelCallError(fmt.Errorf("wrapped: %w", &Error{Code: code})), code)
	}
	for _, err := range []error{
		&Error{Code: "upstream.network.failed"},
		&Error{Code: CodeResponseInvalid},
		&Error{Code: CodeMaaSAccountingRecovery, Cause: &Error{Code: upstreamTimeoutCode}},
		context.DeadlineExceeded,
		errors.New("rate limited"),
	} {
		require.False(t, isTransientModelCallError(err), err)
	}
}

func TestDefaultModelRetryWaitHonorsCanceledContext(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	err := (&Runtime{}).waitForModelRetry(ctx, time.Hour)

	var runtimeErr *Error
	require.ErrorAs(t, err, &runtimeErr)
	require.Equal(t, CodeCanceled, runtimeErr.Code)
}