package server

import (
	"context"
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/identity/trusted"
	"github.com/tidbcloud/AgentTeam/internal/service/memory"
	memorystore "github.com/tidbcloud/AgentTeam/internal/store/resource/memory"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

func TestBackfillAgentCloudMemoryEnablesEligibleUnboundAgents(t *testing.T) {
	ctx := context.Background()
	agentteamCoreDB := newBackfillSQLiteDB(t, "agentteam")
	meshxDB := newBackfillSQLiteDB(t, "meshx")
	require.NoError(t, meshxDB.AutoMigrate(&memorystore.Memory{}))
	seedBackfillAgentTeamCoreTables(t, agentteamCoreDB)

	boundAgentID := "agt_bound"
	require.NoError(t, meshxDB.Create(&memorystore.Memory{
		ID:       "mem_bound",
		AgentID:  &boundAgentID,
		Name:     "bound memory",
		IsActive: true,
	}).Error)

	enabler := &recordingBackfillEnabler{}
	result, err := backfillAgentCloudMemory(ctx, agentteamCoreDB, meshxDB, enabler, AgentCloudMemoryBackfillOptions{Limit: 10})
	require.NoError(t, err)
	require.Equal(t, 2, result.Scanned)
	require.Equal(t, 1, result.AlreadyBound)
	require.Equal(t, 1, result.Enabled)
	require.Equal(t, 0, result.Failed)
	require.Len(t, enabler.calls, 1)
	require.Equal(t, "ws_1", enabler.calls[0].workspaceID)
	require.Equal(t, "usr_owner", enabler.calls[0].subject)
	require.Equal(t, "agt_new", enabler.calls[0].agentID)
	require.NotNil(t, enabler.calls[0].input.ImportExistingLocalContent)
	require.True(t, *enabler.calls[0].input.ImportExistingLocalContent)
	require.Equal(t, "usr_owner", enabler.calls[0].identity.Subject)
	require.Equal(t, "ws_1", enabler.calls[0].identity.RawClaims["agent_ownership_workspace_id"])
}

func newBackfillSQLiteDB(t *testing.T, suffix string) *gorm.DB {
	t.Helper()
	db, err := gorm.Open(sqlite.Open("file:"+filepath.Join(t.TempDir(), suffix+".db")+"?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	return db
}

func seedBackfillAgentTeamCoreTables(t *testing.T, db *gorm.DB) {
	t.Helper()
	require.NoError(t, db.Exec(`CREATE TABLE agents (id TEXT PRIMARY KEY, workspace_id TEXT, name TEXT, member_id TEXT, machine_id TEXT, status TEXT)`).Error)
	require.NoError(t, db.Exec(`CREATE TABLE members (id TEXT PRIMARY KEY, display_name TEXT)`).Error)
	require.NoError(t, db.Exec(`CREATE TABLE machines (id TEXT, workspace_id TEXT, created_by_user_id TEXT, status TEXT)`).Error)
	require.NoError(t, db.Exec(`INSERT INTO members (id, display_name) VALUES
		('mem_new', 'New Agent'),
		('mem_bound', 'Bound Agent'),
		('mem_archived', 'Archived Agent'),
		('mem_deleted_machine', 'Deleted Machine Agent')`).Error)
	require.NoError(t, db.Exec(`INSERT INTO machines (id, workspace_id, created_by_user_id, status) VALUES
		('mac_owner', 'ws_1', 'usr_owner', 'active'),
		('mac_deleted', 'ws_1', 'usr_owner', 'archived')`).Error)
	require.NoError(t, db.Exec(`INSERT INTO agents (id, workspace_id, name, member_id, machine_id, status) VALUES
		('agt_new', 'ws_1', 'new-agent', 'mem_new', 'mac_owner', 'active'),
		('agt_bound', 'ws_1', 'bound-agent', 'mem_bound', 'mac_owner', 'active'),
		('agt_archived', 'ws_1', 'archived-agent', 'mem_archived', 'mac_owner', 'archived'),
		('agt_deleted_machine', 'ws_1', 'deleted-machine-agent', 'mem_deleted_machine', 'mac_deleted', 'active')`).Error)
}

type recordingBackfillEnabler struct {
	calls []recordedBackfillEnableCall
}

type recordedBackfillEnableCall struct {
	workspaceID string
	subject     string
	agentID     string
	input       memory.AgentCloudMemoryEnableMutation
	identity    trusted.Identity
}

func (e *recordingBackfillEnabler) EnableAgentCloudMemory(ctx context.Context, workspaceID, subject, agentID string, input memory.AgentCloudMemoryEnableMutation) (memory.AgentCloudMemory, error) {
	identity, _ := trusted.IdentityFromContext(ctx)
	e.calls = append(e.calls, recordedBackfillEnableCall{
		workspaceID: workspaceID,
		subject:     subject,
		agentID:     agentID,
		input:       input,
		identity:    identity,
	})
	return memory.AgentCloudMemory{AgentID: agentID, Enabled: true}, nil
}