package knowledgeqa

import (
	"errors"
	"testing"
	"time"

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

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

func TestApplyTurnSnapshotAcceptsUnchangedRetrySnapshot(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()

	now := time.Date(2026, 8, 1, 0, 50, 0, 0, time.UTC)
	snapshot := TurnSnapshot{
		EffectiveModel: "model-a", EffectiveReasonLevel: "high",
		OrganizationPolicyVersion: 2,
		RuntimeVersion:            "knowledge-qa-v13",
		PromptVersion:             "native-tool-calling-v11",
	}
	mock.ExpectBegin()
	mock.ExpectExec(`(?s)UPDATE knowledge_qa_turns.*status IN \('preparing','generating'\)`).
		WithArgs(
			[]byte(`[]`), "model-a", "high", int64(2), "knowledge-qa-v13",
			"native-tool-calling-v11", now, "wsp_1", "kqt_1", "kqa_retry",
		).
		WillReturnResult(sqlmock.NewResult(0, 0))
	mock.ExpectQuery(`(?s)SELECT authorized_roots_json.*FROM knowledge_qa_turns.*FOR UPDATE`).
		WithArgs("wsp_1", "kqt_1", "kqa_retry").
		WillReturnRows(sqlmock.NewRows([]string{
			"authorized_roots_json", "effective_model", "effective_reason_level",
			"organization_policy_version", "runtime_version", "prompt_version",
		}).AddRow(
			[]byte(`[]`), "model-a", "high", int64(2),
			"knowledge-qa-v13", "native-tool-calling-v11",
		))
	mock.ExpectCommit()

	err = New(coredb.WrapSQLDB(sqlDB)).ApplyTurnSnapshot(
		t.Context(), "wsp_1", "kqt_1", "kqa_retry", snapshot, now,
	)
	if err != nil {
		t.Fatal(err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestApplyTurnSnapshotRejectsUnsatisfiedSnapshotAfterNoChange(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()

	now := time.Date(2026, 8, 1, 0, 50, 0, 0, time.UTC)
	snapshot := TurnSnapshot{
		EffectiveModel: "model-a", EffectiveReasonLevel: "high",
		OrganizationPolicyVersion: 2,
		RuntimeVersion:            "knowledge-qa-v13",
		PromptVersion:             "native-tool-calling-v11",
	}
	mock.ExpectBegin()
	mock.ExpectExec(`(?s)UPDATE knowledge_qa_turns.*status IN \('preparing','generating'\)`).
		WithArgs(
			[]byte(`[]`), "model-a", "high", int64(2), "knowledge-qa-v13",
			"native-tool-calling-v11", now, "wsp_1", "kqt_1", "kqa_retry",
		).
		WillReturnResult(sqlmock.NewResult(0, 0))
	mock.ExpectQuery(`(?s)SELECT authorized_roots_json.*FROM knowledge_qa_turns.*FOR UPDATE`).
		WithArgs("wsp_1", "kqt_1", "kqa_retry").
		WillReturnRows(sqlmock.NewRows([]string{
			"authorized_roots_json", "effective_model", "effective_reason_level",
			"organization_policy_version", "runtime_version", "prompt_version",
		}).AddRow(
			[]byte(`[]`), "model-b", "high", int64(2),
			"knowledge-qa-v13", "native-tool-calling-v11",
		))
	mock.ExpectRollback()

	err = New(coredb.WrapSQLDB(sqlDB)).ApplyTurnSnapshot(
		t.Context(), "wsp_1", "kqt_1", "kqa_retry", snapshot, now,
	)
	if !errors.Is(err, ErrInvalidState) {
		t.Fatalf("ApplyTurnSnapshot() error = %v, want ErrInvalidState", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}