package knowledgeqa

import (
	"context"
	"crypto/ed25519"
	"encoding/base64"
	"encoding/json"
	"errors"
	"log/slog"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/forwarder"
	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"
	"github.com/tidbcloud/AgentTeam/internal/maas/capability"
	knowledgeqastore "github.com/tidbcloud/AgentTeam/internal/store/knowledgeqa"
	"github.com/tidbcloud/AgentTeam/internal/store/message"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
)

type workerTestStore struct {
	turn           knowledgeqastore.Turn
	session        knowledgeqastore.Session
	expired        []knowledgeqastore.QueuedAttempt
	snapshots      []knowledgeqastore.TurnSnapshot
	transitions    []string
	events         []protocol.RuntimeActivity
	eventMu        sync.Mutex
	completion     knowledgeqastore.Completion
	completions    []knowledgeqastore.Completion
	runtimeContext knowledgeqastore.RuntimeContext
	recoveries     []knowledgeqastore.AccountingRecovery
	resolved       []knowledgeqastore.AccountingRecoveryResolution
	resolveErr     error
	retries        []workerTestAccountingRetry
	completeErrs   []error
	citations      []knowledgeqastore.Citation
}

type workerTestAccountingRetry struct {
	attemptID string
	next      time.Time
	expired   bool
}

func TestTurnExecutionTimeoutIsTwentyMinutes(t *testing.T) {
	require.Equal(t, 20*time.Minute, DefaultTurnExecutionTimeout)
	require.Equal(t, DefaultTurnExecutionTimeout, turnExecutionTimeout)
}

func TestKnowledgeQACapabilityClaimsCoverCompleteTurnExecutionWindow(t *testing.T) {
	now := time.Date(2026, 8, 2, 11, 45, 0, 123456789, time.UTC)
	_, privateKey, err := ed25519.GenerateKey(nil)
	require.NoError(t, err)
	worker := &Worker{
		signer: capability.Signer{KeyID: "knowledge-qa-test", PrivateKey: privateKey},
		now:    func() time.Time { return now },
	}
	turn := knowledgeqastore.Turn{
		CurrentAttemptID: "kqa_attempt", OrganizationID: "org_1", WorkspaceID: "wsp_1",
		RequesterUserID: "usr_1", ID: "kqt_1",
	}

	token, err := worker.issueCapability(turn)
	require.NoError(t, err)
	parts := strings.Split(token, ".")
	require.Len(t, parts, 3)
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	require.NoError(t, err)
	var claims struct {
		Subject             string `json:"sub"`
		ExpiresAt           int64  `json:"exp"`
		NotBefore           int64  `json:"nbf"`
		IssuedAt            int64  `json:"iat"`
		CapabilitySessionID string `json:"capability_session_id"`
		OrganizationID      string `json:"organization_id"`
		WorkspaceID         string `json:"workspace_id"`
		RequesterUserID     string `json:"requester_user_id"`
		TurnID              string `json:"turn_id"`
		WorkloadKind        string `json:"workload_kind"`
	}
	require.NoError(t, json.Unmarshal(payload, &claims))

	require.Equal(t, now.Unix(), claims.NotBefore)
	require.Equal(t, now.Unix(), claims.IssuedAt)
	require.Equal(t, now.Add(knowledgeQACapabilityTTL).Unix(), claims.ExpiresAt)
	require.Greater(t, claims.ExpiresAt, now.Add(turnExecutionTimeout).Unix())
	require.Equal(t, int64(knowledgeQACapabilityTTL/time.Second), claims.ExpiresAt-claims.IssuedAt)
	require.Equal(t, "kqa_attempt", claims.Subject)
	require.Equal(t, "kqa_attempt", claims.CapabilitySessionID)
	require.Equal(t, "org_1", claims.OrganizationID)
	require.Equal(t, "wsp_1", claims.WorkspaceID)
	require.Equal(t, "usr_1", claims.RequesterUserID)
	require.Equal(t, "kqt_1", claims.TurnID)
	require.Equal(t, capability.WorkloadKnowledgeQA, claims.WorkloadKind)
}

type workerTestTurnUpdates struct {
	updates []protocol.TurnUpdate
}

func (updates *workerTestTurnUpdates) NotifyKnowledgeQATurnUpdated(_ context.Context, update protocol.TurnUpdate) {
	updates.updates = append(updates.updates, update)
}

type workerTestTx struct{}

func (workerTestTx) QueryRow(context.Context, string, ...any) coredb.Row {
	return nil
}

func (workerTestTx) Query(context.Context, string, ...any) (coredb.Rows, error) {
	return nil, nil
}

func (workerTestTx) Exec(context.Context, string, ...any) (coredb.Result, error) {
	return nil, nil
}

func (workerTestTx) Commit(context.Context) error   { return nil }
func (workerTestTx) Rollback(context.Context) error { return nil }

func (store *workerTestStore) ReadRuntimeContext(context.Context, string, string, string, string) (knowledgeqastore.RuntimeContext, error) {
	return store.runtimeContext, nil
}

func (store *workerTestStore) ListFairQueuedAttempts(context.Context, int) ([]knowledgeqastore.QueuedAttempt, error) {
	return nil, nil
}

func (store *workerTestStore) ListExpiredRunningAttempts(context.Context, time.Time, int) ([]knowledgeqastore.QueuedAttempt, error) {
	return store.expired, nil
}

func (store *workerTestStore) ListPendingAccountingRecoveries(context.Context, time.Time, int) ([]knowledgeqastore.AccountingRecovery, error) {
	return store.recoveries, nil
}

func (store *workerTestStore) ClaimQueuedAttempt(context.Context, string, string, string, time.Time) (knowledgeqastore.Turn, error) {
	return store.turn, nil
}

func (store *workerTestStore) LoadSessionForOwner(context.Context, string, string, string) (knowledgeqastore.Session, error) {
	return store.session, nil
}

func (store *workerTestStore) ApplyTurnSnapshot(_ context.Context, _, _, _ string, snapshot knowledgeqastore.TurnSnapshot, _ time.Time) error {
	store.snapshots = append(store.snapshots, snapshot)
	return nil
}

func (store *workerTestStore) TransitionTurn(_ context.Context, _, _, _, from, to string, _ time.Time) error {
	store.transitions = append(store.transitions, from+"->"+to)
	return nil
}

func (store *workerTestStore) AppendNextEvent(_ context.Context, _, _, _ string, activity protocol.RuntimeActivity, _ map[string]any, _ time.Time) error {
	store.eventMu.Lock()
	defer store.eventMu.Unlock()
	store.events = append(store.events, activity)
	return nil
}

func (store *workerTestStore) eventSnapshot() []protocol.RuntimeActivity {
	store.eventMu.Lock()
	defer store.eventMu.Unlock()
	return append([]protocol.RuntimeActivity(nil), store.events...)
}

func (store *workerTestStore) CompleteAttempt(_ context.Context, _, _, _ string, completion knowledgeqastore.Completion) error {
	store.completion = completion
	store.completions = append(store.completions, completion)
	if len(store.completeErrs) > 0 {
		err := store.completeErrs[0]
		store.completeErrs = store.completeErrs[1:]
		return err
	}
	return nil
}

func (store *workerTestStore) CompleteAttemptWithAssistant(
	ctx context.Context,
	_, _, _ string,
	completion knowledgeqastore.Completion,
	citations []knowledgeqastore.Citation,
	appendAssistant func(context.Context, coredb.Tx) (string, error),
) error {
	assistantMessageID, err := appendAssistant(ctx, workerTestTx{})
	if err != nil {
		return err
	}
	completion.AssistantMessageID = assistantMessageID
	store.citations = append([]knowledgeqastore.Citation(nil), citations...)
	store.completion = completion
	store.completions = append(store.completions, completion)
	return nil
}

func (store *workerTestStore) ResolveAccountingRecovery(_ context.Context, resolution knowledgeqastore.AccountingRecoveryResolution) error {
	store.resolved = append(store.resolved, resolution)
	return store.resolveErr
}

func (store *workerTestStore) RecordAccountingRecoveryRetry(_ context.Context, attemptID string, next time.Time, expired bool) error {
	store.retries = append(store.retries, workerTestAccountingRetry{attemptID: attemptID, next: next, expired: expired})
	return nil
}

type workerTestPolicy struct {
	effective protocol.EffectivePolicy
}

func (policy workerTestPolicy) ResolveEffectivePolicy(_ context.Context, _, _, _, dataPolicyRequirementRef string) (protocol.EffectivePolicy, error) {
	if strings.TrimSpace(dataPolicyRequirementRef) != "" {
		return protocol.EffectivePolicy{}, protocol.ErrDataPolicyUnsupported
	}
	return policy.effective, nil
}

type workerTestRuntime struct {
	prepareRequest  forwarder.KnowledgeQARuntimePrepareRequest
	prepared        forwarder.KnowledgeQARuntimePrepareResult
	request         forwarder.KnowledgeQARuntimeRequest
	result          forwarder.KnowledgeQARuntimeResult
	executeErr      error
	recoveryRequest forwarder.KnowledgeQAAccountingRecoveryRequest
	recoveryResult  forwarder.KnowledgeQAAccountingRecoveryResult
	recoveryErr     error
}

func (runtime *workerTestRuntime) PrepareKnowledgeQATurn(_ context.Context, request forwarder.KnowledgeQARuntimePrepareRequest) (forwarder.KnowledgeQARuntimePrepareResult, error) {
	runtime.prepareRequest = request
	return runtime.prepared, nil
}

func (runtime *workerTestRuntime) ExecuteKnowledgeQATurn(_ context.Context, request forwarder.KnowledgeQARuntimeRequest) (forwarder.KnowledgeQARuntimeResult, error) {
	runtime.request = request
	return runtime.result, runtime.executeErr
}

func (runtime *workerTestRuntime) RecoverKnowledgeQATurnAccounting(_ context.Context, request forwarder.KnowledgeQAAccountingRecoveryRequest) (forwarder.KnowledgeQAAccountingRecoveryResult, error) {
	runtime.recoveryRequest = request
	return runtime.recoveryResult, runtime.recoveryErr
}

type workerTestProvisioner struct{}

func (workerTestProvisioner) EnsureWorkspace(context.Context, string, string, time.Time) (model.RuntimeAppInstallation, error) {
	return model.RuntimeAppInstallation{}, nil
}

type workerTestMessages struct {
	err          error
	metadataSink *map[string]any
}

type workerTestMessageFanout struct {
	appMessages []message.AppMessage
	deliveries  []message.MessageDelivery
	onDeliver   func()
	err         error
}

func (fanout *workerTestMessageFanout) DeliverCommittedMessage(_ context.Context, _ string, appMessage message.AppMessage, delivery message.MessageDelivery) error {
	if fanout.onDeliver != nil {
		fanout.onDeliver()
	}
	fanout.appMessages = append(fanout.appMessages, appMessage)
	fanout.deliveries = append(fanout.deliveries, delivery)
	return fanout.err
}

func (messages workerTestMessages) AppendKnowledgeAssistantMessageTx(_ context.Context, _ coredb.Tx, _, _, _ string, _ model.RuntimeAppInstallation, _ string, metadata map[string]any, _ string) (message.MessageDelivery, error) {
	if messages.metadataSink != nil {
		*messages.metadataSink = metadata
	}
	if messages.err != nil {
		return message.MessageDelivery{}, messages.err
	}
	return message.MessageDelivery{
		AppMessage: &message.AppMessage{ID: "msg_answer"},
	}, nil
}

func TestWorkerPublishesConfirmableSourceBranchChangeMetadata(t *testing.T) {
	now := time.Date(2026, 8, 4, 13, 0, 0, 0, time.UTC)
	_, privateKey, err := ed25519.GenerateKey(nil)
	require.NoError(t, err)
	store := &workerTestStore{
		turn: knowledgeqastore.Turn{
			ID: "kqt_branch", WorkspaceID: "wsp_1", OrganizationID: "org_1",
			SessionID: "kqs_1", RequesterUserID: "usr_1", Question: "Switch to release",
			CurrentAttemptID: "kqa_branch",
		},
		session: knowledgeqastore.Session{
			ID: "kqs_1", OwnerMemberID: "mem_1", ThreadChannelID: "thr_1",
		},
	}
	runtime := &workerTestRuntime{
		prepared: forwarder.KnowledgeQARuntimePrepareResult{
			Roots: []forwarder.KnowledgeQARuntimeRoot{{
				HandleID: "kqh_1", KnowledgeBaseID: "kb_1", SourceID: "src_1",
				RootKind: "primary", RepoRole: "mirror", Revision: strings.Repeat("a", 40),
			}},
			RuntimeVersion: "knowledge-qa-v18", PromptVersion: "native-tool-calling-v16",
		},
		result: forwarder.KnowledgeQARuntimeResult{
			Status: "completed",
			Claims: []forwarder.KnowledgeQARuntimeClaim{{
				Text: "The release branch is available. [E1]", EvidenceIDs: []string{"E1"},
			}},
			Evidence: []forwarder.KnowledgeQARuntimeEvidence{{
				ID: "E1", Citation: forwarder.KnowledgeQARuntimeCitation{
					KnowledgeBaseID: "kb_1", SourceID: "src_1", RootKind: "primary",
					RepoRole: "mirror", Revision: strings.Repeat("b", 40), Path: "README.md",
					ResolvedRef: "release", ResolvedRefKind: "branch",
					LineStart: 1, LineEnd: 1,
				},
			}},
			SourceBranchChange: &forwarder.KnowledgeQASourceBranchChange{
				KnowledgeBaseID: "kb_1", SourceID: "src_1", CurrentRef: "main", TargetRef: "release",
			},
			SearchedRoots: 1,
		},
	}
	var metadata map[string]any
	worker, err := NewWorker(WorkerOptions{
		Store: store, Policies: workerTestPolicy{effective: protocol.EffectivePolicy{
			Model: "model-a", ReasonLevel: "high", WorkloadKind: capability.WorkloadKnowledgeQA,
		}}, Runtime: runtime, Signer: capability.Signer{KeyID: "test", PrivateKey: privateKey},
		Provisioner: workerTestProvisioner{}, Messages: workerTestMessages{metadataSink: &metadata},
		MessageFanout: &workerTestMessageFanout{}, Now: func() time.Time { return now },
	})
	require.NoError(t, err)

	worker.execute(context.Background(), knowledgeqastore.QueuedAttempt{
		WorkspaceID: "wsp_1", TurnID: "kqt_branch", AttemptID: "kqa_branch",
	})

	require.Equal(t, "completed", store.completion.Status)
	require.Equal(t, "release", store.citations[0].ResolvedRef)
	require.Equal(t, "branch", store.citations[0].ResolvedRefKind)
	require.Equal(t, map[string]any{
		"knowledgeQATurn": map[string]any{
			"kind": "knowledge_qa_turn", "schemaVersion": 2, "sessionId": "kqs_1", "turnId": "kqt_branch", "citationCount": 1,
			"sourceBranchChange": map[string]any{
				"knowledgeBaseId": "kb_1", "sourceId": "src_1", "currentRef": "main", "targetRef": "release",
			},
		},
		"knowledgeQATurnId": "kqt_branch",
		"citationCount":     1,
	}, metadata)
}

func TestCancelKnowledgeQATurnWithoutAttemptCancelsDeletedSessionExecution(t *testing.T) {
	firstContext, firstCancel := context.WithCancel(context.Background())
	secondContext, secondCancel := context.WithCancel(context.Background())
	otherContext, otherCancel := context.WithCancel(context.Background())
	t.Cleanup(firstCancel)
	t.Cleanup(secondCancel)
	t.Cleanup(otherCancel)

	worker := &Worker{cancel: map[string]context.CancelFunc{
		workKey("kqt_deleted", "kqa_1"): firstCancel,
		workKey("kqt_deleted", "kqa_2"): secondCancel,
		workKey("kqt_other", "kqa_3"):   otherCancel,
	}}
	worker.CancelKnowledgeQATurn(context.Background(), "kqt_deleted", "")

	require.ErrorIs(t, firstContext.Err(), context.Canceled)
	require.ErrorIs(t, secondContext.Err(), context.Canceled)
	require.NoError(t, otherContext.Err())
}

func TestWorkerResolvesTerminalAccountingCheckpointAfterTurnTerminalization(t *testing.T) {
	now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC)
	store := &workerTestStore{recoveries: []knowledgeqastore.AccountingRecovery{{
		WorkspaceID: "wsp_1", OrganizationID: "org_1", RequesterID: "usr_1",
		TurnID: "kqt_1", AttemptID: "kqa_1", DeadlineAt: now.Add(time.Hour),
	}}}
	runtime := &workerTestRuntime{recoveryResult: forwarder.KnowledgeQAAccountingRecoveryResult{
		Found: true, Terminal: true, Generation: 2,
		Usage:         map[string]any{"input_tokens": float64(17)},
		SettlementRef: "usage_2", SettlementRefs: []string{"usage_1", "usage_2"},
	}}
	worker := &Worker{
		store: store, runtime: runtime, capacity: 8, logger: slog.Default(),
		now: func() time.Time { return now },
	}

	worker.reconcilePendingAccounting(context.Background())

	require.Len(t, store.resolved, 1)
	require.Empty(t, store.retries)
	require.Equal(t, "kqa_1", store.resolved[0].AttemptID)
	require.Equal(t, []string{"usage_1", "usage_2"}, store.resolved[0].SettlementRefs)
	require.Equal(t, "usr_1", runtime.recoveryRequest.Principal.RequesterID)
	require.Equal(t, "kqa_1", runtime.recoveryRequest.AttemptID)
}

func TestWorkerRecoversQuiescentCheckpointAfterMeshXProcessFailure(t *testing.T) {
	now := time.Date(2026, 7, 28, 10, 10, 0, 0, time.UTC)
	startedAt := now.Add(-turnExecutionTimeout - accountingRecoveryCrashGrace)
	store := &workerTestStore{recoveries: []knowledgeqastore.AccountingRecovery{{
		WorkspaceID: "wsp_1", OrganizationID: "org_1", RequesterID: "usr_1",
		TurnID: "kqt_1", AttemptID: "kqa_1", StartedAt: &startedAt,
		DeadlineAt: now.Add(time.Hour),
	}}}
	runtime := &workerTestRuntime{recoveryResult: forwarder.KnowledgeQAAccountingRecoveryResult{
		Found: true, Terminal: false, Generation: 1,
		SettlementRef: "usage_before_crash", SettlementRefs: []string{"usage_before_crash"},
	}}
	worker := &Worker{
		store: store, runtime: runtime, capacity: 8, logger: slog.Default(),
		now: func() time.Time { return now },
	}

	worker.reconcilePendingAccounting(context.Background())

	require.Len(t, store.resolved, 1)
	require.Equal(t, []string{"usage_before_crash"}, store.resolved[0].SettlementRefs)
	require.Empty(t, store.retries)
}

func TestWorkerExpiresMissingAccountingCheckpointOnlyAfterRecoveryWindow(t *testing.T) {
	now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC)
	store := &workerTestStore{recoveries: []knowledgeqastore.AccountingRecovery{{
		WorkspaceID: "wsp_1", OrganizationID: "org_1", RequesterID: "usr_1",
		TurnID: "kqt_1", AttemptID: "kqa_1", Attempts: 7, DeadlineAt: now,
	}}}
	worker := &Worker{
		store: store, runtime: &workerTestRuntime{}, capacity: 8, logger: slog.Default(),
		now: func() time.Time { return now },
	}

	worker.reconcilePendingAccounting(context.Background())

	require.Empty(t, store.resolved)
	require.Len(t, store.retries, 1)
	require.True(t, store.retries[0].expired)
}

func TestWorkerBoundsCheckpointMergeFailuresWithRetryState(t *testing.T) {
	now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC)
	store := &workerTestStore{
		recoveries: []knowledgeqastore.AccountingRecovery{{
			WorkspaceID: "wsp_1", OrganizationID: "org_1", RequesterID: "usr_1",
			TurnID: "kqt_1", AttemptID: "kqa_1", Attempts: 2, DeadlineAt: now.Add(time.Hour),
		}},
		resolveErr: errors.New("database unavailable"),
	}
	worker := &Worker{
		store: store,
		runtime: &workerTestRuntime{recoveryResult: forwarder.KnowledgeQAAccountingRecoveryResult{
			Found: true, Terminal: true,
		}},
		capacity: 8, logger: slog.Default(), now: func() time.Time { return now },
	}

	worker.reconcilePendingAccounting(context.Background())

	require.Len(t, store.resolved, 1)
	require.Len(t, store.retries, 1)
	require.Equal(t, now.Add(30*time.Second), store.retries[0].next)
	require.False(t, store.retries[0].expired)
}

func TestWorkerExpiresAbandonedRunningAttemptAfterDurableDeadline(t *testing.T) {
	now := time.Date(2026, 7, 25, 8, 5, 1, 0, time.UTC)
	store := &workerTestStore{expired: []knowledgeqastore.QueuedAttempt{{
		WorkspaceID: "wsp_1", TurnID: "kqt_expired", AttemptID: "kqa_expired", Status: "generating",
	}}}
	worker := &Worker{
		store: store, capacity: 8, logger: slog.Default(),
		now: func() time.Time { return now },
	}

	worker.scan(context.Background())

	require.Len(t, store.completions, 1)
	require.Equal(t, "failed", store.completions[0].Status)
	require.Equal(t, "knowledge_qa_runtime_timeout", store.completions[0].ErrorCode)
	require.Equal(t, "Knowledge Q&A execution exceeded the twenty minute runtime limit", store.completions[0].ErrorDetail)
	require.True(t, store.completions[0].RecoverAccounting)
	require.Equal(t, now, store.completions[0].Now)
}

func TestWorkerDoesNotRecoverAccountingForExpiredPreInferenceAttempt(t *testing.T) {
	now := time.Date(2026, 7, 25, 8, 5, 1, 0, time.UTC)
	store := &workerTestStore{expired: []knowledgeqastore.QueuedAttempt{{
		WorkspaceID: "wsp_1", TurnID: "kqt_expired", AttemptID: "kqa_expired", Status: "preparing",
	}}}
	worker := &Worker{
		store: store, capacity: 8, logger: slog.Default(),
		now: func() time.Time { return now },
	}

	worker.scan(context.Background())

	require.Len(t, store.completions, 1)
	require.False(t, store.completions[0].RecoverAccounting)
}

func TestKnowledgeQARuntimeErrorCodePreservesMaaSFailureCategory(t *testing.T) {
	require.Equal(t, "resolver.credit.insufficient", knowledgeQARuntimeErrorCode(
		forwarder.UpstreamHTTPError{
			StatusCode: 402,
			Code:       "resolver.credit.insufficient",
			Message:    "organization credits are insufficient",
		},
	))
	require.Equal(t, "knowledge_qa_runtime_failed", knowledgeQARuntimeErrorCode(
		forwarder.UpstreamHTTPError{StatusCode: 500, Code: "untrusted.code"},
	))
}

func TestScopeChangeSummaryDisclosesCountsWithoutObjectNames(t *testing.T) {
	summary, details := scopeChangeSummary(
		knowledgeqastore.Turn{
			RequestedKnowledgeBaseIDs: []string{"kb_ready", "kb_lost"},
			RequestedSourceIDs:        []string{"src_ready", "src_lost", "src_lost"},
		},
		[]forwarder.KnowledgeQARuntimeRoot{{
			KnowledgeBaseID: "kb_ready", SourceID: "src_ready",
		}},
		forwarder.KnowledgeQARuntimeScopeDiagnostics{},
	)

	require.Equal(t, 1, details["missingKnowledgeBases"])
	require.Equal(t, 1, details["missingSources"])
	require.Equal(
		t,
		"Knowledge scope changed: 1 selected knowledge base(s) and 1 selected source(s) were unavailable or did not have complete ready revisions",
		summary,
	)
	require.NotContains(t, summary, "kb_lost")
	require.NotContains(t, summary, "src_lost")
}

func TestScopeChangeSummaryDisclosesPartialReadyRootsInDefaultScope(t *testing.T) {
	summary, details := scopeChangeSummary(
		knowledgeqastore.Turn{},
		[]forwarder.KnowledgeQARuntimeRoot{{
			KnowledgeBaseID: "kb_ready", SourceID: "src_partial",
		}},
		forwarder.KnowledgeQARuntimeScopeDiagnostics{
			AuthorizedKnowledgeBaseCount:  2,
			AuthorizedSourceCount:         3,
			UnavailableKnowledgeBaseCount: 1,
			UnavailableSourceCount:        1,
			PartiallyAvailableSourceCount: 1,
		},
	)

	require.Contains(t, summary, "1 accessible knowledge base(s)")
	require.Contains(t, summary, "1 accessible source(s)")
	require.Contains(t, summary, "partial ready-root coverage")
	require.Equal(t, 2, details["authorizedKnowledgeBaseCount"])
	require.Equal(t, 3, details["authorizedSourceCount"])
	require.NotContains(t, summary, "kb_ready")
	require.NotContains(t, summary, "src_partial")
}

func TestWorkerExecutesDurableServerHostedNoAnswerTurn(t *testing.T) {
	now := time.Date(2026, 7, 25, 8, 0, 0, 0, time.UTC)
	_, privateKey, err := ed25519.GenerateKey(nil)
	require.NoError(t, err)
	store := &workerTestStore{
		turn: knowledgeqastore.Turn{
			ID: "kqt_1", WorkspaceID: "wsp_1", OrganizationID: "org_1",
			SessionID: "kqs_1", RequesterUserID: "usr_1", Question: "What is policy?",
			RequestedKnowledgeBaseIDs: []string{"kb_1"}, CurrentAttemptID: "kqa_1",
		},
		session: knowledgeqastore.Session{
			ID: "kqs_1", OwnerMemberID: "mem_1", ThreadChannelID: "thr_1",
		},
	}
	turnUpdates := &workerTestTurnUpdates{}
	runtime := &workerTestRuntime{
		prepared: forwarder.KnowledgeQARuntimePrepareResult{
			Roots: []forwarder.KnowledgeQARuntimeRoot{{
				HandleID: "kqh_1", KnowledgeBaseID: "kb_1", RootKind: "primary",
				RepoRole: "primary", Revision: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
			}},
			RuntimeVersion: "knowledge-qa-v3", PromptVersion: "exploration-v2",
		},
		result: forwarder.KnowledgeQARuntimeResult{
			Status: "no_answer",
			Roots: []forwarder.KnowledgeQARuntimeRoot{{
				HandleID: "kqh_1", KnowledgeBaseID: "kb_1", RootKind: "primary",
				RepoRole: "primary", Revision: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
			}},
			SearchedRoots: 1,
			Usage:         map[string]any{"input_tokens": 12}, SettlementRef: "settlement_2",
			SettlementRefs: []string{"settlement_1", "settlement_2"},
			RuntimeVersion: "knowledge-qa-v3", PromptVersion: "exploration-v2",
		}}
	worker, err := NewWorker(WorkerOptions{
		Store: store,
		Policies: workerTestPolicy{effective: protocol.EffectivePolicy{
			Model: "model-a", ReasonLevel: "high", WorkloadKind: capability.WorkloadKnowledgeQA, OrganizationPolicyVersion: 4,
		}},
		Runtime: runtime, Signer: capability.Signer{KeyID: "test", PrivateKey: privateKey},
		Provisioner: workerTestProvisioner{}, Messages: workerTestMessages{}, MessageFanout: &workerTestMessageFanout{},
		TurnUpdates: turnUpdates,
		Now:         func() time.Time { return now },
	})
	require.NoError(t, err)

	worker.execute(context.Background(), knowledgeqastore.QueuedAttempt{
		WorkspaceID: "wsp_1", TurnID: "kqt_1", AttemptID: "kqa_1",
	})

	require.Equal(t, []string{
		"preparing->generating",
	}, store.transitions)
	require.Len(t, store.snapshots, 2)
	require.Equal(t, "model-a", store.snapshots[0].EffectiveModel)
	require.Equal(t, "kqh_1", store.snapshots[1].AuthorizedRoots[0].RootHandleRef)
	require.Equal(t, "no_answer", store.completion.Status)
	require.Equal(t, "settlement_2", store.completion.SettlementRef)
	require.Equal(t, []string{"settlement_1", "settlement_2"}, store.completion.SettlementRefs)
	require.Equal(t, "usr_1", runtime.request.Principal.RequesterID)
	require.Equal(t, []string{"kqh_1"}, runtime.request.RootHandleIDs)
	require.Equal(t, "kqt_1", runtime.request.Principal.TurnID)
	require.NotEmpty(t, runtime.request.MaaSCapability)
	require.Equal(t, "kqa_1", runtime.request.IdempotencyKey)
	require.NotEmpty(t, turnUpdates.updates)
	require.Equal(t, protocol.TurnUpdate{
		WorkspaceID: "wsp_1", OwnerMemberID: "mem_1",
		ThreadChannelID: "thr_1", TurnID: "kqt_1",
	}, turnUpdates.updates[len(turnUpdates.updates)-1])
	require.Condition(t, func() bool {
		for _, event := range store.events {
			if event.Kind == "tool" && event.ToolName == "grep" &&
				event.Status == "completed" &&
				event.Summary == "Searched 1 authorized evidence roots; no sufficient evidence was found" {
				return true
			}
		}
		return false
	})
}

func TestWorkerPersistsExplicitDebugModelOutputOnRuntimeFailure(t *testing.T) {
	now := time.Date(2026, 7, 30, 4, 0, 0, 0, time.UTC)
	debugModelOutput := `{"model_output":"` + strings.Repeat("full-debug-output-", 5000) + `"}`
	require.Greater(t, len(debugModelOutput), 64<<10)
	_, privateKey, err := ed25519.GenerateKey(nil)
	require.NoError(t, err)
	store := &workerTestStore{
		turn: knowledgeqastore.Turn{
			ID: "kqt_debug", WorkspaceID: "wsp_1", OrganizationID: "org_1",
			SessionID: "kqs_1", RequesterUserID: "usr_1", Question: "What is policy?",
			CurrentAttemptID: "kqa_debug",
		},
		session: knowledgeqastore.Session{ID: "kqs_1", OwnerMemberID: "mem_1", ThreadChannelID: "thr_1"},
	}
	runtime := &workerTestRuntime{
		prepared: forwarder.KnowledgeQARuntimePrepareResult{
			Roots: []forwarder.KnowledgeQARuntimeRoot{{
				HandleID: "kqh_1", KnowledgeBaseID: "kb_1", SourceID: "src_1",
				RootKind: "primary", RepoRole: "primary", Revision: "sha-fixed",
			}},
		},
		result: forwarder.KnowledgeQARuntimeResult{
			Status: "failed", ErrorCode: "knowledge_qa.resource_exhausted", SearchedRoots: 1,
			DebugModelOutput: debugModelOutput,
		},
	}
	worker, err := NewWorker(WorkerOptions{
		Store: store, Policies: workerTestPolicy{effective: protocol.EffectivePolicy{
			Model: "model-a", ReasonLevel: "high", WorkloadKind: capability.WorkloadKnowledgeQA,
		}}, Runtime: runtime, Signer: capability.Signer{KeyID: "test", PrivateKey: privateKey},
		Provisioner: workerTestProvisioner{}, Messages: workerTestMessages{}, MessageFanout: &workerTestMessageFanout{},
		Now: func() time.Time { return now },
	})
	require.NoError(t, err)

	worker.execute(context.Background(), knowledgeqastore.QueuedAttempt{
		WorkspaceID: "wsp_1", TurnID: "kqt_debug", AttemptID: "kqa_debug",
	})

	require.Equal(t, "failed", store.completion.Status)
	require.Equal(t, "knowledge_qa.resource_exhausted", store.completion.ErrorCode)
	require.Equal(t, "Knowledge Q&A exploration reached its resource limit", store.completion.ErrorDetail)
	require.Equal(t, runtime.result.DebugModelOutput, store.completion.DebugModelOutput)
	require.NotContains(t, store.completion.DebugModelOutput, "debug output truncated")
}

func TestWorkerRecoversAccountingWhenSettledTerminalWriteFails(t *testing.T) {
	now := time.Date(2026, 7, 28, 8, 0, 0, 0, time.UTC)
	_, privateKey, err := ed25519.GenerateKey(nil)
	require.NoError(t, err)
	store := &workerTestStore{
		turn: knowledgeqastore.Turn{
			ID: "kqt_terminal_write", WorkspaceID: "wsp_1", OrganizationID: "org_1",
			SessionID: "kqs_1", RequesterUserID: "usr_1", Question: "What is policy?",
			CurrentAttemptID: "kqa_terminal_write",
		},
		session:      knowledgeqastore.Session{ID: "kqs_1", OwnerMemberID: "mem_1"},
		completeErrs: []error{errors.New("terminal write unavailable")},
	}
	runtime := &workerTestRuntime{
		prepared: forwarder.KnowledgeQARuntimePrepareResult{
			Roots: []forwarder.KnowledgeQARuntimeRoot{{
				HandleID: "kqh_1", KnowledgeBaseID: "kb_1", SourceID: "src_1",
				RootKind: "primary", RepoRole: "primary", Revision: "sha-fixed",
			}},
			RuntimeVersion: "knowledge-qa-v3", PromptVersion: "exploration-v2",
		},
		result: forwarder.KnowledgeQARuntimeResult{
			Status: "no_answer", SearchedRoots: 1,
			Usage: map[string]any{"input_tokens": 12}, SettlementRef: "settlement_2",
			SettlementRefs: []string{"settlement_1", "settlement_2"},
		},
	}
	worker, err := NewWorker(WorkerOptions{
		Store: store, Policies: workerTestPolicy{effective: protocol.EffectivePolicy{
			Model: "model-a", ReasonLevel: "high", WorkloadKind: capability.WorkloadKnowledgeQA,
		}},
		Runtime: runtime, Signer: capability.Signer{KeyID: "test", PrivateKey: privateKey},
		Provisioner: workerTestProvisioner{}, Messages: workerTestMessages{}, MessageFanout: &workerTestMessageFanout{},
		Now: func() time.Time { return now },
	})
	require.NoError(t, err)

	worker.execute(context.Background(), knowledgeqastore.QueuedAttempt{
		WorkspaceID: "wsp_1", TurnID: "kqt_terminal_write", AttemptID: "kqa_terminal_write",
	})

	require.Len(t, store.completions, 2)
	require.Equal(t, map[string]any{"input_tokens": 12}, store.completions[0].Usage)
	require.Equal(t, []string{"settlement_1", "settlement_2"}, store.completions[0].SettlementRefs)
	require.Equal(t, "failed", store.completions[1].Status)
	require.True(t, store.completions[1].RecoverAccounting)
}

func TestWorkerTerminalizesTransportFailureBeforeSchedulingAccountingRecovery(t *testing.T) {
	now := time.Date(2026, 7, 28, 8, 0, 0, 0, time.UTC)
	_, privateKey, err := ed25519.GenerateKey(nil)
	require.NoError(t, err)
	store := &workerTestStore{
		turn: knowledgeqastore.Turn{
			ID: "kqt_transport", WorkspaceID: "wsp_1", OrganizationID: "org_1",
			SessionID: "kqs_1", RequesterUserID: "usr_1", Question: "What is policy?",
			CurrentAttemptID: "kqa_transport",
		},
		session: knowledgeqastore.Session{ID: "kqs_1", OwnerMemberID: "mem_1"},
	}
	runtime := &workerTestRuntime{
		prepared: forwarder.KnowledgeQARuntimePrepareResult{
			Roots: []forwarder.KnowledgeQARuntimeRoot{{
				HandleID: "kqh_1", KnowledgeBaseID: "kb_1", SourceID: "src_1",
				RootKind: "primary", RepoRole: "primary", Revision: "sha-fixed",
			}},
			RuntimeVersion: "knowledge-qa-v3", PromptVersion: "exploration-v2",
		},
		executeErr: context.DeadlineExceeded,
	}
	worker, err := NewWorker(WorkerOptions{
		Store: store, Policies: workerTestPolicy{effective: protocol.EffectivePolicy{
			Model: "model-a", ReasonLevel: "high", WorkloadKind: capability.WorkloadKnowledgeQA,
		}},
		Runtime: runtime, Signer: capability.Signer{KeyID: "test", PrivateKey: privateKey},
		Provisioner: workerTestProvisioner{}, Messages: workerTestMessages{}, MessageFanout: &workerTestMessageFanout{},
		Now: func() time.Time { return now },
	})
	require.NoError(t, err)

	worker.execute(context.Background(), knowledgeqastore.QueuedAttempt{
		WorkspaceID: "wsp_1", TurnID: "kqt_transport", AttemptID: "kqa_transport",
	})

	require.Equal(t, "failed", store.completion.Status)
	require.Equal(t, "knowledge_qa_runtime_timeout", store.completion.ErrorCode)
	require.True(t, store.completion.RecoverAccounting)
}

func TestGroundingDiagnosticCodeBoundsRuntimeValues(t *testing.T) {
	require.Equal(t,
		"knowledge_qa.grounding.claim_evidence_missing",
		groundingDiagnosticCode(" knowledge_qa.grounding.claim_evidence_missing "),
	)
	require.Equal(t, "knowledge_qa.grounding.warning", groundingDiagnosticCode("answer text"))
}

func TestWorkerRetainsSettledUsageWhenAnswerPersistenceFails(t *testing.T) {
	now := time.Date(2026, 7, 25, 8, 0, 0, 0, time.UTC)
	_, privateKey, err := ed25519.GenerateKey(nil)
	require.NoError(t, err)
	store := &workerTestStore{
		turn: knowledgeqastore.Turn{
			ID: "kqt_failed_reply", WorkspaceID: "wsp_1", OrganizationID: "org_1",
			SessionID: "kqs_1", RequesterUserID: "usr_1", Question: "What is policy?",
			CurrentAttemptID: "kqa_failed_reply",
		},
		session: knowledgeqastore.Session{
			ID: "kqs_1", OwnerMemberID: "mem_1", ThreadChannelID: "thr_1",
		},
	}
	runtime := &workerTestRuntime{
		prepared: forwarder.KnowledgeQARuntimePrepareResult{
			Roots: []forwarder.KnowledgeQARuntimeRoot{{
				HandleID: "kqh_1", KnowledgeBaseID: "kb_1", SourceID: "src_1",
				RootKind: "primary", RepoRole: "primary", Revision: "sha-fixed",
			}},
			RuntimeVersion: "knowledge-qa-v3", PromptVersion: "exploration-v2",
		},
		result: forwarder.KnowledgeQARuntimeResult{
			Status: "completed",
			Claims: []forwarder.KnowledgeQARuntimeClaim{{
				Text: "Grounded answer. [E1]", EvidenceIDs: []string{"E1"},
			}},
			Evidence: []forwarder.KnowledgeQARuntimeEvidence{{
				ID: "E1", Citation: forwarder.KnowledgeQARuntimeCitation{
					KnowledgeBaseID: "kb_1", SourceID: "src_1", RootKind: "primary",
					RepoRole: "primary", Revision: "sha-fixed", Path: "docs/policy.md",
					LineStart: 1, LineEnd: 2,
				},
			}},
			SearchedRoots: 1, Usage: map[string]any{"input_tokens": 9},
			SettlementRef:  "settlement_failed_reply_2",
			SettlementRefs: []string{"settlement_failed_reply_1", "settlement_failed_reply_2"},
		},
	}
	worker, err := NewWorker(WorkerOptions{
		Store: store, Policies: workerTestPolicy{effective: protocol.EffectivePolicy{
			Model: "model-a", ReasonLevel: "high", WorkloadKind: capability.WorkloadKnowledgeQA, OrganizationPolicyVersion: 1,
		}}, Runtime: runtime, Signer: capability.Signer{KeyID: "test", PrivateKey: privateKey},
		Provisioner:   workerTestProvisioner{},
		Messages:      workerTestMessages{err: errors.New("reply write failed")},
		MessageFanout: &workerTestMessageFanout{},
		Now:           func() time.Time { return now },
	})
	require.NoError(t, err)

	worker.execute(context.Background(), knowledgeqastore.QueuedAttempt{
		WorkspaceID: "wsp_1", TurnID: "kqt_failed_reply", AttemptID: "kqa_failed_reply",
	})

	require.Equal(t, "failed", store.completion.Status)
	require.Equal(t, "knowledge_qa_reply_failed", store.completion.ErrorCode)
	require.Equal(t, 9, store.completion.Usage["input_tokens"])
	require.Equal(t, "settlement_failed_reply_2", store.completion.SettlementRef)
	require.Equal(t, []string{"settlement_failed_reply_1", "settlement_failed_reply_2"}, store.completion.SettlementRefs)
}

func TestRuntimeFailureMessagesMatchFailureCategory(t *testing.T) {
	tests := []struct {
		name, code, summary, detail string
	}{
		{
			name: "grounding", code: "knowledge_qa.grounded_response_invalid",
			summary: "the generated answer failed groundedness validation",
			detail:  "The generated answer did not pass groundedness validation",
		},
		{
			name: "resource limit", code: "knowledge_qa.resource_exhausted",
			summary: "exploration reached its resource limit",
			detail:  "Knowledge Q&A exploration reached its resource limit",
		},
		{
			name: "repository", code: "knowledge_qa.repository_failed",
			summary: "an authorized knowledge source could not be read",
			detail:  "Knowledge Q&A could not read an authorized knowledge source",
		},
		{
			name: "inference", code: "maas.unavailable",
			summary: "execution failed before a grounded answer was available",
			detail:  "Knowledge Q&A model execution failed",
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			summary, detail := runtimeFailureMessages(test.code)
			require.Equal(t, test.summary, summary)
			require.Equal(t, test.detail, detail)
		})
	}
}