package runtime

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

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

func TestModelCallAccountingAcceptsExactSafeIntegerAggregateBoundary(t *testing.T) {
	now := time.Date(2026, 8, 3, 14, 0, 0, 0, time.UTC)
	inference := &fakeInference{
		results: []InferenceResult{
			{Usage: map[string]any{"input_tokens": maxInferenceUsageSafeInteger - 7}, SettlementRef: "usage_first"},
			{Usage: map[string]any{"input_tokens": 7}, 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),
	}

	_, 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, float64(maxInferenceUsageSafeInteger), accounting.usage["input_tokens"])
	require.Len(t, checkpoints.writes, 2)
	require.Equal(t, float64(maxInferenceUsageSafeInteger), checkpoints.writes[1].Usage["input_tokens"])
	require.Equal(t, []string{"usage_first", "usage_second"}, checkpoints.writes[1].SettlementRefs)
	encoded, err := json.Marshal(checkpoints.writes[1].Usage)
	require.NoError(t, err)
	require.JSONEq(t, `{"input_tokens":9007199254740991}`, string(encoded))
}

func TestModelCallAccountingRejectsAggregateSafeIntegerOverflowAfterCheckpointingReference(t *testing.T) {
	now := time.Date(2026, 8, 3, 14, 0, 0, 0, time.UTC)
	inference := &fakeInference{
		results: []InferenceResult{
			{Usage: map[string]any{"input_tokens": maxInferenceUsageSafeInteger - 7}, SettlementRef: "usage_first"},
			{Usage: map[string]any{"input_tokens": 8}, SettlementRef: "usage_second"},
			{Usage: map[string]any{"input_tokens": 1}, SettlementRef: "usage_must_not_run"},
		},
		errs: []error{
			&Error{Code: upstreamUnavailableCode, Message: "upstream unavailable"}, nil, 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),
	}

	_, 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, 2, inference.calls)
	require.Equal(t, float64(maxInferenceUsageSafeInteger-7), accounting.usage["input_tokens"])
	require.Len(t, checkpoints.writes, 2)
	require.Equal(t, float64(maxInferenceUsageSafeInteger-7), checkpoints.writes[1].Usage["input_tokens"])
	require.Equal(t, []string{"usage_first", "usage_second"}, checkpoints.writes[1].SettlementRefs)
}