package knowledgeqa

import (
	"context"
	"errors"
	"fmt"
	"slices"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/core/identifiers"
	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
)

// ListFairQueuedAttempts returns at most one oldest queued attempt per
// organization. Repeated worker passes therefore preserve FIFO within an
// organization without letting one organization monopolize global workers.
func (store *Store) ListFairQueuedAttempts(ctx context.Context, limit int) ([]QueuedAttempt, error) {
	if store == nil || store.db == nil {
		return nil, fmt.Errorf("knowledge Q&A store is not configured")
	}
	if limit <= 0 {
		limit = 32
	}
	rows, err := store.db.Query(ctx, `
SELECT workspace_id, id, current_attempt_id
FROM (
  SELECT workspace_id, id, current_attempt_id, organization_id, queued_at,
         ROW_NUMBER() OVER (
           PARTITION BY organization_id
           ORDER BY queued_at ASC, id ASC
         ) AS organization_queue_position
  FROM knowledge_qa_turns
  WHERE status='queued'
) queued
WHERE organization_queue_position=1
ORDER BY queued_at ASC, organization_id ASC, id ASC
LIMIT $1
`, limit)
	if err != nil {
		return nil, fmt.Errorf("list fair Knowledge Q&A queue: %w", err)
	}
	defer rows.Close()
	result := make([]QueuedAttempt, 0)
	for rows.Next() {
		var item QueuedAttempt
		if err := rows.Scan(&item.WorkspaceID, &item.TurnID, &item.AttemptID); err != nil {
			return nil, fmt.Errorf("scan fair Knowledge Q&A queue: %w", err)
		}
		result = append(result, item)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate fair Knowledge Q&A queue: %w", err)
	}
	return result, nil
}

// ListExpiredRunningAttempts returns attempts whose execution deadline has
// elapsed. The worker normally completes these from its context deadline, but
// this durable scan also releases sessions left non-terminal by a server
// process crash. Completion remains conditional on the same current attempt,
// so concurrent Stop/Retry and multiple server replicas are safe.
func (store *Store) ListExpiredRunningAttempts(ctx context.Context, startedBefore time.Time, limit int) ([]QueuedAttempt, error) {
	if store == nil || store.db == nil {
		return nil, fmt.Errorf("knowledge Q&A store is not configured")
	}
	if limit <= 0 {
		limit = 32
	}
	rows, err := store.db.Query(ctx, `
SELECT workspace_id, id, current_attempt_id, status
FROM knowledge_qa_turns
WHERE status IN ('preparing','generating')
  AND started_at IS NOT NULL AND started_at <= $1
ORDER BY started_at ASC, id ASC
LIMIT $2
`, startedBefore.UTC(), limit)
	if err != nil {
		return nil, fmt.Errorf("list expired Knowledge Q&A attempts: %w", err)
	}
	defer rows.Close()
	result := make([]QueuedAttempt, 0)
	for rows.Next() {
		var item QueuedAttempt
		if err := rows.Scan(&item.WorkspaceID, &item.TurnID, &item.AttemptID, &item.Status); err != nil {
			return nil, fmt.Errorf("scan expired Knowledge Q&A attempt: %w", err)
		}
		result = append(result, item)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate expired Knowledge Q&A attempts: %w", err)
	}
	return result, nil
}

func (store *Store) CreateSession(ctx context.Context, params CreateSessionParams) (Session, error) {
	if store == nil || store.db == nil {
		return Session{}, fmt.Errorf("knowledge Q&A store is not configured")
	}
	params = normalizeCreateSession(params)
	if err := validateCreateSession(params); err != nil {
		return Session{}, err
	}
	if params.ID == "" {
		id, err := identifiers.NewEntityID("kqs")
		if err != nil {
			return Session{}, fmt.Errorf("generate knowledge Q&A session id: %w", err)
		}
		params.ID = id
	}
	if existing, err := store.loadSession(ctx, params.WorkspaceID, params.ID); err == nil {
		if sessionMatchesCreate(existing, params) {
			return existing, nil
		}
		return Session{}, ErrForbidden
	} else if !errors.Is(err, ErrNotFound) {
		return Session{}, err
	}
	if err := store.validateSessionThread(ctx, params); err != nil {
		return Session{}, err
	}
	kbJSON, err := model.MarshalJSONObject("requested Knowledge Q&A knowledge base IDs", params.RequestedKnowledgeBaseIDs)
	if err != nil {
		return Session{}, fmt.Errorf("marshal requested knowledge base ids: %w", err)
	}
	sourceJSON, err := model.MarshalJSONObject("requested Knowledge Q&A source IDs", params.RequestedSourceIDs)
	if err != nil {
		return Session{}, fmt.Errorf("marshal requested source ids: %w", err)
	}
	_, err = store.db.Exec(ctx, `
INSERT INTO knowledge_qa_sessions (
  id, workspace_id, organization_id, owner_user_id, owner_member_id,
  root_message_id, thread_channel_id, requested_knowledge_base_ids_json,
  requested_source_ids_json, requested_model, requested_reason_level,
  web_search_enabled, active_turn_id, status, last_activity_at,
  created_at, updated_at
) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,'','active',$13,$13,$13)
`, params.ID, params.WorkspaceID, params.OrganizationID, params.OwnerUserID,
		params.OwnerMemberID, params.RootMessageID, params.ThreadChannelID,
		kbJSON, sourceJSON, params.RequestedModel, params.RequestedReasonLevel,
		params.WebSearchEnabled, params.Now)
	if err != nil {
		if existing, loadErr := store.loadSession(ctx, params.WorkspaceID, params.ID); loadErr == nil &&
			sessionMatchesCreate(existing, params) {
			return existing, nil
		}
		return Session{}, fmt.Errorf("create knowledge Q&A session: %w", err)
	}
	return store.LoadSessionForOwner(ctx, params.WorkspaceID, params.OwnerUserID, params.ID)
}

func sessionMatchesCreate(session Session, params CreateSessionParams) bool {
	return session.WorkspaceID == params.WorkspaceID &&
		session.OrganizationID == params.OrganizationID &&
		session.OwnerUserID == params.OwnerUserID &&
		session.OwnerMemberID == params.OwnerMemberID &&
		session.RootMessageID == params.RootMessageID &&
		session.ThreadChannelID == params.ThreadChannelID &&
		session.RequestedModel == params.RequestedModel &&
		session.RequestedReasonLevel == params.RequestedReasonLevel &&
		session.WebSearchEnabled == params.WebSearchEnabled
}

func (store *Store) LoadSessionForOwner(ctx context.Context, workspaceID, ownerUserID, sessionID string) (Session, error) {
	session, err := store.loadSession(ctx, workspaceID, sessionID)
	if err != nil {
		return Session{}, err
	}
	if strings.TrimSpace(session.OwnerUserID) != strings.TrimSpace(ownerUserID) {
		return Session{}, ErrForbidden
	}
	return session, nil
}

// LoadCitationForOwner resolves citation metadata only through its owning
// session. Historical metadata remains readable to that owner, but opening the
// cited bytes is separately authorized against current MeshX access.
func (store *Store) LoadCitationForOwner(ctx context.Context, workspaceID, ownerUserID, sessionID, citationID string) (Citation, error) {
	if store == nil || store.db == nil {
		return Citation{}, fmt.Errorf("knowledge Q&A store is not configured")
	}
	var citation Citation
	err := store.db.QueryRow(ctx, `
	SELECT citation.id, citation.workspace_id, citation.turn_id,
	       citation.attempt_id, citation.citation_kind, citation.knowledge_base_id,
	       COALESCE(citation.source_id,''), citation.root_kind,
	       citation.repo_role, citation.revision, citation.resolved_ref,
	       citation.resolved_ref_kind, citation.path,
	       citation.line_start, citation.line_end, COALESCE(citation.url,''),
	       COALESCE(citation.title,''), COALESCE(citation.snippet,''),
	       citation.fetched_at, citation.created_at
	FROM knowledge_qa_citations AS citation
JOIN knowledge_qa_turns AS turn ON turn.id=citation.turn_id
JOIN knowledge_qa_sessions AS session ON session.id=turn.session_id
WHERE citation.workspace_id=$1 AND citation.id=$2
  AND session.id=$3 AND session.workspace_id=$1 AND session.owner_user_id=$4
  AND turn.status='completed' AND citation.attempt_id=turn.current_attempt_id
`, strings.TrimSpace(workspaceID), strings.TrimSpace(citationID),
		strings.TrimSpace(sessionID), strings.TrimSpace(ownerUserID)).Scan(
		&citation.ID, &citation.WorkspaceID, &citation.TurnID,
		&citation.AttemptID, &citation.CitationKind, &citation.KnowledgeBaseID,
		&citation.SourceID,
		&citation.RootKind, &citation.RepoRole, &citation.Revision,
		&citation.ResolvedRef, &citation.ResolvedRefKind, &citation.Path,
		&citation.LineStart, &citation.LineEnd,
		&citation.URL, &citation.Title, &citation.Snippet,
		&citation.FetchedAt,
		&citation.CreatedAt,
	)
	if err != nil {
		if coredb.IsNotFound(err) {
			return Citation{}, ErrNotFound
		}
		return Citation{}, fmt.Errorf("load Knowledge Q&A citation: %w", err)
	}
	return citation, nil
}

func (store *Store) CreateTurn(ctx context.Context, params CreateTurnParams) (Turn, error) {
	if store == nil || store.db == nil {
		return Turn{}, fmt.Errorf("knowledge Q&A store is not configured")
	}
	params = normalizeCreateTurn(params)
	if err := validateCreateTurn(params); err != nil {
		return Turn{}, err
	}
	var err error
	if params.ID == "" {
		params.ID, err = identifiers.NewEntityID("kqt")
		if err != nil {
			return Turn{}, fmt.Errorf("generate knowledge Q&A turn id: %w", err)
		}
	}
	if params.AttemptID == "" {
		params.AttemptID, err = identifiers.NewEntityID("kqa")
		if err != nil {
			return Turn{}, fmt.Errorf("generate knowledge Q&A attempt id: %w", err)
		}
	}

	var turn Turn
	err = coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		session, err := lockSession(ctx, tx, params.WorkspaceID, params.SessionID)
		if err != nil {
			return err
		}
		if session.OwnerUserID != params.RequesterUserID ||
			session.OrganizationID != params.OrganizationID {
			return ErrForbidden
		}
		if session.Status != "active" {
			return ErrNotFound
		}

		existing, found, err := loadTurnByIdempotencyKey(ctx, tx, session.ID, params.IdempotencyKey)
		if err != nil {
			return err
		}
		if found {
			turn = existing
			return nil
		}
		if session.ActiveTurnID != "" {
			return ErrTurnBusy
		}
		if err := lockAdmissionPrincipals(ctx, tx, session.OrganizationID, session.OwnerUserID); err != nil {
			return err
		}
		if err := enforceSubmissionAdmission(ctx, tx, session.OrganizationID, session.OwnerUserID, params.Now); err != nil {
			return err
		}
		loadedQuestion, err := loadUserQuestion(ctx, tx, session, params.UserMessageID)
		if err != nil {
			return err
		}
		params.Question = loadedQuestion

		requestedKBs := params.RequestedKnowledgeBaseIDs
		requestedSources := params.RequestedSourceIDs
		if requestedKBs == nil {
			requestedKBs = session.RequestedKnowledgeBaseIDs
		}
		if requestedSources == nil {
			requestedSources = session.RequestedSourceIDs
		}
		kbJSON, err := model.MarshalJSONObject("requested Knowledge Q&A knowledge base IDs", requestedKBs)
		if err != nil {
			return fmt.Errorf("marshal requested knowledge base ids: %w", err)
		}
		sourceJSON, err := model.MarshalJSONObject("requested Knowledge Q&A source IDs", requestedSources)
		if err != nil {
			return fmt.Errorf("marshal requested source ids: %w", err)
		}
		if _, err := tx.Exec(ctx, `
INSERT INTO knowledge_qa_turns (
  id, workspace_id, organization_id, session_id, requester_user_id,
  user_message_id, assistant_message_id, question, idempotency_key,
  requested_knowledge_base_ids_json, requested_source_ids_json,
  authorized_roots_json, requested_model, requested_reason_level,
  web_search_enabled, current_attempt_id, status, queued_at, updated_at
) VALUES ($1,$2,$3,$4,$5,$6,'',$7,$8,$9,$10,JSON_ARRAY(),$11,$12,$13,$14,'queued',$15,$15)
`, params.ID, session.WorkspaceID, session.OrganizationID, session.ID,
			session.OwnerUserID, params.UserMessageID, params.Question,
			params.IdempotencyKey, kbJSON, sourceJSON, session.RequestedModel,
			session.RequestedReasonLevel, session.WebSearchEnabled, params.AttemptID,
			params.Now); err != nil {
			return fmt.Errorf("create knowledge Q&A turn: %w", err)
		}
		if _, err := tx.Exec(ctx, `
INSERT INTO knowledge_qa_turn_attempts (
  id, workspace_id, organization_id, turn_id, attempt_number,
  idempotency_key, status, usage_json, settlement_refs_json, queued_at, updated_at
) VALUES ($1,$2,$3,$4,1,$5,'queued',JSON_OBJECT(),JSON_ARRAY(),$6,$6)
`, params.AttemptID, session.WorkspaceID, session.OrganizationID,
			params.ID, params.IdempotencyKey, params.Now); err != nil {
			return fmt.Errorf("create knowledge Q&A turn attempt: %w", err)
		}
		if _, err := tx.Exec(ctx, `
UPDATE knowledge_qa_sessions
SET active_turn_id=$2, last_activity_at=$3, updated_at=$3
WHERE id=$1
`, session.ID, params.ID, params.Now); err != nil {
			return fmt.Errorf("activate knowledge Q&A turn: %w", err)
		}
		turn, err = loadTurn(ctx, tx, session.WorkspaceID, params.ID)
		return err
	})
	return turn, err
}

func (store *Store) RetryTurn(ctx context.Context, params RetryTurnParams) (Turn, error) {
	if store == nil || store.db == nil {
		return Turn{}, fmt.Errorf("knowledge Q&A store is not configured")
	}
	params = normalizeRetryTurn(params)
	if params.WorkspaceID == "" || params.SessionID == "" || params.TurnID == "" ||
		params.RequesterUserID == "" || params.IdempotencyKey == "" {
		return Turn{}, fmt.Errorf("workspace, session, turn, requester, and idempotency key are required")
	}
	var err error
	if params.AttemptID == "" {
		params.AttemptID, err = identifiers.NewEntityID("kqa")
		if err != nil {
			return Turn{}, fmt.Errorf("generate knowledge Q&A attempt id: %w", err)
		}
	}
	var turn Turn
	err = coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		session, err := lockSession(ctx, tx, params.WorkspaceID, params.SessionID)
		if err != nil {
			return err
		}
		if session.OwnerUserID != params.RequesterUserID {
			return ErrForbidden
		}
		if session.ActiveTurnID != "" && session.ActiveTurnID != params.TurnID {
			return ErrTurnBusy
		}
		turn, err = loadTurn(ctx, tx, params.WorkspaceID, params.TurnID)
		if err != nil {
			return err
		}
		if turn.SessionID != session.ID {
			return ErrNotFound
		}
		var replayAttemptID string
		err = tx.QueryRow(ctx, `
SELECT id FROM knowledge_qa_turn_attempts
WHERE turn_id=$1 AND idempotency_key=$2
`, turn.ID, params.IdempotencyKey).Scan(&replayAttemptID)
		switch {
		case err == nil:
			// The first request may already have queued or even completed this
			// retry before its HTTP response was lost. Replay the durable turn
			// instead of rejecting the same idempotency key based on its new
			// non-terminal/terminal status.
			return nil
		case coredb.IsNotFound(err):
		default:
			return fmt.Errorf("load knowledge Q&A retry idempotency key: %w", err)
		}
		retryable, err := retryableKnowledgeQATurn(ctx, tx, turn)
		if err != nil {
			return err
		}
		if !retryable {
			return ErrInvalidState
		}
		latestTerminalTurnID, err := latestTerminalKnowledgeQATurn(ctx, tx, session.WorkspaceID, session.ID)
		if err != nil {
			return err
		}
		if latestTerminalTurnID != turn.ID {
			return ErrInvalidState
		}
		if err := lockAdmissionPrincipals(ctx, tx, session.OrganizationID, session.OwnerUserID); err != nil {
			return err
		}
		if err := enforceSubmissionAdmission(ctx, tx, session.OrganizationID, session.OwnerUserID, params.Now); err != nil {
			return err
		}
		var attemptNumber int64
		if err := tx.QueryRow(ctx, `
SELECT COALESCE(MAX(attempt_number),0)+1
FROM knowledge_qa_turn_attempts WHERE turn_id=$1
`, turn.ID).Scan(&attemptNumber); err != nil {
			return fmt.Errorf("allocate knowledge Q&A attempt number: %w", err)
		}
		if _, err := tx.Exec(ctx, `
INSERT INTO knowledge_qa_turn_attempts (
  id, workspace_id, organization_id, turn_id, attempt_number,
  idempotency_key, status, usage_json, settlement_refs_json, queued_at, updated_at
) VALUES ($1,$2,$3,$4,$5,$6,'queued',JSON_OBJECT(),JSON_ARRAY(),$7,$7)
`, params.AttemptID, session.WorkspaceID, session.OrganizationID, turn.ID,
			attemptNumber, params.IdempotencyKey, params.Now); err != nil {
			return fmt.Errorf("create Knowledge Q&A retry attempt: %w", err)
		}
		if _, err := tx.Exec(ctx, `
UPDATE knowledge_qa_turns
SET current_attempt_id=$3, assistant_message_id='', status='queued',
    error_code='', error_detail='', debug_model_output=NULL, queued_at=$4, started_at=NULL,
    completed_at=NULL, updated_at=$4
WHERE workspace_id=$1 AND id=$2
`, session.WorkspaceID, turn.ID, params.AttemptID, params.Now); err != nil {
			return fmt.Errorf("queue Knowledge Q&A retry: %w", err)
		}
		if _, err := tx.Exec(ctx, `
UPDATE knowledge_qa_sessions
SET active_turn_id=$2, last_activity_at=$3, updated_at=$3
WHERE id=$1
`, session.ID, turn.ID, params.Now); err != nil {
			return fmt.Errorf("activate Knowledge Q&A retry: %w", err)
		}
		turn, err = loadTurn(ctx, tx, session.WorkspaceID, turn.ID)
		return err
	})
	return turn, err
}

func latestTerminalKnowledgeQATurn(ctx context.Context, tx coredb.Tx, workspaceID, sessionID string) (string, error) {
	var turnID string
	if err := tx.QueryRow(ctx, `
SELECT id FROM knowledge_qa_turns
WHERE workspace_id=$1 AND session_id=$2
  AND status IN ('completed','no_answer','failed','canceled')
ORDER BY queued_at DESC, id DESC
LIMIT 1
`, workspaceID, sessionID).Scan(&turnID); err != nil {
		if coredb.IsNotFound(err) {
			return "", ErrInvalidState
		}
		return "", fmt.Errorf("load latest terminal Knowledge Q&A turn: %w", err)
	}
	return turnID, nil
}

func retryableKnowledgeQATurn(ctx context.Context, tx coredb.Tx, turn Turn) (bool, error) {
	switch turn.Status {
	case "failed", "canceled", "no_answer":
		return true, nil
	case "completed":
		var warningCount int64
		if err := tx.QueryRow(ctx, `
SELECT COUNT(*) FROM knowledge_qa_events
WHERE attempt_id=$1 AND turn_id=$2 AND kind='diagnostic' AND status='warning'
`, turn.CurrentAttemptID, turn.ID).Scan(&warningCount); err != nil {
			return false, fmt.Errorf("check Knowledge Q&A retry diagnostics: %w", err)
		}
		return warningCount > 0, nil
	default:
		return false, nil
	}
}

func (store *Store) ApplyTurnSnapshot(ctx context.Context, workspaceID, turnID, attemptID string, snapshot TurnSnapshot, now time.Time) error {
	now = normalizeNow(now)
	snapshot = normalizeTurnSnapshot(snapshot)
	if snapshot.EffectiveModel == "" || snapshot.EffectiveReasonLevel == "" ||
		snapshot.OrganizationPolicyVersion <= 0 ||
		snapshot.RuntimeVersion == "" || snapshot.PromptVersion == "" {
		return fmt.Errorf("complete Knowledge Q&A policy and runtime snapshot is required")
	}
	rootsJSON, err := marshalAuthorizedRoots(snapshot.AuthorizedRoots)
	if err != nil {
		return fmt.Errorf("marshal Knowledge Q&A roots: %w", err)
	}
	return coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		result, err := tx.Exec(ctx, `
UPDATE knowledge_qa_turns
SET authorized_roots_json=$4, effective_model=$5,
	    effective_reason_level=$6, organization_policy_version=$7,
	    runtime_version=$8, prompt_version=$9, updated_at=$10
WHERE workspace_id=$1 AND id=$2 AND current_attempt_id=$3
  AND status IN ('preparing','generating')
`, strings.TrimSpace(workspaceID), strings.TrimSpace(turnID),
			strings.TrimSpace(attemptID), rootsJSON, snapshot.EffectiveModel,
			snapshot.EffectiveReasonLevel, snapshot.OrganizationPolicyVersion,
			snapshot.RuntimeVersion, snapshot.PromptVersion, now)
		if err != nil {
			return fmt.Errorf("snapshot Knowledge Q&A turn: %w", err)
		}
		if result.RowsAffected() == 1 {
			return nil
		}
		if result.RowsAffected() != 0 {
			return ErrInvalidState
		}

		// TiDB reports changed rows rather than matched rows. A retry can
		// legitimately persist the same snapshot and second-precision timestamp
		// as the previous attempt, so distinguish that idempotent replay from a
		// stale attempt or state transition.
		var persisted TurnSnapshot
		var persistedRootsJSON []byte
		err = tx.QueryRow(ctx, `
SELECT authorized_roots_json, effective_model, effective_reason_level,
       organization_policy_version, runtime_version, prompt_version
FROM knowledge_qa_turns
WHERE workspace_id=$1 AND id=$2 AND current_attempt_id=$3
  AND status IN ('preparing','generating')
FOR UPDATE
`, strings.TrimSpace(workspaceID), strings.TrimSpace(turnID),
			strings.TrimSpace(attemptID)).Scan(
			&persistedRootsJSON, &persisted.EffectiveModel,
			&persisted.EffectiveReasonLevel, &persisted.OrganizationPolicyVersion,
			&persisted.RuntimeVersion, &persisted.PromptVersion,
		)
		if err != nil {
			if coredb.IsNotFound(err) {
				return ErrInvalidState
			}
			return fmt.Errorf("verify Knowledge Q&A turn snapshot: %w", err)
		}
		persisted.AuthorizedRoots, err = unmarshalAuthorizedRoots(persistedRootsJSON)
		if err != nil {
			return fmt.Errorf("decode persisted Knowledge Q&A roots: %w", err)
		}
		persisted = normalizeTurnSnapshot(persisted)
		if !slices.Equal(persisted.AuthorizedRoots, snapshot.AuthorizedRoots) ||
			persisted.EffectiveModel != snapshot.EffectiveModel ||
			persisted.EffectiveReasonLevel != snapshot.EffectiveReasonLevel ||
			persisted.OrganizationPolicyVersion != snapshot.OrganizationPolicyVersion ||
			persisted.RuntimeVersion != snapshot.RuntimeVersion ||
			persisted.PromptVersion != snapshot.PromptVersion {
			return ErrInvalidState
		}
		return nil
	})
}

func (store *Store) TransitionTurn(ctx context.Context, workspaceID, turnID, attemptID, from, to string, now time.Time) error {
	now = normalizeNow(now)
	if !validTurnTransition(from, to) {
		return ErrInvalidState
	}
	result, err := store.db.Exec(ctx, `
UPDATE knowledge_qa_turns AS turn
JOIN knowledge_qa_turn_attempts AS attempt
  ON attempt.turn_id=turn.id AND attempt.id=$3
SET turn.status=$5,
    turn.started_at=CASE WHEN $5='preparing' THEN COALESCE(turn.started_at,$6) ELSE turn.started_at END,
    turn.updated_at=$6,
    attempt.status=$5,
    attempt.started_at=CASE WHEN $5='preparing' THEN COALESCE(attempt.started_at,$6) ELSE attempt.started_at END,
    attempt.updated_at=$6
WHERE turn.workspace_id=$1 AND turn.id=$2 AND turn.current_attempt_id=$3
  AND turn.status=$4 AND attempt.status=$4
`, strings.TrimSpace(workspaceID), strings.TrimSpace(turnID),
		strings.TrimSpace(attemptID), strings.TrimSpace(from),
		strings.TrimSpace(to), now)
	if err != nil {
		return fmt.Errorf("transition Knowledge Q&A turn: %w", err)
	}
	if result.RowsAffected() == 0 {
		return ErrInvalidState
	}
	return nil
}

// ClaimQueuedAttempt atomically enforces per-user and per-organization running
// limits at the moment a worker leaves the durable queue. A saturated worker
// leaves the attempt queued; credits never change queue priority.
func (store *Store) ClaimQueuedAttempt(ctx context.Context, workspaceID, turnID, attemptID string, now time.Time) (Turn, error) {
	now = normalizeNow(now)
	var claimed Turn
	err := coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		var organizationID, requesterUserID, currentAttemptID, status string
		if err := tx.QueryRow(ctx, `
SELECT organization_id, requester_user_id, current_attempt_id, status
FROM knowledge_qa_turns
WHERE workspace_id=$1 AND id=$2
FOR UPDATE
`, strings.TrimSpace(workspaceID), strings.TrimSpace(turnID)).Scan(
			&organizationID, &requesterUserID, &currentAttemptID, &status,
		); err != nil {
			if coredb.IsNotFound(err) {
				return ErrNotFound
			}
			return fmt.Errorf("lock queued Knowledge Q&A turn: %w", err)
		}
		if status != "queued" || currentAttemptID != strings.TrimSpace(attemptID) {
			return ErrInvalidState
		}
		if err := lockAdmissionPrincipals(ctx, tx, organizationID, requesterUserID); err != nil {
			return err
		}
		var userRunning, organizationRunning int
		if err := tx.QueryRow(ctx, `
SELECT
  COALESCE(SUM(CASE WHEN requester_user_id=$2 THEN 1 ELSE 0 END),0),
  COUNT(*)
FROM knowledge_qa_turns
WHERE organization_id=$1
  AND status IN ('preparing','generating')
`, organizationID, requesterUserID).Scan(&userRunning, &organizationRunning); err != nil {
			return fmt.Errorf("load Knowledge Q&A running capacity: %w", err)
		}
		if userRunning >= MaxUserRunningTurns ||
			organizationRunning >= MaxOrganizationRunningTurns {
			return ErrQueueFull
		}
		result, err := tx.Exec(ctx, `
UPDATE knowledge_qa_turns AS turn
JOIN knowledge_qa_turn_attempts AS attempt
  ON attempt.turn_id=turn.id AND attempt.id=$3
SET turn.status='preparing', turn.started_at=COALESCE(turn.started_at,$4),
    turn.updated_at=$4, attempt.status='preparing',
    attempt.started_at=COALESCE(attempt.started_at,$4), attempt.updated_at=$4
WHERE turn.workspace_id=$1 AND turn.id=$2 AND turn.current_attempt_id=$3
  AND turn.status='queued' AND attempt.status='queued'
`, workspaceID, turnID, attemptID, now)
		if err != nil {
			return fmt.Errorf("claim queued Knowledge Q&A attempt: %w", err)
		}
		if result.RowsAffected() == 0 {
			return ErrInvalidState
		}
		claimed, err = loadTurn(ctx, tx, workspaceID, turnID)
		return err
	})
	return claimed, err
}

func (store *Store) CompleteAttempt(ctx context.Context, workspaceID, turnID, attemptID string, completion Completion) error {
	if completion.Status == "completed" && strings.TrimSpace(completion.AssistantMessageID) == "" {
		return fmt.Errorf("completed Knowledge Q&A turn requires an assistant message")
	}
	return store.completeAttempt(ctx, workspaceID, turnID, attemptID, completion, nil, nil)
}

// CompleteAttemptWithAssistant commits the stable citations, formal answer
// message, and completed turn in one database transaction. Stop and completion
// therefore have a single serialization point: a canceled turn cannot leave
// orphaned citations or an assistant answer in the Human-visible thread.
func (store *Store) CompleteAttemptWithAssistant(
	ctx context.Context,
	workspaceID, turnID, attemptID string,
	completion Completion,
	citations []Citation,
	appendAssistant func(context.Context, coredb.Tx) (string, error),
) error {
	if completion.Status != "completed" || appendAssistant == nil {
		return ErrInvalidState
	}
	preparedCitations, err := PrepareCitations(citations)
	if err != nil {
		return err
	}
	completion.AssistantMessageID = ""
	return store.completeAttempt(
		ctx,
		workspaceID,
		turnID,
		attemptID,
		completion,
		preparedCitations,
		appendAssistant,
	)
}

func (store *Store) completeAttempt(
	ctx context.Context,
	workspaceID, turnID, attemptID string,
	completion Completion,
	citations []Citation,
	appendAssistant func(context.Context, coredb.Tx) (string, error),
) error {
	completion, usageJSON, settlementRefsJSON, err := prepareCompletion(completion)
	if err != nil {
		return err
	}
	return coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		var sessionID, currentStatus string
		if err := tx.QueryRow(ctx, `
SELECT session_id, status FROM knowledge_qa_turns
WHERE workspace_id=$1 AND id=$2 AND current_attempt_id=$3
FOR UPDATE
`, workspaceID, turnID, attemptID).Scan(&sessionID, &currentStatus); err != nil {
			if coredb.IsNotFound(err) {
				return ErrNotFound
			}
			return fmt.Errorf("lock Knowledge Q&A turn completion: %w", err)
		}
		if isTerminalTurnStatus(currentStatus) {
			return completeTerminalAttemptAccountingTx(
				ctx, tx, turnID, attemptID, completion, usageJSON, settlementRefsJSON,
			)
		}
		if appendAssistant != nil {
			if _, err := tx.Exec(ctx, `DELETE FROM knowledge_qa_citations WHERE turn_id=$1`, turnID); err != nil {
				return fmt.Errorf("replace Knowledge Q&A citations: %w", err)
			}
			if err := insertCitations(
				ctx, tx, workspaceID, turnID, attemptID, citations, completion.Now,
			); err != nil {
				return err
			}
			assistantMessageID, err := appendAssistant(ctx, tx)
			if err != nil {
				return err
			}
			completion.AssistantMessageID = strings.TrimSpace(assistantMessageID)
			if completion.AssistantMessageID == "" {
				return fmt.Errorf("completed Knowledge Q&A turn requires an assistant message")
			}
		}
		if completion.Status != "completed" {
			if _, err := tx.Exec(ctx, `
DELETE FROM knowledge_qa_citations
WHERE turn_id=$1 AND attempt_id=$2
`, turnID, attemptID); err != nil {
				return fmt.Errorf("discard non-completed Knowledge Q&A citations: %w", err)
			}
		}
		if _, err := tx.Exec(ctx, `
UPDATE knowledge_qa_turn_attempts
SET status=$2, error_code=$3, error_detail=$4, usage_json=CAST($5 AS JSON),
    settlement_ref=$6, settlement_refs_json=CAST($7 AS JSON), completed_at=$8, updated_at=$8
WHERE id=$1
`, attemptID, completion.Status, strings.TrimSpace(completion.ErrorCode),
			strings.TrimSpace(completion.ErrorDetail), usageJSON,
			completion.SettlementRef, settlementRefsJSON, completion.Now); err != nil {
			return fmt.Errorf("complete Knowledge Q&A attempt: %w", err)
		}
		if _, err := tx.Exec(ctx, `
UPDATE knowledge_qa_turns
SET assistant_message_id=$3, status=$4, error_code=$5, error_detail=$6,
    debug_model_output=NULLIF($7,''), completed_at=$8, updated_at=$8
WHERE workspace_id=$1 AND id=$2
`, workspaceID, turnID, strings.TrimSpace(completion.AssistantMessageID),
			completion.Status, strings.TrimSpace(completion.ErrorCode),
			strings.TrimSpace(completion.ErrorDetail), completion.DebugModelOutput, completion.Now); err != nil {
			return fmt.Errorf("complete Knowledge Q&A turn: %w", err)
		}
		if _, err := tx.Exec(ctx, `
UPDATE knowledge_qa_sessions
SET active_turn_id='', last_activity_at=$2, updated_at=$2
WHERE id=$1 AND active_turn_id=$3
`, sessionID, completion.Now, turnID); err != nil {
			return fmt.Errorf("release Knowledge Q&A session turn: %w", err)
		}
		if completion.RuntimeContext != nil {
			if err := applyRuntimeContextUpdateTx(
				ctx, tx, workspaceID, sessionID, turnID, *completion.RuntimeContext, completion.Now,
			); err != nil {
				return err
			}
		}
		if completion.RecoverAccounting {
			if err := markAccountingRecoveryPendingTx(ctx, tx, attemptID, completion.Now); err != nil {
				return err
			}
		}
		return nil
	})
}

func completeTerminalAttemptAccountingTx(
	ctx context.Context,
	tx coredb.Tx,
	turnID, attemptID string,
	completion Completion,
	usageJSON, settlementRefsJSON []byte,
) error {
	if len(completion.Usage) > 0 || completion.SettlementRef != "" {
		if _, err := tx.Exec(ctx, `
UPDATE knowledge_qa_turn_attempts
SET usage_json=CAST($2 AS JSON), settlement_ref=$3,
    settlement_refs_json=CAST($4 AS JSON), updated_at=$5
WHERE id=$1 AND turn_id=$6
`, attemptID, usageJSON, completion.SettlementRef, settlementRefsJSON,
			completion.Now, turnID); err != nil {
			return fmt.Errorf("record terminal Knowledge Q&A usage: %w", err)
		}
		if err := resolveAccountingRecoveryTx(ctx, tx, attemptID, completion.Now); err != nil {
			return err
		}
	}
	if completion.RecoverAccounting {
		return markAccountingRecoveryPendingTx(ctx, tx, attemptID, completion.Now)
	}
	return nil
}

func (store *Store) StopTurnForOwner(ctx context.Context, workspaceID, ownerUserID, turnID string, now time.Time) (Turn, error) {
	now = normalizeNow(now)
	var turn Turn
	err := coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		var sessionID, attemptID, status string
		if err := tx.QueryRow(ctx, `
SELECT turn.session_id, turn.current_attempt_id, turn.status
FROM knowledge_qa_turns AS turn
JOIN knowledge_qa_sessions AS session ON session.id=turn.session_id
WHERE turn.workspace_id=$1 AND turn.id=$2 AND session.owner_user_id=$3
FOR UPDATE
`, strings.TrimSpace(workspaceID), strings.TrimSpace(turnID),
			strings.TrimSpace(ownerUserID)).Scan(&sessionID, &attemptID, &status); err != nil {
			if coredb.IsNotFound(err) {
				return ErrNotFound
			}
			return fmt.Errorf("lock Knowledge Q&A turn stop: %w", err)
		}
		if !isTerminalTurnStatus(status) {
			if _, err := tx.Exec(ctx, `
UPDATE knowledge_qa_turns
SET status='canceled', error_code='', error_detail='', debug_model_output=NULL,
    assistant_message_id='', completed_at=$3, updated_at=$3
WHERE workspace_id=$1 AND id=$2
`, workspaceID, turnID, now); err != nil {
				return fmt.Errorf("cancel Knowledge Q&A turn: %w", err)
			}
			attemptUpdate := `
UPDATE knowledge_qa_turn_attempts
SET status='canceled', error_code='', error_detail='',
    completed_at=$2, updated_at=$2 WHERE id=$1
`
			attemptArgs := []any{attemptID, now}
			if status == "generating" {
				attemptUpdate = `
UPDATE knowledge_qa_turn_attempts
SET status='canceled', error_code='', error_detail='',
    accounting_recovery_status='pending', accounting_recovery_attempts=0,
    accounting_recovery_next_at=$2, accounting_recovery_deadline_at=$3,
    completed_at=$2, updated_at=$2 WHERE id=$1
`
				attemptArgs = append(attemptArgs, now.Add(AccountingRecoveryWindow))
			}
			if _, err := tx.Exec(ctx, attemptUpdate, attemptArgs...); err != nil {
				return fmt.Errorf("cancel Knowledge Q&A attempt: %w", err)
			}
			if _, err := tx.Exec(ctx, `
DELETE FROM knowledge_qa_citations
WHERE turn_id=$1 AND attempt_id=$2
`, turnID, attemptID); err != nil {
				return fmt.Errorf("discard canceled Knowledge Q&A citations: %w", err)
			}
			if _, err := tx.Exec(ctx, `
UPDATE knowledge_qa_sessions
SET active_turn_id='', last_activity_at=$2, updated_at=$2
WHERE id=$1 AND active_turn_id=$3
`, sessionID, now, turnID); err != nil {
				return fmt.Errorf("release canceled Knowledge Q&A turn: %w", err)
			}
		}
		var err error
		turn, err = loadTurn(ctx, tx, workspaceID, turnID)
		return err
	})
	return turn, err
}

func (store *Store) AppendEvent(ctx context.Context, workspaceID, turnID, attemptID string, activity protocol.RuntimeActivity, payload map[string]any, now time.Time) error {
	now = normalizeNow(now)
	if activity.Sequence <= 0 {
		return ErrEventSequence
	}
	payloadJSON, err := model.MarshalJSONObject("Knowledge Q&A event", nonNilMap(payload))
	if err != nil {
		return fmt.Errorf("marshal Knowledge Q&A event: %w", err)
	}
	_, err = store.db.Exec(ctx, `
INSERT INTO knowledge_qa_events (
  turn_id, sequence, workspace_id, attempt_id, kind, status,
  summary, tool_name, payload_json, created_at
) SELECT $1,$2,$3,$4,$5,$6,$7,$8,$9,$10
FROM knowledge_qa_turns
WHERE id=$1 AND workspace_id=$3 AND current_attempt_id=$4
`, strings.TrimSpace(turnID), activity.Sequence, strings.TrimSpace(workspaceID),
		strings.TrimSpace(attemptID), strings.TrimSpace(activity.Kind),
		strings.TrimSpace(activity.Status), strings.TrimSpace(activity.Summary),
		strings.TrimSpace(activity.ToolName), payloadJSON, now)
	if err != nil {
		return fmt.Errorf("append Knowledge Q&A event: %w", err)
	}
	return nil
}

func (store *Store) AppendNextEvent(ctx context.Context, workspaceID, turnID, attemptID string, activity protocol.RuntimeActivity, payload map[string]any, now time.Time) error {
	now = normalizeNow(now)
	payloadJSON, err := model.MarshalJSONObject("Knowledge Q&A event", nonNilMap(payload))
	if err != nil {
		return fmt.Errorf("marshal Knowledge Q&A event: %w", err)
	}
	return coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		var currentAttemptID, status string
		if err := tx.QueryRow(ctx, `
SELECT current_attempt_id, status FROM knowledge_qa_turns
WHERE workspace_id=$1 AND id=$2
FOR UPDATE
`, strings.TrimSpace(workspaceID), strings.TrimSpace(turnID)).Scan(&currentAttemptID, &status); err != nil {
			if coredb.IsNotFound(err) {
				return ErrNotFound
			}
			return fmt.Errorf("lock Knowledge Q&A event sequence: %w", err)
		}
		if currentAttemptID != strings.TrimSpace(attemptID) || isTerminalTurnStatus(status) {
			return ErrInvalidState
		}
		var sequence int64
		if err := tx.QueryRow(ctx, `
SELECT COALESCE(MAX(sequence),0)+1 FROM knowledge_qa_events WHERE turn_id=$1
`, strings.TrimSpace(turnID)).Scan(&sequence); err != nil {
			return fmt.Errorf("allocate Knowledge Q&A event sequence: %w", err)
		}
		_, err := tx.Exec(ctx, `
INSERT INTO knowledge_qa_events (
  turn_id, sequence, workspace_id, attempt_id, kind, status,
  summary, tool_name, payload_json, created_at
) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10)
`, strings.TrimSpace(turnID), sequence, strings.TrimSpace(workspaceID),
			strings.TrimSpace(attemptID), strings.TrimSpace(activity.Kind),
			strings.TrimSpace(activity.Status), strings.TrimSpace(activity.Summary),
			strings.TrimSpace(activity.ToolName), payloadJSON, now)
		if err != nil {
			return fmt.Errorf("append Knowledge Q&A event: %w", err)
		}
		return nil
	})
}