package app

import (
	"context"
	"errors"
	"os"
	"regexp"
	"testing"

	"github.com/go-sql-driver/mysql"
	"github.com/stretchr/testify/require"

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

type knowledgeBaseGrantCleanupResult struct{}

func (knowledgeBaseGrantCleanupResult) RowsAffected() int64 { return 1 }

type knowledgeBaseGrantCleanupFakeExecer struct {
	agentIDs []string
	err      error
}

func (f *knowledgeBaseGrantCleanupFakeExecer) Exec(_ context.Context, query string, args ...any) (coredb.Result, error) {
	if f.err != nil {
		return nil, f.err
	}
	if !regexp.MustCompile(`(?s)DELETE FROM knowledge_base_grants.*subject_type = 'agent'`).MatchString(query) {
		return nil, errors.New("unexpected cleanup query")
	}
	if len(args) >= 2 {
		if agentID, ok := args[1].(string); ok {
			f.agentIDs = append(f.agentIDs, agentID)
		}
	}
	return knowledgeBaseGrantCleanupResult{}, nil
}

func (*knowledgeBaseGrantCleanupFakeExecer) Query(context.Context, string, ...any) (coredb.Rows, error) {
	panic("unused")
}

func (*knowledgeBaseGrantCleanupFakeExecer) QueryRow(context.Context, string, ...any) coredb.Row {
	panic("unused")
}

func TestDeleteKnowledgeBaseAgentGrantsForAgentIDs(t *testing.T) {
	t.Parallel()

	execer := &knowledgeBaseGrantCleanupFakeExecer{}
	require.NoError(t, deleteKnowledgeBaseAgentGrantsForAgentIDs(context.Background(), execer, " ws_1 ", []string{" agt_1 ", "agt_1", "", "agt_2"}))
	require.Equal(t, []string{"agt_1", "agt_2"}, execer.agentIDs)
}

func TestDeleteKnowledgeBaseAgentGrantsForAgentIDsAllowsMissingMeshXTable(t *testing.T) {
	t.Parallel()

	execer := &knowledgeBaseGrantCleanupFakeExecer{err: &mysql.MySQLError{Number: 1146, Message: "Table 'test.knowledge_base_grants' doesn't exist"}}
	require.NoError(t, deleteKnowledgeBaseAgentGrantsForAgentIDs(context.Background(), execer, "ws_1", []string{"agt_1"}))
}

func TestAgentDeletionPathsInvokeKnowledgeBaseGrantCleanup(t *testing.T) {
	t.Parallel()

	runtimeSource, err := os.ReadFile("app_runtime.go")
	require.NoError(t, err)
	require.Regexp(t, regexp.MustCompile(`(?s)deleteKnowledgeBaseAgentGrantsForAgentIDs\(ctx, tx, workspaceID, \[\]string\{agentID\}\).*DELETE FROM agents`), string(runtimeSource))

	workspacesSource, err := os.ReadFile("app_workspaces.go")
	require.NoError(t, err)
	require.Regexp(t, regexp.MustCompile(`(?s)agentIDs := make\(\[\]string, 0, len\(agents\)\).*agentIDs = append\(agentIDs, item\.AgentID\).*deleteKnowledgeBaseAgentGrantsForAgentIDs\(ctx, tx, workspaceID, agentIDs\).*UPDATE agents a`), string(workspacesSource))
}