package knowledgeqa

import (
	"context"
	"errors"
	"strings"
	"testing"
	"time"

	"github.com/DATA-DOG/go-sqlmock"

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

func TestPrepareCitationsAllocatesStableIDs(t *testing.T) {
	prepared, err := PrepareCitations([]Citation{{
		KnowledgeBaseID: "kb_1", SourceID: "src_1", RootKind: "primary",
		RepoRole: "canonical", Revision: "sha_1", Path: "docs/design.md",
		LineStart: 4, LineEnd: 6,
	}})
	if err != nil {
		t.Fatal(err)
	}
	if len(prepared) != 1 || !strings.HasPrefix(prepared[0].ID, "kqc_") {
		t.Fatalf("PrepareCitations() = %#v", prepared)
	}
}

func TestCompleteAttemptWithAssistantRollsBackCitationsWhenMessageFails(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", "canonical", "sha_1",
			"", "", "docs/design.md", int64(4), int64(6),
			"", "", "", nil, now,
		).
		WillReturnResult(sqlmock.NewResult(0, 1))
	mock.ExpectRollback()

	err = New(coredb.WrapSQLDB(sqlDB)).CompleteAttemptWithAssistant(
		t.Context(), "wsp_1", "kqt_1", "kqa_1", Completion{Status: "completed", Now: now},
		[]Citation{{
			ID: "kqc_1", KnowledgeBaseID: "kb_1", SourceID: "src_1",
			RootKind: "primary", RepoRole: "canonical", Revision: "sha_1",
			Path: "docs/design.md", LineStart: 4, LineEnd: 6,
		}},
		func(context.Context, coredb.Tx) (string, error) {
			return "", errors.New("message write failed")
		},
	)
	if err == nil || !strings.Contains(err.Error(), "message write failed") {
		t.Fatalf("CompleteAttemptWithAssistant() error = %v", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}