package knowledgeqa

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

	"context"
	"net/http"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/api/meshxapi"
	"github.com/tidbcloud/AgentTeam/internal/identity/trusted"
	"github.com/tidbcloud/AgentTeam/internal/service/knowledgeqa/runtime"
)

type recordingAccountingCheckpointStore struct {
	principal protocol.Principal
	attemptID string
}

func (*recordingAccountingCheckpointStore) PutAccountingCheckpoint(context.Context, protocol.AccountingCheckpoint) error {
	return nil
}

func (store *recordingAccountingCheckpointStore) LoadAccountingCheckpoint(
	_ context.Context,
	principal protocol.Principal,
	attemptID string,
	_ time.Time,
) (protocol.AccountingCheckpoint, bool, error) {
	store.principal = principal
	store.attemptID = attemptID
	return protocol.AccountingCheckpoint{
		Principal: principal, AttemptID: attemptID, Generation: 2,
		State:          protocol.AccountingCheckpointTerminal,
		SettlementRefs: []string{"usage_1", "usage_2"},
	}, true, nil
}

func TestRecoverAccountingDerivesRequesterAuthorityFromTrustedIdentity(t *testing.T) {
	checkpoints := &recordingAccountingCheckpointStore{}
	handler := &Handler{Runtime: &runtime.Runtime{
		Checkpoints: checkpoints,
		Clock:       func() time.Time { return time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC) },
	}}
	httpRequest, err := http.NewRequestWithContext(t.Context(), http.MethodPost, "/", nil)
	require.NoError(t, err)
	httpRequest.Header.Set(meshxapi.HeaderWorkspaceID, "wsp_trusted")

	response, err := handler.RecoverAccounting(meshxapi.AuthenticatedJSON[RecoverAccountingRequest]{
		AuthenticatedRequest: meshxapi.AuthenticatedRequest{
			Request: httpRequest,
			Identity: trusted.Identity{Subject: "usr_trusted", RawClaims: map[string]any{
				"agentteam_internal_purpose": "knowledge_qa_runtime_accounting",
			}},
		},
		Body: RecoverAccountingRequest{
			OrganizationID: "org_1", TurnID: "kqt_1", AttemptID: "kqa_1",
		},
	})

	require.NoError(t, err)
	require.Equal(t, http.StatusOK, response.Status)
	require.Equal(t, protocol.Principal{
		OrganizationID: "org_1", WorkspaceID: "wsp_trusted",
		RequesterID: "usr_trusted", TurnID: "kqt_1",
	}, checkpoints.principal)
	require.Equal(t, "kqa_1", checkpoints.attemptID)
	result, ok := response.Body.(runtime.RecoverAccountingResult)
	require.True(t, ok)
	require.True(t, result.Terminal)
	require.Equal(t, []string{"usage_1", "usage_2"}, result.SettlementRefs)
}

func TestRecoverAccountingRejectsMissingInternalPurpose(t *testing.T) {
	handler := &Handler{Runtime: &runtime.Runtime{Checkpoints: &recordingAccountingCheckpointStore{}}}
	_, err := handler.RecoverAccounting(meshxapi.AuthenticatedJSON[RecoverAccountingRequest]{})

	var httpErr meshxapi.HTTPError
	require.ErrorAs(t, err, &httpErr)
	require.Equal(t, http.StatusForbidden, httpErr.Status)
}

func TestContextWithRuntimeDeadlineDefaultsToTwentyMinutes(t *testing.T) {
	startedAt := time.Now()
	ctx, cancel := contextWithRuntimeDeadline(context.Background())
	defer cancel()

	deadline, ok := ctx.Deadline()
	require.True(t, ok)
	require.WithinDuration(t, startedAt.Add(20*time.Minute), deadline, time.Second)
}

func TestContextWithRuntimeDeadlinePreservesCallerDeadline(t *testing.T) {
	parent, parentCancel := context.WithTimeout(context.Background(), time.Minute)
	defer parentCancel()
	parentDeadline, ok := parent.Deadline()
	require.True(t, ok)

	ctx, cancel := contextWithRuntimeDeadline(parent)
	defer cancel()
	deadline, ok := ctx.Deadline()
	require.True(t, ok)
	require.Equal(t, parentDeadline, deadline)
}