package knowledgeqa

import (
	"database/sql"
	"errors"
	"testing"

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

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

func TestLockLegacySessionConversationForOwnerAuthorizesLegacyAgentThread(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	metadata := []byte(`{"agentThreadSession":{"kind":"agent_thread_session","agentId":"agt_1","subsystem":"kb"}}`)
	mock.ExpectQuery(`(?s)SELECT thread\.workspace_id, workspace\.organization_id,.*FROM scopes AS thread.*parent_profile\.subsystem=.*FOR UPDATE`).
		WithArgs("wsp_1", "thr_legacy", "kb", "usr_1").
		WillReturnRows(sqlmock.NewRows([]string{
			"workspace_id", "organization_id", "owner_user_id", "owner_member_id", "root_message_id", "thread_channel_id", "metadata",
		}).AddRow("wsp_1", "org_1", "usr_1", "mem_1", "msg_root", "thr_legacy", metadata))

	session, err := lockLegacySessionConversationForOwner(t.Context(), coredb.WrapSQLDB(sqlDB), "wsp_1", "usr_1", "thr_legacy")
	if err != nil {
		t.Fatal(err)
	}
	if session.WorkspaceID != "wsp_1" || session.OrganizationID != "org_1" || session.OwnerUserID != "usr_1" ||
		session.OwnerMemberID != "mem_1" || session.RootMessageID != "msg_root" || session.ThreadChannelID != "thr_legacy" ||
		session.Status != "active" {
		t.Fatalf("session = %#v", session)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestLockLegacySessionConversationForOwnerRejectsNonQAMetadataAndWrongOwner(t *testing.T) {
	for _, test := range []struct {
		name     string
		metadata []byte
		noRow    bool
	}{
		{
			name:     "non Q&A metadata",
			metadata: []byte(`{"productOpaque":{"knowledgeBaseIds":["kb_1"]}}`),
		},
		{
			name:  "wrong owner",
			noRow: true,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			sqlDB, mock, err := sqlmock.New()
			if err != nil {
				t.Fatal(err)
			}
			defer sqlDB.Close()
			expect := mock.ExpectQuery(`(?s)SELECT thread\.workspace_id, workspace\.organization_id,.*FROM scopes AS thread.*FOR UPDATE`).
				WithArgs("wsp_1", "thr_legacy", "kb", "usr_1")
			if test.noRow {
				expect.WillReturnError(sql.ErrNoRows)
			} else {
				expect.WillReturnRows(sqlmock.NewRows([]string{
					"workspace_id", "organization_id", "owner_user_id", "owner_member_id", "root_message_id", "thread_channel_id", "metadata",
				}).AddRow("wsp_1", "org_1", "usr_1", "mem_1", "msg_root", "thr_legacy", test.metadata))
			}

			_, err = lockLegacySessionConversationForOwner(t.Context(), coredb.WrapSQLDB(sqlDB), "wsp_1", "usr_1", "thr_legacy")
			if !errors.Is(err, ErrNotFound) {
				t.Fatalf("lockLegacySessionConversationForOwner() error = %v, want ErrNotFound", err)
			}
			if err := mock.ExpectationsWereMet(); err != nil {
				t.Fatal(err)
			}
		})
	}
}

func TestLegacyKnowledgeQASessionIDFromMetadataAcceptsHostedSessionRoot(t *testing.T) {
	gotID, gotFound, err := legacyKnowledgeQASessionIDFromMetadata([]byte(`{"knowledgeQASession":{"kind":"knowledge_qa_session","sessionId":" kqs_legacy "}}`))
	if err != nil {
		t.Fatal(err)
	}
	if gotID != "kqs_legacy" || !gotFound {
		t.Fatalf("legacyKnowledgeQASessionIDFromMetadata() = (%q, %v), want hosted session", gotID, gotFound)
	}
}