package repository

import (
	"context"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/stretchr/testify/require"
)

func TestMemoryAgentPolicyUsesPerTurnIdentityAndRejectsDeletedAgents(t *testing.T) {
	db, sessions := newOwnerScopeDB(t, t.Name())
	addMemoryWebOwner(t, db)
	require.NoError(t, db.AutoMigrate(&types.CustomAgent{}))
	require.NoError(t, db.Model(&types.Session{}).Where("id = ?", sessions["alice"]).Update("user_id", "alice").Error)
	disabled := false
	agents := []*types.CustomAgent{
		{ID: "allowed", TenantID: 7, Enabled: true},
		{ID: "disabled", TenantID: 7, Enabled: false},
		{ID: "opted-out", TenantID: 7, Enabled: true, Config: types.CustomAgentConfig{MemoryEnabled: &disabled}},
		{ID: "deleted", TenantID: 7, Enabled: true},
		{ID: "foreign", TenantID: 8, Enabled: true},
	}
	for _, agent := range agents {
		require.NoError(t, db.Create(agent).Error)
	}
	require.NoError(t, db.Model(agents[1]).Update("enabled", false).Error)
	require.NoError(t, db.Delete(agents[3]).Error)
	base := time.Now().Add(-time.Hour)
	var inputs []*types.Message
	for i, agentID := range []string{"allowed", "disabled", "opted-out", "deleted", "missing", "foreign", ""} {
		input := &types.Message{ID: "user-" + agentID, SessionID: sessions["alice"], Role: "user", RequestID: "request-" + agentID, CreatedAt: base.Add(time.Duration(i) * time.Minute)}
		answer := &types.Message{ID: "answer-" + agentID, SessionID: input.SessionID, Role: "assistant", RequestID: input.RequestID, AgentID: agentID, IsCompleted: true, CreatedAt: input.CreatedAt.Add(time.Second)}
		require.NoError(t, db.Create(input).Error)
		require.NoError(t, db.Create(answer).Error)
		inputs = append(inputs, input)
	}
	repo := NewMessageRepository(db)
	scope := interfaces.MemoryScope{TenantID: 7, SubjectID: "web_user:alice"}
	allowed, err := repo.MemoryAllowedMessageIDs(context.Background(), scope, sessions["alice"], inputs)
	require.NoError(t, err)
	require.Equal(t, map[string]bool{inputs[0].ID: true, inputs[6].ID: true}, allowed)
	require.NoError(t, db.Model(agents[0]).Update("enabled", false).Error)
	allowed, err = repo.MemoryAllowedMessageIDs(context.Background(), scope, sessions["alice"], inputs)
	require.NoError(t, err)
	require.Equal(t, map[string]bool{inputs[6].ID: true}, allowed)
	allowed, err = repo.MemoryAllowedMessageIDs(context.Background(), interfaces.MemoryScope{TenantID: 7, SubjectID: "web_user:bob"}, sessions["alice"], inputs)
	require.NoError(t, err)
	require.Empty(t, allowed)
}

func TestMemorySharedAgentRechecksExactSourceAndCurrentMembership(t *testing.T) {
	db, sessions := newOwnerScopeDB(t, t.Name())
	addMemoryWebOwner(t, db)
	require.NoError(t, db.AutoMigrate(&types.CustomAgent{}))
	require.NoError(t, db.Model(&types.Session{}).Where("id = ?", sessions["alice"]).Update("user_id", "alice").Error)
	for _, query := range []string{
		`CREATE TABLE agent_shares (id TEXT PRIMARY KEY, agent_id TEXT, organization_id TEXT, source_tenant_id INTEGER, deleted_at DATETIME)`,
		`CREATE TABLE organization_tenant_members (organization_id TEXT, tenant_id INTEGER)`,
		`CREATE TABLE organizations (id TEXT PRIMARY KEY, deleted_at DATETIME)`,
		`INSERT INTO organizations(id) VALUES ('org')`,
		`INSERT INTO organization_tenant_members VALUES ('org', 7)`,
		`INSERT INTO agent_shares(id, agent_id, organization_id, source_tenant_id) VALUES ('share', 'shared', 'org', 84)`,
	} {
		require.NoError(t, db.Exec(query).Error)
	}
	require.NoError(t, db.Create(&types.CustomAgent{ID: "shared", TenantID: 84, Enabled: true}).Error)
	require.NoError(t, db.Create(&types.CustomAgent{ID: "shared", TenantID: 42, Enabled: true}).Error)
	input := &types.Message{ID: "input", SessionID: sessions["alice"], Role: "user", RequestID: "request", CreatedAt: time.Now().Add(-time.Minute)}
	answer := &types.Message{ID: "answer", SessionID: input.SessionID, Role: "assistant", RequestID: input.RequestID, AgentID: "shared", AgentTenantID: 84, IsCompleted: true, CreatedAt: time.Now()}
	require.NoError(t, db.Create(input).Error)
	require.NoError(t, db.Create(answer).Error)
	repo := NewMessageRepository(db)
	scope := interfaces.MemoryScope{TenantID: 7, SubjectID: "web_user:alice"}
	check := func(want bool) {
		t.Helper()
		allowed, err := repo.MemoryAllowedMessageIDs(context.Background(), scope, input.SessionID, []*types.Message{input})
		require.NoError(t, err)
		require.Equal(t, want, allowed[input.ID])
	}
	check(true)
	require.NoError(t, db.Model(answer).Update("agent_tenant_id", 42).Error)
	check(false)
	require.NoError(t, db.Model(answer).Update("agent_tenant_id", 84).Error)
	require.NoError(t, db.Exec("DELETE FROM organization_tenant_members").Error)
	check(false)
}