package knowledgeqa

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"strings"
	"testing"
	"time"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"
)

func TestNormalizeCreateSessionScopesAndTrims(t *testing.T) {
	now := time.Date(2026, 7, 16, 10, 0, 0, 0, time.FixedZone("CST", 8*60*60))
	params := normalizeCreateSession(CreateSessionParams{
		ID: " kqs_1 ", WorkspaceID: " wsp_1 ", OrganizationID: " org_1 ",
		OwnerUserID: " usr_1 ", OwnerMemberID: " mem_1 ",
		RootMessageID: " msg_1 ", ThreadChannelID: " thr_1 ",
		RequestedKnowledgeBaseIDs: []string{" kb_2 ", "kb_1", "kb_2", " "},
		RequestedSourceIDs:        []string{" src_2 ", "src_1", "src_2"},
		RequestedModel:            " agentteam/model-a ",
		RequestedReasonLevel:      " medium ",
		Now:                       now,
	})
	if err := validateCreateSession(params); err != nil {
		t.Fatalf("validateCreateSession() error = %v", err)
	}
	if params.ID != "kqs_1" || params.WorkspaceID != "wsp_1" || params.OrganizationID != "org_1" {
		t.Fatalf("normalizeCreateSession() did not trim stable identifiers: %#v", params)
	}
	if got, want := params.RequestedKnowledgeBaseIDs, []string{"kb_2", "kb_1"}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] {
		t.Fatalf("RequestedKnowledgeBaseIDs = %#v, want %#v", got, want)
	}
	if got, want := params.RequestedSourceIDs, []string{"src_2", "src_1"}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] {
		t.Fatalf("RequestedSourceIDs = %#v, want %#v", got, want)
	}
	if params.Now.Location() != time.UTC {
		t.Fatalf("Now location = %v, want UTC", params.Now.Location())
	}
}

func TestLoadOrganizationPolicyDefaultsToDisabledWithoutMaterializing(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	mock.ExpectQuery(`(?s)SELECT organization.id, COALESCE\(policy.enabled,FALSE\).*LEFT JOIN organization_knowledge_qa_policy`).
		WithArgs("org_1").
		WillReturnRows(sqlmock.NewRows([]string{
			"organization_id", "enabled", "policy_version", "updated_by_user_id", "created_at", "updated_at",
		}).AddRow("org_1", false, int64(0), "", nil, nil))

	policy, err := New(coredb.WrapSQLDB(sqlDB)).LoadOrganizationPolicy(t.Context(), "org_1")
	if err != nil {
		t.Fatal(err)
	}
	if policy.Enabled || policy.PolicyVersion != 0 {
		t.Fatalf("policy = %#v, want sparse disabled default", policy)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestEnforceSubmissionAdmissionHasNoOrganizationQueueCeiling(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	now := time.Date(2026, 8, 2, 4, 0, 0, 0, time.UTC)
	mock.ExpectBegin()
	// Organization admission still observes active execution capacity, but
	// intentionally does not count the durable organization-wide queued backlog.
	mock.ExpectQuery(`(?s)SELECT\s+COALESCE\(SUM\(CASE WHEN requester_user_id=\? AND status='queued'.*COALESCE\(SUM\(CASE WHEN requester_user_id=\? AND status IN \('preparing','generating'\).*COALESCE\(SUM\(CASE WHEN status IN \('preparing','generating'\).*FROM knowledge_qa_turns WHERE organization_id=\?`).
		WillReturnRows(sqlmock.NewRows([]string{
			"user_queued", "user_running", "organization_running",
		}).AddRow(MaxUserQueuedTurns-1, MaxUserRunningTurns, MaxOrganizationRunningTurns))
	mock.ExpectQuery(`(?s)FROM knowledge_qa_turn_attempts AS attempt.*attempt.organization_id=\? AND attempt.queued_at >= \?`).
		WillReturnRows(sqlmock.NewRows([]string{
			"user_recent", "organization_recent",
		}).AddRow(MaxUserSubmissionsPerMinute-1, MaxOrgSubmissionsPerMinute-1))
	mock.ExpectCommit()

	err = coredb.WithTx(t.Context(), coredb.WrapSQLDB(sqlDB), func(tx coredb.Tx) error {
		return enforceSubmissionAdmission(t.Context(), tx, "org_1", "usr_1", now)
	})
	if err != nil {
		t.Fatalf("enforceSubmissionAdmission() error = %v, want admitted", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestListExpiredRunningAttemptsExcludesQueueWait(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	cutoff := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	mock.ExpectQuery(`(?s)SELECT workspace_id, id, current_attempt_id, status.*started_at IS NOT NULL AND started_at <=`).
		WithArgs(cutoff, 16).
		WillReturnRows(sqlmock.NewRows([]string{
			"workspace_id", "id", "current_attempt_id", "status",
		}).AddRow("wsp_1", "kqt_1", "kqa_1", "generating"))

	store := New(coredb.WrapSQLDB(sqlDB))
	items, err := store.ListExpiredRunningAttempts(t.Context(), cutoff, 16)
	if err != nil {
		t.Fatal(err)
	}
	if len(items) != 1 || items[0].AttemptID != "kqa_1" || items[0].Status != "generating" {
		t.Fatalf("items = %#v", items)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestListPendingAccountingRecoveriesUsesBoundedDueScan(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC)
	startedAt := now.Add(-time.Minute)
	deadline := now.Add(time.Hour)
	mock.ExpectQuery(`(?s)FROM knowledge_qa_turn_attempts AS attempt.*accounting_recovery_status='pending'.*accounting_recovery_next_at <=.*LIMIT`).
		WithArgs(now, 8).
		WillReturnRows(sqlmock.NewRows([]string{
			"workspace_id", "organization_id", "requester_user_id", "turn_id",
			"attempt_id", "attempts", "started_at", "deadline_at",
		}).AddRow("wsp_1", "org_1", "usr_1", "kqt_1", "kqa_1", int64(2), startedAt, deadline))

	items, err := New(coredb.WrapSQLDB(sqlDB)).ListPendingAccountingRecoveries(t.Context(), now, 8)
	if err != nil {
		t.Fatal(err)
	}
	if len(items) != 1 || items[0].AttemptID != "kqa_1" || items[0].Attempts != 2 {
		t.Fatalf("items = %#v", items)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestResolveAccountingRecoveryMergesWithoutReopeningTerminalTurn(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC)
	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)SELECT status, accounting_recovery_status, usage_json.*FOR UPDATE`).
		WithArgs("wsp_1", "kqt_1", "kqa_1").
		WillReturnRows(sqlmock.NewRows([]string{
			"status", "accounting_recovery_status", "usage_json", "settlement_ref", "settlement_refs_json",
		}).AddRow("canceled", "pending", []byte(`{"input_tokens":10}`), "usage_1", []byte(`["usage_1"]`)))
	mock.ExpectExec(`(?s)UPDATE knowledge_qa_turn_attempts.*accounting_recovery_status='resolved'`).
		WithArgs(
			[]byte(`{"input_tokens":10,"output_tokens":4}`),
			"usage_2",
			[]byte(`["usage_1","usage_2"]`),
			now,
			"kqa_1",
		).
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectCommit()

	err = New(coredb.WrapSQLDB(sqlDB)).ResolveAccountingRecovery(t.Context(), AccountingRecoveryResolution{
		WorkspaceID: "wsp_1", TurnID: "kqt_1", AttemptID: "kqa_1",
		Usage:         map[string]any{"input_tokens": float64(7), "output_tokens": float64(4)},
		SettlementRef: "usage_2", SettlementRefs: []string{"usage_1", "usage_2"}, Now: now,
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestResolveAccountingRecoveryRejectsSettlementLineageOutsideStorageEnvelope(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()
	now := time.Date(2026, 8, 3, 13, 0, 0, 0, time.UTC)
	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)SELECT status, accounting_recovery_status, usage_json.*FOR UPDATE`).
		WithArgs("wsp_1", "kqt_1", "kqa_1").
		WillReturnRows(sqlmock.NewRows([]string{
			"status", "accounting_recovery_status", "usage_json", "settlement_ref", "settlement_refs_json",
		}).AddRow("failed", "pending", []byte(`{}`), "", []byte(`[]`)))
	mock.ExpectRollback()

	err = New(coredb.WrapSQLDB(sqlDB)).ResolveAccountingRecovery(t.Context(), AccountingRecoveryResolution{
		WorkspaceID: "wsp_1", TurnID: "kqt_1", AttemptID: "kqa_1",
		SettlementRefs: []string{
			strings.Repeat("x", maxPersistedSettlementReferenceBytes+1),
		},
		Now: now,
	})

	require.ErrorContains(t, err, "storage boundary")
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestResolveAccountingRecoveryDecodesTiDBOpaqueJSONAndRewritesDocuments(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	now := time.Date(2026, 8, 2, 12, 0, 0, 0, time.UTC)
	usageJSON, err := json.Marshal("base64:type253:" + base64.StdEncoding.EncodeToString([]byte(`{"input_tokens":10}`)))
	if err != nil {
		t.Fatal(err)
	}
	refsJSON, err := json.Marshal("base64:type253:" + base64.StdEncoding.EncodeToString([]byte(`["usage_1"]`)))
	if err != nil {
		t.Fatal(err)
	}

	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)SELECT status, accounting_recovery_status, usage_json.*FOR UPDATE`).
		WithArgs("wsp_1", "kqt_1", "kqa_1").
		WillReturnRows(sqlmock.NewRows([]string{
			"status", "accounting_recovery_status", "usage_json", "settlement_ref", "settlement_refs_json",
		}).AddRow("failed", "pending", usageJSON, "usage_1", refsJSON))
	mock.ExpectExec(`(?s)SET usage_json=CAST\(\? AS JSON\).*settlement_refs_json=CAST\(\? AS JSON\).*accounting_recovery_status='resolved'`).
		WithArgs(
			[]byte(`{"input_tokens":10,"output_tokens":4}`),
			"usage_2",
			[]byte(`["usage_1","usage_2"]`),
			now,
			"kqa_1",
		).
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectCommit()

	err = New(coredb.WrapSQLDB(sqlDB)).ResolveAccountingRecovery(t.Context(), AccountingRecoveryResolution{
		WorkspaceID: "wsp_1", TurnID: "kqt_1", AttemptID: "kqa_1",
		Usage:         map[string]any{"input_tokens": float64(7), "output_tokens": float64(4)},
		SettlementRef: "usage_2", SettlementRefs: []string{"usage_1", "usage_2"}, Now: now,
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestValidateSessionThreadRequiresOwnerPrivateKBChannel(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	mock.ExpectQuery(`(?s)FROM scopes AS thread.*JOIN scopes AS parent.*parent_profile\.visibility='private'.*parent_profile\.subsystem=.*parent_profile\.owner_member_id=.*NOT EXISTS.*JOIN members AS other_human.*other_human\.id<>root\.author_member_id`).
		WithArgs("wsp_1", "thr_1", "msg_1", "mem_1", "org_1", "kb", "mem_1").
		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))

	store := New(coredb.WrapSQLDB(sqlDB))
	err = store.validateSessionThread(t.Context(), CreateSessionParams{
		WorkspaceID: "wsp_1", ThreadChannelID: "thr_1",
		RootMessageID: "msg_1", OwnerMemberID: "mem_1",
		OrganizationID: "org_1",
	})
	if !errors.Is(err, ErrForbidden) {
		t.Fatalf("validateSessionThread() error = %v, want ErrForbidden", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestAppendNextEventRejectsTerminalTurn(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)SELECT current_attempt_id, status FROM knowledge_qa_turns.*FOR UPDATE`).
		WithArgs("wsp_1", "kqt_1").
		WillReturnRows(sqlmock.NewRows([]string{"current_attempt_id", "status"}).
			AddRow("kqa_1", "canceled"))
	mock.ExpectRollback()

	store := New(coredb.WrapSQLDB(sqlDB))
	err = store.AppendNextEvent(
		t.Context(), "wsp_1", "kqt_1", "kqa_1",
		protocol.RuntimeActivity{Kind: "tool", Status: "completed"},
		nil, time.Now(),
	)
	if !errors.Is(err, ErrInvalidState) {
		t.Fatalf("AppendNextEvent() error = %v, want ErrInvalidState", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestCompleteAttemptWithAssistantCommitsMessageAtTerminalFence(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	now := time.Date(2026, 7, 25, 10, 5, 0, 0, time.UTC)
	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)SELECT session_id, status FROM knowledge_qa_turns.*FOR UPDATE`).
		WithArgs("wsp_1", "kqt_1", "kqa_1").
		WillReturnRows(sqlmock.NewRows([]string{"session_id", "status"}).
			AddRow("kqs_1", "generating"))
	mock.ExpectExec(`DELETE FROM knowledge_qa_citations`).
		WithArgs("kqt_1").
		WillReturnResult(sqlmock.NewResult(0, 0))
	mock.ExpectExec(`INSERT INTO knowledge_qa_citations`).
		WithArgs(
			"kqc_1", "wsp_1", "kqt_1", "kqa_1", CitationKindKnowledge,
			"kb_1", "src_1", "primary", "primary", "sha_1",
			"", "", "docs/design.md", int64(4), int64(6),
			"", "", "", nil, now,
		).
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectExec(`(?s)UPDATE knowledge_qa_turn_attempts.*SET status=.*usage_json=CAST\(\? AS JSON\).*settlement_refs_json=CAST\(\? AS JSON\)`).
		WithArgs(
			"completed",
			"",
			"",
			[]byte(`{"input_tokens":3}`),
			"settlement_2",
			[]byte(`["settlement_1","settlement_2"]`),
			now,
			now,
			"kqa_1",
		).
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectExec(`UPDATE knowledge_qa_turns`).
		WithArgs(
			"msg_answer",
			"completed",
			"",
			"",
			"",
			now,
			now,
			"wsp_1",
			"kqt_1",
		).
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectExec(`UPDATE knowledge_qa_sessions`).
		WithArgs(now, now, "kqs_1", "kqt_1").
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectCommit()

	appended := false
	store := New(coredb.WrapSQLDB(sqlDB))
	err = store.CompleteAttemptWithAssistant(
		t.Context(),
		"wsp_1",
		"kqt_1",
		"kqa_1",
		Completion{
			Status: "completed", Usage: map[string]any{"input_tokens": 3},
			SettlementRef:  "settlement_2",
			SettlementRefs: []string{"settlement_1", "settlement_2"}, Now: now,
		},
		[]Citation{{
			ID: "kqc_1", KnowledgeBaseID: "kb_1", SourceID: "src_1",
			RootKind: "primary", RepoRole: "primary", Revision: "sha_1",
			Path: "docs/design.md", LineStart: 4, LineEnd: 6,
		}},
		func(context.Context, coredb.Tx) (string, error) {
			appended = true
			return "msg_answer", nil
		},
	)
	if err != nil {
		t.Fatal(err)
	}
	if !appended {
		t.Fatal("assistant message appender was not called")
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestCompleteAttemptWithAssistantDoesNotAppendAfterStopWins(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	now := time.Date(2026, 7, 25, 10, 5, 0, 0, time.UTC)
	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)SELECT session_id, status FROM knowledge_qa_turns.*FOR UPDATE`).
		WithArgs("wsp_1", "kqt_1", "kqa_1").
		WillReturnRows(sqlmock.NewRows([]string{"session_id", "status"}).
			AddRow("kqs_1", "canceled"))
	mock.ExpectExec(`(?s)UPDATE knowledge_qa_turn_attempts.*SET usage_json=CAST\(\? AS JSON\).*settlement_refs_json=CAST\(\? AS JSON\)`).
		WithArgs(
			[]byte(`{"input_tokens":5}`),
			"settlement_canceled",
			[]byte(`["settlement_canceled"]`),
			now,
			"kqa_1",
			"kqt_1",
		).
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectExec(`(?s)UPDATE knowledge_qa_turn_attempts.*accounting_recovery_status='resolved'`).
		WithArgs(now, "kqa_1").
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectCommit()

	appended := false
	store := New(coredb.WrapSQLDB(sqlDB))
	err = store.CompleteAttemptWithAssistant(
		t.Context(),
		"wsp_1",
		"kqt_1",
		"kqa_1",
		Completion{
			Status: "completed", Usage: map[string]any{"input_tokens": 5},
			SettlementRef: "settlement_canceled", Now: now,
		},
		[]Citation{{
			ID: "kqc_unused", KnowledgeBaseID: "kb_1", SourceID: "src_1",
			RootKind: "primary", RepoRole: "primary", Revision: "sha_1",
			Path: "docs/design.md", LineStart: 4, LineEnd: 6,
		}},
		func(context.Context, coredb.Tx) (string, error) {
			appended = true
			return "msg_answer", nil
		},
	)
	if err != nil {
		t.Fatal(err)
	}
	if appended {
		t.Fatal("assistant message was appended after the turn was canceled")
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestCompleteTerminalTransportFailureSchedulesAccountingRecovery(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	now := time.Date(2026, 7, 28, 10, 5, 0, 0, time.UTC)
	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)SELECT session_id, status FROM knowledge_qa_turns.*FOR UPDATE`).
		WithArgs("wsp_1", "kqt_1", "kqa_1").
		WillReturnRows(sqlmock.NewRows([]string{"session_id", "status"}).
			AddRow("kqs_1", "canceled"))
	mock.ExpectExec(`(?s)UPDATE knowledge_qa_turn_attempts.*accounting_recovery_status='pending'`).
		WithArgs(now, now.Add(AccountingRecoveryWindow), "kqa_1").
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectCommit()

	err = New(coredb.WrapSQLDB(sqlDB)).CompleteAttempt(
		t.Context(), "wsp_1", "kqt_1", "kqa_1",
		Completion{Status: "canceled", RecoverAccounting: true, Now: now},
	)
	if err != nil {
		t.Fatal(err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestCompleteFailedAttemptDiscardsProvisionalCitations(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	now := time.Date(2026, 7, 25, 10, 5, 0, 0, time.UTC)
	debugModelOutput := `{"model_output":"` + strings.Repeat("full-debug-output-", 5000) + `"}`
	if len(debugModelOutput) <= 64<<10 {
		t.Fatal("debug fixture must exceed the former excerpt limit")
	}
	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)SELECT session_id, status FROM knowledge_qa_turns.*FOR UPDATE`).
		WithArgs("wsp_1", "kqt_1", "kqa_1").
		WillReturnRows(sqlmock.NewRows([]string{"session_id", "status"}).
			AddRow("kqs_1", "generating"))
	mock.ExpectExec(`DELETE FROM knowledge_qa_citations`).
		WithArgs("kqt_1", "kqa_1").
		WillReturnResult(sqlmock.NewResult(0, 2))
	mock.ExpectExec(`UPDATE knowledge_qa_turn_attempts`).
		WithArgs(
			"failed",
			"knowledge_qa_reply_failed",
			"reply failed",
			[]byte(`{}`),
			"",
			[]byte(`[]`),
			now,
			now,
			"kqa_1",
		).
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectExec(`UPDATE knowledge_qa_turns`).
		WithArgs(
			"",
			"failed",
			"knowledge_qa_reply_failed",
			"reply failed",
			debugModelOutput,
			now,
			now,
			"wsp_1",
			"kqt_1",
		).
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectExec(`UPDATE knowledge_qa_sessions`).
		WithArgs(now, now, "kqs_1", "kqt_1").
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectCommit()

	store := New(coredb.WrapSQLDB(sqlDB))
	err = store.CompleteAttempt(t.Context(), "wsp_1", "kqt_1", "kqa_1", Completion{
		Status: "failed", ErrorCode: "knowledge_qa_reply_failed",
		ErrorDetail: "reply failed", DebugModelOutput: debugModelOutput, Now: now,
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestListTurnsForOwnerRestoresOrderedActivitiesAndCompleteLatestDebugOutput(t *testing.T) {
	now := time.Date(2026, 7, 16, 10, 0, 0, 0, time.UTC)
	debugModelOutput := `{"model_output":"` + strings.Repeat("full-debug-output-", 5000) + `"}`
	for _, testCase := range []struct {
		name, status, wantDebug  string
		expectDebug, returnDebug bool
	}{
		{"actual latest failed", "failed", debugModelOutput, true, true},
		{"bounded window tail failed but not latest", "failed", "", true, false},
		{"non-failed window tail", "completed", "", false, false},
	} {
		t.Run(testCase.name, func(t *testing.T) {
			sqlDB, mock, err := sqlmock.New()
			if err != nil {
				t.Fatal(err)
			}
			defer sqlDB.Close()
			mock.ExpectQuery(`SELECT COUNT\(\*\) FROM knowledge_qa_sessions`).
				WithArgs("wsp_1", "usr_1", "thr_1").
				WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
			mock.ExpectQuery(`(?s)SELECT turn\.id.*FROM knowledge_qa_turns`).
				WithArgs("wsp_1", "usr_1", "thr_1", 100).
				WillReturnRows(sqlmock.NewRows([]string{
					"id", "workspace_id", "organization_id", "session_id", "requester_user_id",
					"user_message_id", "assistant_message_id", "question", "idempotency_key",
					"requested_kbs", "requested_sources", "authorized_roots", "requested_model",
					"requested_reason_level", "web_search_enabled", "effective_model", "effective_reason_level",
					"organization_policy_version", "runtime_version", "prompt_version",
					"current_attempt_id", "status", "error_code", "error_detail", "queued_at",
					"started_at", "completed_at", "updated_at",
				}).AddRow(
					"kqt_1", "wsp_1", "org_1", "kqs_1", "usr_1", "msg_q", "", "question",
					"idem_1", []byte(`["kb_1"]`), []byte(`["src_1"]`), []byte(`[]`),
					"model-a", "high", false, "model-a", "high", int64(4), "runtime-v3", "prompt-v3",
					"kqa_1", testCase.status, "knowledge_qa.response_invalid", "internal detail",
					now, now, now, now,
				))
			mock.ExpectQuery(`(?s)SELECT event\.turn_id.*FROM knowledge_qa_events.*current_attempt_id=event\.attempt_id`).
				WithArgs("wsp_1", "usr_1", "thr_1", 100).
				WillReturnRows(sqlmock.NewRows([]string{
					"turn_id", "sequence", "kind", "status", "summary", "tool_name", "payload_json", "created_at",
				}).AddRow("kqt_1", int64(1), "tool", "completed", "visible summary", "history_read", []byte(`{"stage":"history_read","tailTurnCount":2}`), now).
					AddRow("kqt_1", int64(2), "tool", "completed", "", "model", []byte(`{"stage":"exploration_progress","modelCalls":4}`), now).
					AddRow("kqt_1", int64(3), "tool", "completed", "", "grep", []byte(`{"stage":"retrieval_complete","grepCalls":1,"readCalls":0}`), now).
					AddRow("kqt_1", int64(4), "diagnostic", "warning", "", "knowledge_qa.grounding.evidence_unavailable", []byte(`{"stage":"grounding_diagnostic","code":"knowledge_qa.grounding.evidence_unavailable","query":"secret"}`), now))
			if testCase.expectDebug {
				rows := sqlmock.NewRows([]string{"debug_model_output"})
				if testCase.returnDebug {
					rows.AddRow(debugModelOutput)
				}
				mock.ExpectQuery(`(?s)SELECT COALESCE\(turn\.debug_model_output,''\).*turn\.id=\?.*NOT EXISTS.*later\.workspace_id=turn\.workspace_id.*later\.session_id=turn\.session_id.*later\.queued_at > turn\.queued_at.*later\.queued_at=turn\.queued_at AND later\.id > turn\.id`).
					WithArgs("wsp_1", "usr_1", "thr_1", "kqt_1").
					WillReturnRows(rows)
			}
			turns, err := New(coredb.WrapSQLDB(sqlDB)).ListTurnsForOwner(context.Background(), "wsp_1", "usr_1", "thr_1", 100)
			if err != nil {
				t.Fatal(err)
			}
			if len(turns) != 1 || len(turns[0].Activities) != 4 || turns[0].DebugModelOutput != testCase.wantDebug {
				t.Fatalf("turns = %#v, want one turn with four activities and debug %q", turns, testCase.wantDebug)
			}
			if turns[0].Activities[0].Summary != "visible summary" || turns[0].Activities[0].Stage != "history_read" ||
				turns[0].Activities[0].Counts == nil || turns[0].Activities[0].Counts.ContextTurns == nil || *turns[0].Activities[0].Counts.ContextTurns != 2 ||
				turns[0].Activities[1].ModelCalls != 4 || turns[0].Activities[1].Stage != "exploration_progress" ||
				turns[0].Activities[2].ToolName != "grep" || turns[0].Activities[2].Counts == nil || turns[0].Activities[2].Counts.ReadCalls == nil || *turns[0].Activities[2].Counts.ReadCalls != 0 ||
				turns[0].Activities[3].Status != "warning" || turns[0].Activities[3].DiagnosticCode != "knowledge_qa.grounding.evidence_unavailable" {
				t.Fatalf("activities = %#v", turns[0].Activities)
			}
			if err := mock.ExpectationsWereMet(); err != nil {
				t.Fatal(err)
			}
		})
	}
}

func TestLoadCitationForOwnerRequiresCompletedAttempt(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	now := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	mock.ExpectQuery(`(?s)SELECT citation\.id.*turn\.status='completed'.*citation\.attempt_id=turn\.current_attempt_id`).
		WithArgs("wsp_1", "kqc_1", "kqs_1", "wsp_1", "usr_1").
		WillReturnRows(sqlmock.NewRows([]string{
			"id", "workspace_id", "turn_id", "attempt_id", "citation_kind",
			"knowledge_base_id", "source_id", "root_kind", "repo_role", "revision",
			"resolved_ref", "resolved_ref_kind", "path", "line_start", "line_end",
			"url", "title", "snippet", "fetched_at", "created_at",
		}).AddRow(
			"kqc_1", "wsp_1", "kqt_1", "kqa_1", CitationKindKnowledge, "kb_1", "src_1",
			"primary", "primary", "abc123", "release-7.5", "branch", "docs/policy.md",
			int64(10), int64(15), "", "", "", nil, now,
		))

	store := New(coredb.WrapSQLDB(sqlDB))
	citation, err := store.LoadCitationForOwner(t.Context(), "wsp_1", "usr_1", "kqs_1", "kqc_1")
	if err != nil {
		t.Fatal(err)
	}
	if citation.TurnID != "kqt_1" || citation.AttemptID != "kqa_1" ||
		citation.Revision != "abc123" || citation.ResolvedRef != "release-7.5" ||
		citation.ResolvedRefKind != "branch" || citation.Path != "docs/policy.md" {
		t.Fatalf("citation = %#v", citation)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestAppendEventUsesMonotonicTurnSequence(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	mock.ExpectExec(`INSERT INTO knowledge_qa_events`).
		WithArgs(
			"kqt_1", int64(1), "wsp_1", "kqa_1", "tool", "completed",
			"visible summary", "history_read", []byte(`{}`), sqlmock.AnyArg(),
			"kqt_1", "wsp_1", "kqa_1",
		).
		WillReturnResult(sqlmock.NewResult(1, 1))
	store := New(coredb.WrapSQLDB(sqlDB))
	err = store.AppendEvent(context.Background(), "wsp_1", "kqt_1", "kqa_1", protocol.RuntimeActivity{
		ID: "tool:1", Sequence: 1, Kind: "tool", Status: "completed", Summary: "visible summary", ToolName: "history_read",
	}, nil, time.Now())
	if err != nil {
		t.Fatal(err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestValidateCreateTurnRejectsMissingAuthority(t *testing.T) {
	params := normalizeCreateTurn(CreateTurnParams{
		WorkspaceID: "wsp_1", OrganizationID: "org_1", SessionID: "kqs_1",
		UserMessageID: "msg_1", Question: "question", IdempotencyKey: "idem_1",
	})
	if err := validateCreateTurn(params); err == nil {
		t.Fatal("validateCreateTurn() error = nil, want missing owner error")
	}
}

func TestRetryTurnReplaysQueuedAttemptWithSameIdempotencyKey(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	now := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)SELECT id, workspace_id, organization_id.*FROM knowledge_qa_sessions.*FOR UPDATE`).
		WithArgs("wsp_1", "kqs_1").
		WillReturnRows(sqlmock.NewRows([]string{
			"id", "workspace_id", "organization_id", "owner_user_id",
			"owner_member_id", "root_message_id", "thread_channel_id",
			"requested_kbs", "requested_sources", "requested_model", "requested_reason_level",
			"web_search_enabled", "active_turn_id", "status",
			"last_activity_at", "created_at", "updated_at",
		}).AddRow(
			"kqs_1", "wsp_1", "org_1", "usr_1", "mem_1", "msg_root",
			"thr_1", []byte(`[]`), []byte(`[]`), "model-a", "high", false, "kqt_1", "active",
			now, now, now,
		))
	mock.ExpectQuery(`(?s)SELECT turn\.id.*FROM knowledge_qa_turns.*WHERE turn\.workspace_id`).
		WithArgs("wsp_1", "kqt_1").
		WillReturnRows(sqlmock.NewRows([]string{
			"id", "workspace_id", "organization_id", "session_id",
			"requester_user_id", "user_message_id", "assistant_message_id",
			"question", "idempotency_key", "requested_kbs", "requested_sources",
			"authorized_roots", "requested_model", "requested_reason_level",
			"web_search_enabled", "effective_model", "effective_reason_level", "organization_policy_version",
			"runtime_version", "prompt_version", "current_attempt_id", "status",
			"error_code", "error_detail", "queued_at", "started_at",
			"completed_at", "updated_at",
		}).AddRow(
			"kqt_1", "wsp_1", "org_1", "kqs_1", "usr_1", "msg_q", "",
			"question", "initial-key", []byte(`[]`), []byte(`[]`), []byte(`[]`),
			"model-a", "high", false, "", "", int64(0), "", "", "kqa_retry",
			"queued", "", "", now, nil, nil, now,
		))
	mock.ExpectQuery(`(?s)SELECT id FROM knowledge_qa_turn_attempts.*idempotency_key`).
		WithArgs("kqt_1", "retry-key").
		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("kqa_retry"))
	mock.ExpectCommit()

	store := New(coredb.WrapSQLDB(sqlDB))
	turn, err := store.RetryTurn(context.Background(), RetryTurnParams{
		WorkspaceID: "wsp_1", SessionID: "kqs_1", TurnID: "kqt_1",
		RequesterUserID: "usr_1", IdempotencyKey: "retry-key",
		AttemptID: "kqa_unused", Now: now,
	})
	if err != nil {
		t.Fatal(err)
	}
	if turn.Status != "queued" || turn.CurrentAttemptID != "kqa_retry" {
		t.Fatalf("turn = %#v", turn)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestLoadThreadMessageIDsPreservesScopeSequence(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	mock.ExpectQuery(`(?s)SELECT id.*FROM scope_messages.*ORDER BY seq ASC`).
		WithArgs("wsp_1", "thr_1").
		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("msg_1").AddRow("msg_2"))

	messageIDs, err := loadThreadMessageIDs(context.Background(), coredb.WrapSQLDB(sqlDB), " wsp_1 ", " thr_1 ")
	if err != nil {
		t.Fatal(err)
	}
	if got, want := messageIDs, []string{"msg_1", "msg_2"}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] {
		t.Fatalf("messageIDs = %#v, want %#v", got, want)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestLoadPolicyMetricsReturnsContentFreeOrganizationAggregate(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	mock.ExpectQuery(`(?s)SELECT COUNT\(\*\).*FROM knowledge_qa_turns.*WHERE organization_id=.*queued_at >= DATE_SUB`).
		WithArgs("org_1").
		WillReturnRows(sqlmock.NewRows([]string{
			"total", "completed", "no_answer", "failed", "canceled",
			"running", "queued", "average_latency_ms",
		}).AddRow(int64(20), int64(12), int64(3), int64(2), int64(1),
			int64(1), int64(1), float64(1250.5)))
	mock.ExpectQuery(`(?s)SELECT COUNT\(\*\).*JSON_EXTRACT.*FROM knowledge_qa_turn_attempts.*queued_at >= DATE_SUB`).
		WithArgs("org_1").
		WillReturnRows(sqlmock.NewRows([]string{"turns", "input", "cached", "output"}).
			AddRow(int64(2), int64(170), int64(52), int64(40)))
	mock.ExpectQuery(`(?s)SELECT COALESCE\(NULLIF\(error_code,''\).*FROM knowledge_qa_turn_attempts.*queued_at >= DATE_SUB`).
		WithArgs("org_1").
		WillReturnRows(sqlmock.NewRows([]string{"error_code", "count"}).
			AddRow("maas.credits.insufficient", int64(2)).
			AddRow("knowledge_qa_runtime_timeout", int64(1)))
	store := New(coredb.WrapSQLDB(sqlDB))
	metrics, err := store.LoadPolicyMetrics(t.Context(), "org_1")
	if err != nil {
		t.Fatal(err)
	}
	if metrics.WindowDays != 30 || metrics.TotalTurns != 20 || metrics.CompletedTurns != 12 ||
		metrics.NoAnswerTurns != 3 || metrics.SettledUsageTurns != 2 ||
		metrics.InputTokens != 170 || metrics.CachedInputTokens != 52 ||
		metrics.OutputTokens != 40 || len(metrics.FailureTypes) != 2 ||
		metrics.FailureTypes[0].Code != "maas.credits.insufficient" ||
		metrics.AverageTerminalLatencyMS != 1250.5 {
		t.Fatalf("metrics = %#v", metrics)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestResolveEffectivePolicyUsesRequestedHostedSelection(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	now := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	mock.ExpectQuery(`(?s)SELECT organization.id, COALESCE\(policy.enabled,FALSE\).*FROM organizations AS organization`).
		WithArgs("org_1").
		WillReturnRows(sqlmock.NewRows([]string{
			"organization_id", "enabled", "policy_version", "updated_by_user_id", "created_at", "updated_at",
		}).AddRow("org_1", true, int64(7), "usr_org_admin", now, now))

	store := New(coredb.WrapSQLDB(sqlDB))
	effective, err := store.ResolveEffectivePolicy(context.Background(), "org_1", "model-b", "high", "")
	if err != nil {
		t.Fatal(err)
	}
	if effective.Model != "model-b" || effective.ReasonLevel != "high" {
		t.Fatalf("effective policy = %#v", effective)
	}
	if effective.OrganizationPolicyVersion != 7 {
		t.Fatalf("effective policy versions = %#v", effective)
	}
	if effective.WorkloadKind != "knowledge_qa" ||
		effective.DataPolicyRequirementRef != "" ||
		effective.DataPolicyProfileRef != "" ||
		effective.DataPolicyProfileVersion != "" {
		t.Fatalf("effective policy extension fields = %#v", effective)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestResolveEffectivePolicyRejectsNonEmptyDataPolicyRequirementBeforeStorage(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	store := New(coredb.WrapSQLDB(sqlDB))
	_, err = store.ResolveEffectivePolicy(context.Background(), "org_1", "model-a", "medium", "policy_required")
	if !errors.Is(err, protocol.ErrDataPolicyUnsupported) {
		t.Fatalf("ResolveEffectivePolicy() error = %v, want data policy unsupported", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}