package knowledgeqa

import (
	"fmt"
	"strings"
	"testing"
	"time"

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

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

func TestPersistedAccountingDocumentSizeBoundaries(t *testing.T) {
	require.NoError(t, validatePersistedAccountingDocumentSizes(
		make([]byte, maxPersistedAccountingUsageJSONBytes),
		make([]byte, maxPersistedSettlementLineageJSONBytes),
		maxPersistedSettlementReferences,
	))

	tests := []struct {
		name           string
		usageBytes     int
		lineageBytes   int
		referenceCount int
		wantError      string
	}{
		{
			name: "usage", usageBytes: maxPersistedAccountingUsageJSONBytes + 1,
			lineageBytes: 2, wantError: "usage exceeds",
		},
		{
			name: "encoded lineage", usageBytes: 2,
			lineageBytes:   maxPersistedSettlementLineageJSONBytes + 1,
			referenceCount: 1, wantError: "lineage exceeds",
		},
		{
			name: "reference count", usageBytes: 2, lineageBytes: 2,
			referenceCount: maxPersistedSettlementReferences + 1,
			wantError:      "lineage exceeds",
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			err := validatePersistedAccountingDocumentSizes(
				make([]byte, test.usageBytes), make([]byte, test.lineageBytes), test.referenceCount,
			)
			require.ErrorContains(t, err, test.wantError)
		})
	}
}

func TestPrepareCompletionRejectsOversizedSettlementReference(t *testing.T) {
	_, _, _, err := prepareCompletion(Completion{
		Status: "failed",
		SettlementRefs: []string{
			strings.Repeat("x", maxPersistedSettlementReferenceBytes+1),
		},
		Now: time.Now(),
	})
	require.ErrorContains(t, err, "storage boundary")
}

func TestPrepareCompletionRejectsSettlementReferenceCountOutsideStorageEnvelope(t *testing.T) {
	refs := make([]string, maxPersistedSettlementReferences+1)
	for index := range refs {
		refs[index] = fmt.Sprintf("usage_%d", index)
	}
	_, _, _, err := prepareCompletion(Completion{
		Status: "failed", SettlementRefs: refs, Now: time.Now(),
	})
	require.ErrorContains(t, err, "storage boundary")
}

func TestReadRuntimeContextReturnsSummaryAndAllTerminalTailTurns(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()
	queuedAt := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	mock.ExpectQuery(`(?s)SELECT turn.queued_at,.*knowledge_qa_runtime_contexts`).
		WithArgs("wsp_1", "kqt_current", "kqs_1", "wsp_1", "usr_1").
		WillReturnRows(sqlmock.NewRows([]string{
			"queued_at", "generation", "summary", "compacted_through_turn_id", "summary_version",
		}).AddRow(queuedAt, int64(3), "Earlier deployment discussion.", "kqt_4", "knowledge-qa-context-v1"))
	mock.ExpectQuery(`(?s)SELECT turn.id, turn.status, turn.question,.*LEFT JOIN scope_messages`).
		WithArgs("kqt_4", "wsp_1", "kqs_1", queuedAt, queuedAt, "kqt_current", "kqt_4").
		WillReturnRows(sqlmock.NewRows([]string{"id", "status", "question", "answer"}).
			AddRow("kqt_5", "completed", "What changed?", "The policy changed.").
			AddRow("kqt_6", "no_answer", "Any rollout date?", "").
			AddRow("kqt_7", "failed", "Try another source", "").
			AddRow("kqt_8", "canceled", "Stop this", ""))

	contextState, err := New(coredb.WrapSQLDB(sqlDB)).ReadRuntimeContext(
		t.Context(), "wsp_1", "usr_1", "kqs_1", "kqt_current",
	)
	require.NoError(t, err)
	require.Equal(t, int64(3), contextState.Generation)
	require.Equal(t, "Earlier deployment discussion.", contextState.Summary)
	require.Equal(t, "kqt_4", contextState.CompactedThroughTurnID)
	require.Equal(t, []string{"completed", "no_answer", "failed", "canceled"}, []string{
		contextState.Turns[0].Status, contextState.Turns[1].Status,
		contextState.Turns[2].Status, contextState.Turns[3].Status,
	})
	require.Equal(t, "The policy changed.", contextState.Turns[0].Answer)
	require.Empty(t, contextState.Turns[1].Answer)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestApplyRuntimeContextUpdateUsesGenerationCAS(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()
	now := time.Date(2026, 7, 29, 10, 5, 0, 0, time.UTC)
	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)SELECT COUNT\(\*\).*FROM knowledge_qa_turns AS compacted`).
		WithArgs("kqt_current", "wsp_1", "kqs_1", "kqt_4").
		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
	mock.ExpectQuery(`(?s)SELECT generation FROM knowledge_qa_runtime_contexts.*FOR UPDATE`).
		WithArgs("kqs_1", "wsp_1").
		WillReturnRows(sqlmock.NewRows([]string{"generation"}).AddRow(int64(3)))
	mock.ExpectExec(`(?s)UPDATE knowledge_qa_runtime_contexts.*generation=generation\+1`).
		WithArgs("Updated summary", "kqt_4", "knowledge-qa-context-v1", now, "kqs_1", "wsp_1", int64(3)).
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectRollback()

	tx, err := coredb.WrapSQLDB(sqlDB).BeginTx(t.Context(), coredb.TxOptions{})
	require.NoError(t, err)
	err = applyRuntimeContextUpdateTx(t.Context(), tx, "wsp_1", "kqs_1", "kqt_current", RuntimeContextUpdate{
		ExpectedGeneration: 3, Summary: "Updated summary",
		CompactedThroughTurnID: "kqt_4", SummaryVersion: "knowledge-qa-context-v1",
	}, now)
	require.NoError(t, err)
	require.NoError(t, tx.Rollback(t.Context()))
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestLatestTerminalKnowledgeQATurnSelectsOnlyRetryableHead(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	defer sqlDB.Close()
	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)SELECT id FROM knowledge_qa_turns.*status IN.*ORDER BY queued_at DESC`).
		WithArgs("wsp_1", "kqs_1").
		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("kqt_latest"))
	mock.ExpectRollback()

	tx, err := coredb.WrapSQLDB(sqlDB).BeginTx(t.Context(), coredb.TxOptions{})
	require.NoError(t, err)
	turnID, err := latestTerminalKnowledgeQATurn(t.Context(), tx, "wsp_1", "kqs_1")
	require.NoError(t, err)
	require.Equal(t, "kqt_latest", turnID)
	require.NoError(t, tx.Rollback(t.Context()))
	require.NoError(t, mock.ExpectationsWereMet())
}