package knowledgeqa

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"time"

	"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"
	"github.com/tidbcloud/AgentTeam/internal/observability/logging"
	knowledgeqastore "github.com/tidbcloud/AgentTeam/internal/store/knowledgeqa"
	"github.com/tidbcloud/AgentTeam/internal/store/message"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
)

type WorkerStore interface {
	ListFairQueuedAttempts(context.Context, int) ([]knowledgeqastore.QueuedAttempt, error)
	ListExpiredRunningAttempts(context.Context, time.Time, int) ([]knowledgeqastore.QueuedAttempt, error)
	ListPendingAccountingRecoveries(context.Context, time.Time, int) ([]knowledgeqastore.AccountingRecovery, error)
	ReadRuntimeContext(context.Context, string, string, string, string) (knowledgeqastore.RuntimeContext, error)
	ClaimQueuedAttempt(context.Context, string, string, string, time.Time) (knowledgeqastore.Turn, error)
	LoadSessionForOwner(context.Context, string, string, string) (knowledgeqastore.Session, error)
	ApplyTurnSnapshot(context.Context, string, string, string, knowledgeqastore.TurnSnapshot, time.Time) error
	TransitionTurn(context.Context, string, string, string, string, string, time.Time) error
	AppendNextEvent(context.Context, string, string, string, protocol.RuntimeActivity, map[string]any, time.Time) error
	CompleteAttempt(context.Context, string, string, string, knowledgeqastore.Completion) error
	CompleteAttemptWithAssistant(context.Context, string, string, string, knowledgeqastore.Completion, []knowledgeqastore.Citation, func(context.Context, coredb.Tx) (string, error)) error
	ResolveAccountingRecovery(context.Context, knowledgeqastore.AccountingRecoveryResolution) error
	RecordAccountingRecoveryRetry(context.Context, string, time.Time, bool) error
}

const (
	// DefaultTurnExecutionTimeout is shared with the outbound MeshX request so
	// the transport cannot expire a valid turn before the worker's durable
	// deadline.
	DefaultTurnExecutionTimeout = 20 * time.Minute
	turnExecutionTimeout        = DefaultTurnExecutionTimeout

	// The capability is minted immediately before dispatch and must remain
	// valid for the complete execution window. A bounded extra minute covers
	// verifier clock leeway, Unix-second truncation, and request transit while
	// keeping the token well below MaaS's maximum accepted lifetime.
	knowledgeQACapabilityExpiryMargin = time.Minute
	knowledgeQACapabilityTTL          = turnExecutionTimeout + knowledgeQACapabilityExpiryMargin
)

type WorkerPolicyResolver interface {
	ResolveEffectivePolicy(context.Context, string, string, string, string) (protocol.EffectivePolicy, error)
}

type WorkerRuntime interface {
	PrepareKnowledgeQATurn(context.Context, forwarder.KnowledgeQARuntimePrepareRequest) (forwarder.KnowledgeQARuntimePrepareResult, error)
	ExecuteKnowledgeQATurn(context.Context, forwarder.KnowledgeQARuntimeRequest) (forwarder.KnowledgeQARuntimeResult, error)
	RecoverKnowledgeQATurnAccounting(context.Context, forwarder.KnowledgeQAAccountingRecoveryRequest) (forwarder.KnowledgeQAAccountingRecoveryResult, error)
}

type WorkerProvisioner interface {
	EnsureWorkspace(context.Context, string, string, time.Time) (model.RuntimeAppInstallation, error)
}

type WorkerMessages interface {
	AppendKnowledgeAssistantMessageTx(context.Context, coredb.Tx, string, string, string, model.RuntimeAppInstallation, string, map[string]any, string) (message.MessageDelivery, error)
}

type WorkerMessageFanout interface {
	DeliverCommittedMessage(context.Context, string, message.AppMessage, message.MessageDelivery) error
}

type TurnUpdateNotifier interface {
	NotifyKnowledgeQATurnUpdated(context.Context, protocol.TurnUpdate)
}

type WorkerOptions struct {
	Store         WorkerStore
	Policies      WorkerPolicyResolver
	Runtime       WorkerRuntime
	Signer        capability.Signer
	Provisioner   WorkerProvisioner
	Messages      WorkerMessages
	MessageFanout WorkerMessageFanout
	TurnUpdates   TurnUpdateNotifier
	Capacity      int
	Logger        *slog.Logger
	Now           func() time.Time
}

type Worker struct {
	store                WorkerStore
	policies             WorkerPolicyResolver
	runtime              WorkerRuntime
	signer               capability.Signer
	provisioner          WorkerProvisioner
	messages             WorkerMessages
	messageFanout        WorkerMessageFanout
	turnUpdates          TurnUpdateNotifier
	progressPollInterval time.Duration
	capacity             int
	logger               *slog.Logger
	now                  func() time.Time
	queue                chan knowledgeqastore.QueuedAttempt
	mu                   sync.Mutex
	scheduled            map[string]struct{}
	cancel               map[string]context.CancelFunc
}

func NewWorker(options WorkerOptions) (*Worker, error) {
	if options.Store == nil || options.Policies == nil || options.Runtime == nil ||
		options.Provisioner == nil || options.Messages == nil || options.MessageFanout == nil {
		return nil, fmt.Errorf("Knowledge Q&A worker dependencies are incomplete")
	}
	if options.Capacity <= 0 {
		options.Capacity = 8
	}
	if options.Logger == nil {
		options.Logger = slog.Default()
	}
	if options.Now == nil {
		options.Now = func() time.Time { return time.Now().UTC() }
	}
	return &Worker{
		store: options.Store, policies: options.Policies, runtime: options.Runtime,
		signer: options.Signer, provisioner: options.Provisioner, messages: options.Messages,
		messageFanout: options.MessageFanout, turnUpdates: options.TurnUpdates,
		progressPollInterval: knowledgeQAProgressPollInterval,
		capacity:             options.Capacity, logger: options.Logger, now: options.Now,
		queue:     make(chan knowledgeqastore.QueuedAttempt, 256),
		scheduled: make(map[string]struct{}), cancel: make(map[string]context.CancelFunc),
	}, nil
}

func (worker *Worker) Start(ctx context.Context) {
	if worker == nil {
		return
	}
	for index := 0; index < worker.capacity; index++ {
		logging.Go(worker.logger, func() { worker.runExecutor(ctx) })
	}
	logging.Go(worker.logger, func() { worker.runPoller(ctx) })
	logging.Go(worker.logger, func() { worker.runAccountingReconciler(ctx) })
}

func (worker *Worker) NotifyKnowledgeQATurnQueued(_ context.Context, workspaceID, turnID, attemptID string) {
	worker.enqueue(knowledgeqastore.QueuedAttempt{
		WorkspaceID: strings.TrimSpace(workspaceID),
		TurnID:      strings.TrimSpace(turnID),
		AttemptID:   strings.TrimSpace(attemptID),
	})
}

func (worker *Worker) CancelKnowledgeQATurn(_ context.Context, turnID, attemptID string) {
	if worker == nil {
		return
	}
	turnID = strings.TrimSpace(turnID)
	attemptID = strings.TrimSpace(attemptID)
	if turnID == "" {
		return
	}
	worker.mu.Lock()
	cancels := make([]context.CancelFunc, 0, 1)
	if attemptID != "" {
		if cancel := worker.cancel[workKey(turnID, attemptID)]; cancel != nil {
			cancels = append(cancels, cancel)
		}
	} else {
		prefix := turnID + "\x00"
		for key, cancel := range worker.cancel {
			if strings.HasPrefix(key, prefix) && cancel != nil {
				cancels = append(cancels, cancel)
			}
		}
	}
	worker.mu.Unlock()
	for _, cancel := range cancels {
		cancel()
	}
}

func (worker *Worker) runPoller(ctx context.Context) {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for {
		worker.scan(ctx)
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
	}
}

func (worker *Worker) scan(ctx context.Context) {
	worker.expireAbandonedAttempts(ctx)
	items, err := worker.store.ListFairQueuedAttempts(ctx, worker.capacity*2)
	if err != nil {
		if ctx.Err() == nil {
			worker.logger.Warn("Knowledge Q&A queue scan failed", "error", err)
		}
		return
	}
	for _, item := range items {
		worker.enqueue(item)
	}
}

func (worker *Worker) expireAbandonedAttempts(ctx context.Context) {
	now := worker.now()
	items, err := worker.store.ListExpiredRunningAttempts(
		ctx,
		now.Add(-turnExecutionTimeout),
		worker.capacity*2,
	)
	if err != nil {
		if ctx.Err() == nil {
			worker.logger.Warn("Knowledge Q&A timeout scan failed", "error", err)
		}
		return
	}
	for _, item := range items {
		err := worker.store.CompleteAttempt(
			ctx,
			item.WorkspaceID,
			item.TurnID,
			item.AttemptID,
			knowledgeqastore.Completion{
				Status:            "failed",
				ErrorCode:         "knowledge_qa_runtime_timeout",
				ErrorDetail:       "Knowledge Q&A execution exceeded the twenty minute runtime limit",
				RecoverAccounting: item.Status == "generating",
				Now:               now,
			},
		)
		if err != nil && !errors.Is(err, knowledgeqastore.ErrNotFound) &&
			!errors.Is(err, knowledgeqastore.ErrInvalidState) && ctx.Err() == nil {
			worker.logger.Warn(
				"Knowledge Q&A expired attempt completion failed",
				"turn_id", item.TurnID,
				"error", err,
			)
		}
	}
}

func (worker *Worker) enqueue(item knowledgeqastore.QueuedAttempt) {
	if worker == nil || item.WorkspaceID == "" || item.TurnID == "" || item.AttemptID == "" {
		return
	}
	key := workKey(item.TurnID, item.AttemptID)
	worker.mu.Lock()
	if _, exists := worker.scheduled[key]; exists {
		worker.mu.Unlock()
		return
	}
	worker.scheduled[key] = struct{}{}
	worker.mu.Unlock()
	select {
	case worker.queue <- item:
	default:
		worker.mu.Lock()
		delete(worker.scheduled, key)
		worker.mu.Unlock()
	}
}

func (worker *Worker) runExecutor(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case item := <-worker.queue:
			worker.execute(ctx, item)
			worker.mu.Lock()
			delete(worker.scheduled, workKey(item.TurnID, item.AttemptID))
			worker.mu.Unlock()
		}
	}
}

func (worker *Worker) execute(parent context.Context, item knowledgeqastore.QueuedAttempt) {
	turn, err := worker.store.ClaimQueuedAttempt(parent, item.WorkspaceID, item.TurnID, item.AttemptID, worker.now())
	if err != nil {
		if !errors.Is(err, knowledgeqastore.ErrInvalidState) && !errors.Is(err, knowledgeqastore.ErrQueueFull) &&
			!errors.Is(err, knowledgeqastore.ErrNotFound) && parent.Err() == nil {
			worker.logger.Warn("Knowledge Q&A claim failed", "turn_id", item.TurnID, "error", err)
		}
		return
	}
	ctx, cancel := context.WithTimeout(parent, turnExecutionTimeout)
	key := workKey(turn.ID, turn.CurrentAttemptID)
	worker.mu.Lock()
	worker.cancel[key] = cancel
	worker.mu.Unlock()
	defer func() {
		cancel()
		worker.mu.Lock()
		delete(worker.cancel, key)
		worker.mu.Unlock()
	}()
	session, err := worker.store.LoadSessionForOwner(ctx, turn.WorkspaceID, turn.RequesterUserID, turn.SessionID)
	if err == nil {
		worker.notifyTurnUpdated(ctx, session, turn, false)
		err = worker.executeClaimed(ctx, turn, session)
	}
	if err != nil {
		status := "failed"
		code := knowledgeQARuntimeErrorCode(err)
		if errors.Is(err, context.Canceled) {
			status = "canceled"
			code = ""
		} else if errors.Is(err, context.DeadlineExceeded) {
			code = "knowledge_qa_runtime_timeout"
		}
		var postDispatchErr *postDispatchError
		recoverAccounting := errors.As(err, &postDispatchErr)
		completionCtx, completionCancel := context.WithTimeout(context.WithoutCancel(parent), 5*time.Second)
		defer completionCancel()
		completionErr := worker.store.CompleteAttempt(completionCtx, turn.WorkspaceID, turn.ID, turn.CurrentAttemptID, knowledgeqastore.Completion{
			Status: status, ErrorCode: code, ErrorDetail: err.Error(),
			RecoverAccounting: recoverAccounting, Now: worker.now(),
		})
		if completionErr == nil {
			worker.notifyTurnUpdated(completionCtx, session, turn, false)
		}
	}
}

func (worker *Worker) executeClaimed(ctx context.Context, turn knowledgeqastore.Turn, session knowledgeqastore.Session) error {
	effective, err := worker.policies.ResolveEffectivePolicy(
		ctx, turn.OrganizationID, turn.RequestedModel, turn.RequestedReasonLevel, "",
	)
	if err != nil {
		return err
	}
	if strings.TrimSpace(effective.DataPolicyRequirementRef) != "" ||
		strings.TrimSpace(effective.DataPolicyProfileRef) != "" ||
		strings.TrimSpace(effective.DataPolicyProfileVersion) != "" {
		return protocol.ErrDataPolicyUnsupported
	}
	if strings.TrimSpace(effective.WorkloadKind) != capability.WorkloadKnowledgeQA {
		return fmt.Errorf("Knowledge Q&A policy workload is invalid")
	}
	snapshot := knowledgeqastore.TurnSnapshot{
		EffectiveModel: effective.Model, EffectiveReasonLevel: effective.ReasonLevel,
		OrganizationPolicyVersion: effective.OrganizationPolicyVersion,
		RuntimeVersion:            "knowledge-qa-v18", PromptVersion: "native-tool-calling-v16",
	}
	if err := worker.store.ApplyTurnSnapshot(ctx, turn.WorkspaceID, turn.ID, turn.CurrentAttemptID, snapshot, worker.now()); err != nil {
		return err
	}
	prepared, err := worker.runtime.PrepareKnowledgeQATurn(ctx, forwarder.KnowledgeQARuntimePrepareRequest{
		Principal: forwarder.KnowledgeQARuntimePrincipal{
			OrganizationID: turn.OrganizationID, WorkspaceID: turn.WorkspaceID,
			RequesterID: turn.RequesterUserID, TurnID: turn.ID,
		},
		KnowledgeBaseIDs: turn.RequestedKnowledgeBaseIDs,
		SourceIDs:        turn.RequestedSourceIDs,
	})
	if err != nil {
		return err
	}
	snapshot.AuthorizedRoots = runtimeRoots(prepared.Roots)
	snapshot.RuntimeVersion = prepared.RuntimeVersion
	snapshot.PromptVersion = prepared.PromptVersion
	if err := worker.store.ApplyTurnSnapshot(ctx, turn.WorkspaceID, turn.ID, turn.CurrentAttemptID, snapshot, worker.now()); err != nil {
		return err
	}
	if summary, details := scopeChangeSummary(
		turn,
		prepared.Roots,
		prepared.ScopeDiagnostics,
	); summary != "" {
		worker.appendActivity(ctx, session, turn,
			protocol.RuntimeActivity{
				Kind: "tool", Status: "completed", ToolName: "scope_check", Summary: summary,
			},
			details, worker.now())
	}
	if len(prepared.Roots) == 0 {
		worker.appendActivity(ctx, session, turn,
			protocol.RuntimeActivity{
				Kind: "tool", Status: "completed", ToolName: "grep",
				Summary: "No ready evidence roots were available in the authorized scope",
			},
			map[string]any{"stage": "retrieval_complete", "searchedRoots": 0}, worker.now())
		return worker.completeAttempt(ctx, session, turn, knowledgeqastore.Completion{
			Status: "no_answer", Now: worker.now(),
		})
	}
	runtimeContext, err := worker.store.ReadRuntimeContext(
		ctx, turn.WorkspaceID, turn.RequesterUserID, turn.SessionID, turn.ID,
	)
	if err != nil {
		return err
	}
	if runtimeContext.Bytes > 0 {
		worker.appendActivity(ctx, session, turn,
			protocol.RuntimeActivity{
				Kind: "tool", Status: "completed", ToolName: "history_read",
				Summary: "Loaded conversation context",
			},
			map[string]any{
				"stage": "history_read", "generation": runtimeContext.Generation,
				"tailTurnCount": len(runtimeContext.Turns), "resultBytes": runtimeContext.Bytes,
				"summaryBytes": len(runtimeContext.Summary),
			}, worker.now())
	}

	capabilityToken, err := worker.issueCapability(turn)
	if err != nil {
		return err
	}
	// Generating is the durable send boundary used by timeout/cancellation
	// accounting recovery. Enter it only after all pre-dispatch work succeeds.
	if err := worker.transition(ctx, session, turn, "preparing", "generating"); err != nil {
		return err
	}
	result, err := worker.executeKnowledgeQATurnWithProgress(ctx, session, turn, forwarder.KnowledgeQARuntimeRequest{
		Principal: forwarder.KnowledgeQARuntimePrincipal{
			OrganizationID: turn.OrganizationID, WorkspaceID: turn.WorkspaceID,
			RequesterID: turn.RequesterUserID, TurnID: turn.ID,
		},
		Question: turn.Question, KnowledgeBaseIDs: turn.RequestedKnowledgeBaseIDs,
		SourceIDs: turn.RequestedSourceIDs, RootHandleIDs: runtimeRootHandleIDs(prepared.Roots),
		WebSearchEnabled: turn.WebSearchEnabled, Model: effective.Model,
		ReasonLevel: effective.ReasonLevel, MaaSCapability: capabilityToken,
		IdempotencyKey: turn.CurrentAttemptID,
		RuntimeContext: runtimeContextProjection(runtimeContext),
	})
	if err != nil {
		return &postDispatchError{cause: err}
	}
	retrievalSummary := fmt.Sprintf("Searched %d authorized evidence roots", result.SearchedRoots)
	switch result.Status {
	case "no_answer":
		retrievalSummary += "; no sufficient evidence was found"
	case "completed":
		if len(result.Diagnostics) > 0 {
			retrievalSummary += " and generated an answer with grounding warnings"
		} else {
			retrievalSummary += " and read grounded evidence"
		}
	case "failed":
		summary, _ := runtimeFailureMessages(result.ErrorCode)
		retrievalSummary += "; " + summary
	}
	worker.appendActivity(ctx, session, turn,
		protocol.RuntimeActivity{
			Kind: "tool", Status: "completed", ToolName: "grep", Summary: retrievalSummary,
		},
		map[string]any{
			"stage": "retrieval_complete", "grepCalls": result.ToolUsage.GrepCalls,
			"listCalls": result.ToolUsage.ListCalls, "readCalls": result.ToolUsage.ReadCalls,
			"webSearchCalls": result.ToolUsage.WebSearchCalls,
			"webReadCalls":   result.ToolUsage.WebReadCalls,
			"webSources":     result.ToolUsage.WebSources,
			"refCalls":       result.ToolUsage.RefCalls,
			"modelCalls":     result.ToolUsage.ModelCalls, "resultBytes": result.ToolUsage.ResultBytes,
			"searchedRoots":      result.SearchedRoots,
			"compactedTurnCount": runtimeCompactedTurnCount(result.RuntimeContext),
		}, worker.now())
	if result.Status == "no_answer" {
		err := worker.completeAttempt(ctx, session, turn, knowledgeqastore.Completion{
			Status: "no_answer", Usage: result.Usage, SettlementRef: result.SettlementRef,
			SettlementRefs: result.SettlementRefs, RuntimeContext: runtimeContextUpdate(result.RuntimeContext), Now: worker.now(),
		})
		if err != nil {
			return &postDispatchError{cause: err}
		}
		return nil
	}
	if result.Status == "failed" {
		errorCode := strings.TrimSpace(result.ErrorCode)
		if errorCode == "" {
			errorCode = "knowledge_qa_runtime_failed"
		}
		_, errorDetail := runtimeFailureMessages(errorCode)
		err := worker.completeAttempt(ctx, session, turn, knowledgeqastore.Completion{
			Status: "failed", ErrorCode: errorCode,
			ErrorDetail: errorDetail, DebugModelOutput: result.DebugModelOutput,
			Usage: result.Usage, SettlementRef: result.SettlementRef,
			SettlementRefs: result.SettlementRefs, RuntimeContext: runtimeContextUpdate(result.RuntimeContext), Now: worker.now(),
		})
		if err != nil {
			return &postDispatchError{cause: err}
		}
		return nil
	}
	if result.Status != "completed" || len(result.Claims) == 0 {
		return &postDispatchError{cause: fmt.Errorf("Knowledge Q&A runtime returned an invalid terminal result")}
	}
	if err := worker.persistRuntimeDiagnostics(ctx, session, turn, result.Diagnostics); err != nil {
		return worker.completePostInferenceFailure(
			ctx, session, turn, result, "knowledge_qa_state_failed", err,
		)
	}
	answer, citations, err := renderRuntimeAnswer(session.ID, result.Claims, result.Evidence)
	if err != nil {
		return worker.completePostInferenceFailure(
			ctx, session, turn, result, "knowledge_qa_state_failed", err,
		)
	}
	installation, err := worker.provisioner.EnsureWorkspace(ctx, turn.WorkspaceID, session.OwnerMemberID, worker.now())
	if err != nil {
		return worker.completePostInferenceFailure(
			ctx, session, turn, result, "knowledge_qa_reply_failed", err,
		)
	}
	var assistantDelivery message.MessageDelivery
	err = worker.store.CompleteAttemptWithAssistant(
		ctx,
		turn.WorkspaceID,
		turn.ID,
		turn.CurrentAttemptID,
		knowledgeqastore.Completion{
			Status: "completed", Usage: result.Usage,
			SettlementRef: result.SettlementRef, SettlementRefs: result.SettlementRefs,
			RuntimeContext: runtimeContextUpdate(result.RuntimeContext), Now: worker.now(),
		},
		citations,
		func(ctx context.Context, tx coredb.Tx) (string, error) {
			turnMetadata := map[string]any{
				"kind":          "knowledge_qa_turn",
				"schemaVersion": 2,
				"sessionId":     turn.SessionID,
				"turnId":        turn.ID,
				"citationCount": len(citations),
			}
			if change := normalizedSourceBranchChange(result.SourceBranchChange); change != nil {
				turnMetadata["sourceBranchChange"] = change
			}
			delivery, err := worker.messages.AppendKnowledgeAssistantMessageTx(
				ctx,
				tx,
				turn.RequesterUserID,
				turn.WorkspaceID,
				session.ThreadChannelID,
				installation,
				answer,
				map[string]any{
					"knowledgeQATurn":   turnMetadata,
					"knowledgeQATurnId": turn.ID,
					"citationCount":     len(citations),
				},
				"knowledge-qa-"+turn.CurrentAttemptID,
			)
			if err != nil {
				return "", err
			}
			if delivery.AppMessage == nil {
				return "", fmt.Errorf("knowledge assistant reply did not produce an app message")
			}
			assistantDelivery = delivery
			return delivery.AppMessage.ID, nil
		},
	)
	if err != nil {
		return worker.completePostInferenceFailure(
			ctx, session, turn, result, "knowledge_qa_reply_failed", err,
		)
	}
	worker.notifyTurnUpdated(ctx, session, turn, true)
	fanoutCtx, fanoutCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
	defer fanoutCancel()
	if err := worker.messageFanout.DeliverCommittedMessage(
		fanoutCtx,
		turn.WorkspaceID,
		*assistantDelivery.AppMessage,
		assistantDelivery,
	); err != nil {
		worker.logger.Warn(
			"Knowledge Q&A assistant message fanout failed",
			"workspace_id", turn.WorkspaceID,
			"turn_id", turn.ID,
			"message_id", assistantDelivery.AppMessage.ID,
			"error", err,
		)
	}
	return nil
}

func (worker *Worker) persistRuntimeDiagnostics(
	ctx context.Context,
	session knowledgeqastore.Session,
	turn knowledgeqastore.Turn,
	diagnostics []forwarder.KnowledgeQARuntimeDiagnostic,
) error {
	seen := make(map[string]struct{}, 4)
	for _, diagnostic := range diagnostics {
		if strings.TrimSpace(diagnostic.Severity) != "warning" {
			continue
		}
		code := groundingDiagnosticCode(diagnostic.Code)
		if _, exists := seen[code]; exists {
			continue
		}
		seen[code] = struct{}{}
		if err := worker.store.AppendNextEvent(
			ctx,
			turn.WorkspaceID,
			turn.ID,
			turn.CurrentAttemptID,
			protocol.RuntimeActivity{
				Kind: "diagnostic", Status: "warning", ToolName: code,
			},
			map[string]any{"stage": "grounding_diagnostic", "code": code},
			worker.now(),
		); err != nil {
			return err
		}
		worker.notifyTurnUpdated(ctx, session, turn, false)
	}
	return nil
}

func groundingDiagnosticCode(code string) string {
	switch strings.TrimSpace(code) {
	case "knowledge_qa.grounding.claim_empty",
		"knowledge_qa.grounding.claim_evidence_missing",
		"knowledge_qa.grounding.evidence_unavailable",
		"knowledge_qa.grounding.marker_mismatch":
		return strings.TrimSpace(code)
	default:
		return "knowledge_qa.grounding.warning"
	}
}

func (worker *Worker) completePostInferenceFailure(
	ctx context.Context,
	session knowledgeqastore.Session,
	turn knowledgeqastore.Turn,
	result forwarder.KnowledgeQARuntimeResult,
	errorCode string,
	cause error,
) error {
	detail := "Knowledge Q&A could not persist the generated answer"
	if cause != nil {
		detail = cause.Error()
	}
	completionCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
	defer cancel()
	if err := worker.completeAttempt(
		completionCtx,
		session,
		turn,
		knowledgeqastore.Completion{
			Status: "failed", ErrorCode: errorCode, ErrorDetail: detail,
			Usage: result.Usage, SettlementRef: result.SettlementRef,
			SettlementRefs: result.SettlementRefs,
			RuntimeContext: runtimeContextUpdate(result.RuntimeContext), Now: worker.now(),
		},
	); err != nil {
		return &postDispatchError{cause: errors.Join(cause, err)}
	}
	return nil
}

func runtimeFailureMessages(errorCode string) (string, string) {
	switch strings.TrimSpace(errorCode) {
	case "knowledge_qa.grounded_response_invalid":
		return "the generated answer failed groundedness validation",
			"The generated answer did not pass groundedness validation"
	case "knowledge_qa.resource_exhausted":
		return "exploration reached its resource limit",
			"Knowledge Q&A exploration reached its resource limit"
	case "knowledge_qa.repository_failed":
		return "an authorized knowledge source could not be read",
			"Knowledge Q&A could not read an authorized knowledge source"
	case "knowledge_qa.maas_accounting_recovery":
		return "the model response could not be delivered while usage accounting is pending recovery",
			"Knowledge Q&A could not deliver the model response while usage accounting is pending recovery"
	default:
		return "execution failed before a grounded answer was available",
			"Knowledge Q&A model execution failed"
	}
}

func normalizedSourceBranchChange(change *forwarder.KnowledgeQASourceBranchChange) map[string]any {
	if change == nil {
		return nil
	}
	knowledgeBaseID := strings.TrimSpace(change.KnowledgeBaseID)
	sourceID := strings.TrimSpace(change.SourceID)
	currentRef := strings.TrimSpace(change.CurrentRef)
	targetRef := strings.TrimSpace(change.TargetRef)
	if knowledgeBaseID == "" || sourceID == "" || targetRef == "" ||
		len(knowledgeBaseID) > 64 || len(sourceID) > 64 || len(currentRef) > 255 || len(targetRef) > 255 {
		return nil
	}
	return map[string]any{
		"knowledgeBaseId": knowledgeBaseID,
		"sourceId":        sourceID,
		"currentRef":      currentRef,
		"targetRef":       targetRef,
	}
}

func scopeChangeSummary(
	turn knowledgeqastore.Turn,
	roots []forwarder.KnowledgeQARuntimeRoot,
	diagnostics forwarder.KnowledgeQARuntimeScopeDiagnostics,
) (string, map[string]any) {
	requestedKnowledgeBases := uniqueStrings(turn.RequestedKnowledgeBaseIDs)
	requestedSources := uniqueStrings(turn.RequestedSourceIDs)
	availableKnowledgeBases := make(map[string]struct{}, len(roots))
	availableSources := make(map[string]struct{}, len(roots))
	for _, root := range roots {
		if knowledgeBaseID := strings.TrimSpace(root.KnowledgeBaseID); knowledgeBaseID != "" {
			availableKnowledgeBases[knowledgeBaseID] = struct{}{}
		}
		if sourceID := strings.TrimSpace(root.SourceID); sourceID != "" {
			availableSources[sourceID] = struct{}{}
		}
	}
	missingKnowledgeBases := countMissing(requestedKnowledgeBases, availableKnowledgeBases)
	missingSources := countMissing(requestedSources, availableSources)
	unavailableKnowledgeBases := 0
	if len(requestedKnowledgeBases) == 0 && len(requestedSources) == 0 {
		unavailableKnowledgeBases = diagnostics.UnavailableKnowledgeBaseCount
	}
	unavailableSources := 0
	if len(requestedSources) == 0 {
		unavailableSources = diagnostics.UnavailableSourceCount
	}
	partiallyAvailableSources := diagnostics.PartiallyAvailableSourceCount
	if missingKnowledgeBases == 0 && missingSources == 0 &&
		unavailableKnowledgeBases == 0 && unavailableSources == 0 &&
		partiallyAvailableSources == 0 {
		return "", nil
	}
	parts := make([]string, 0, 5)
	if missingKnowledgeBases > 0 {
		parts = append(parts, fmt.Sprintf("%d selected knowledge base(s)", missingKnowledgeBases))
	}
	if missingSources > 0 {
		parts = append(parts, fmt.Sprintf("%d selected source(s)", missingSources))
	}
	if unavailableKnowledgeBases > 0 {
		parts = append(parts, fmt.Sprintf("%d accessible knowledge base(s)", unavailableKnowledgeBases))
	}
	if unavailableSources > 0 {
		parts = append(parts, fmt.Sprintf("%d accessible source(s)", unavailableSources))
	}
	if partiallyAvailableSources > 0 {
		parts = append(parts, fmt.Sprintf("%d accessible source(s) with partial ready-root coverage", partiallyAvailableSources))
	}
	return "Knowledge scope changed: " + strings.Join(parts, " and ") +
			" were unavailable or did not have complete ready revisions",
		map[string]any{
			"stage":                        "scope_authorized",
			"missingKnowledgeBases":        missingKnowledgeBases,
			"missingSources":               missingSources,
			"unavailableKnowledgeBases":    unavailableKnowledgeBases,
			"unavailableSources":           unavailableSources,
			"partiallyAvailableSources":    partiallyAvailableSources,
			"authorizedKnowledgeBaseCount": diagnostics.AuthorizedKnowledgeBaseCount,
			"authorizedSourceCount":        diagnostics.AuthorizedSourceCount,
		}
}

func uniqueStrings(values []string) []string {
	result := make([]string, 0, len(values))
	seen := make(map[string]struct{}, len(values))
	for _, value := range values {
		value = strings.TrimSpace(value)
		if value == "" {
			continue
		}
		if _, ok := seen[value]; ok {
			continue
		}
		seen[value] = struct{}{}
		result = append(result, value)
	}
	return result
}

func countMissing(requested []string, available map[string]struct{}) int {
	missing := 0
	for _, id := range requested {
		if _, ok := available[id]; !ok {
			missing++
		}
	}
	return missing
}

func runtimeRootHandleIDs(input []forwarder.KnowledgeQARuntimeRoot) []string {
	out := make([]string, 0, len(input))
	for _, root := range input {
		if handleID := strings.TrimSpace(root.HandleID); handleID != "" {
			out = append(out, handleID)
		}
	}
	return out
}

func runtimeHistoryProjection(input []knowledgeqastore.RuntimeHistoryTurn) []forwarder.KnowledgeQARuntimeHistoryTurn {
	result := make([]forwarder.KnowledgeQARuntimeHistoryTurn, 0, len(input))
	for _, item := range input {
		result = append(result, forwarder.KnowledgeQARuntimeHistoryTurn{
			TurnID: item.TurnID, Status: item.Status, Question: item.Question,
			Answer:     protocol.StripCitationDestinations(item.Answer),
			Provenance: "conversation_history", Evidence: false,
		})
	}
	return result
}

func runtimeContextProjection(input knowledgeqastore.RuntimeContext) forwarder.KnowledgeQARuntimeContext {
	return forwarder.KnowledgeQARuntimeContext{
		Generation:             input.Generation,
		Summary:                protocol.StripCitationDestinations(input.Summary),
		CompactedThroughTurnID: input.CompactedThroughTurnID,
		SummaryVersion:         input.SummaryVersion,
		Turns:                  runtimeHistoryProjection(input.Turns),
	}
}

func runtimeContextUpdate(input *forwarder.KnowledgeQARuntimeContextUpdate) *knowledgeqastore.RuntimeContextUpdate {
	if input == nil {
		return nil
	}
	return &knowledgeqastore.RuntimeContextUpdate{
		ExpectedGeneration: input.ExpectedGeneration, Summary: input.Summary,
		CompactedThroughTurnID: input.CompactedThroughTurnID, SummaryVersion: input.SummaryVersion,
	}
}

func runtimeCompactedTurnCount(input *forwarder.KnowledgeQARuntimeContextUpdate) int {
	if input == nil {
		return 0
	}
	return input.CompactedTurnCount
}

func (worker *Worker) transition(ctx context.Context, session knowledgeqastore.Session, turn knowledgeqastore.Turn, from, to string) error {
	if err := worker.store.TransitionTurn(ctx, turn.WorkspaceID, turn.ID, turn.CurrentAttemptID, from, to, worker.now()); err != nil {
		return err
	}
	worker.notifyTurnUpdated(ctx, session, turn, false)
	return nil
}

func (worker *Worker) appendActivity(ctx context.Context, session knowledgeqastore.Session, turn knowledgeqastore.Turn, activity protocol.RuntimeActivity, payload map[string]any, now time.Time) bool {
	if err := worker.store.AppendNextEvent(
		ctx, turn.WorkspaceID, turn.ID, turn.CurrentAttemptID, activity, payload, now,
	); err == nil {
		worker.notifyTurnUpdated(ctx, session, turn, false)
		return true
	}
	return false
}

func (worker *Worker) completeAttempt(ctx context.Context, session knowledgeqastore.Session, turn knowledgeqastore.Turn, completion knowledgeqastore.Completion) error {
	if err := worker.store.CompleteAttempt(
		ctx, turn.WorkspaceID, turn.ID, turn.CurrentAttemptID, completion,
	); err != nil {
		return err
	}
	worker.notifyTurnUpdated(ctx, session, turn, false)
	return nil
}

func (worker *Worker) notifyTurnUpdated(ctx context.Context, session knowledgeqastore.Session, turn knowledgeqastore.Turn, messageProjectionChanged bool) {
	if worker.turnUpdates == nil {
		return
	}
	worker.turnUpdates.NotifyKnowledgeQATurnUpdated(ctx, protocol.TurnUpdate{
		WorkspaceID:              turn.WorkspaceID,
		OwnerMemberID:            session.OwnerMemberID,
		ThreadChannelID:          session.ThreadChannelID,
		TurnID:                   turn.ID,
		MessageProjectionChanged: messageProjectionChanged,
	})
}

func (worker *Worker) issueCapability(turn knowledgeqastore.Turn) (string, error) {
	now := worker.now()
	return worker.signer.Sign(capability.Claims{
		Issuer: capability.DefaultIssuer, Subject: turn.CurrentAttemptID,
		Audience:  capability.Audience{capability.DefaultAudience},
		NotBefore: now.Unix(), ExpiresAt: now.Add(knowledgeQACapabilityTTL).Unix(),
		CapabilitySessionID: turn.CurrentAttemptID,
		OrganizationID:      turn.OrganizationID, WorkspaceID: turn.WorkspaceID,
		WorkloadKind:    capability.WorkloadKnowledgeQA,
		RequesterUserID: turn.RequesterUserID, TurnID: turn.ID,
		AuthorityGeneration: 0, IssuedAt: now.Unix(),
	})
}

func runtimeRoots(input []forwarder.KnowledgeQARuntimeRoot) []knowledgeqastore.AuthorizedRoot {
	out := make([]knowledgeqastore.AuthorizedRoot, 0, len(input))
	for _, root := range input {
		out = append(out, knowledgeqastore.AuthorizedRoot{
			ID: root.HandleID, RootHandleRef: root.HandleID,
			KnowledgeBaseID: root.KnowledgeBaseID, SourceID: root.SourceID,
			RootKind: root.RootKind, RepoRole: root.RepoRole, Revision: root.Revision,
			FreshnessStatus: root.FreshnessStatus,
		})
	}
	return out
}

func workKey(turnID, attemptID string) string {
	return strings.TrimSpace(turnID) + "\x00" + strings.TrimSpace(attemptID)
}

func knowledgeQARuntimeErrorCode(err error) string {
	if errors.Is(err, protocol.ErrDataPolicyUnsupported) {
		return "maas.data_policy_requirement.unsupported"
	}
	var upstream forwarder.UpstreamHTTPError
	if errors.As(err, &upstream) {
		code := strings.TrimSpace(upstream.Code)
		if len(code) <= 128 {
			for _, prefix := range []string{"resolver.", "maas.", "capability.", "knowledge_qa."} {
				if strings.HasPrefix(code, prefix) {
					return code
				}
			}
		}
	}
	return "knowledge_qa_runtime_failed"
}