package tasks

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/tidbcloud/observability/pkg/common/message"
)

func TestBuildZeroPrefixFallbackRule(t *testing.T) {
	rule := BuildZeroPrefixFallbackRule(180)
	assert.Equal(t, zeroPrefixFallbackRuleName, rule.RuleName)
	assert.Equal(t, "0/", rule.Prefix)
	assert.Equal(t, int32(180), rule.ExpirationDays)
	assert.Equal(t, int32(0), rule.ColdDays)
}

func TestBuildDiagnosisDataFallbackRule(t *testing.T) {
	rule := BuildDiagnosisDataFallbackRule(180)
	assert.Equal(t, diagnosisDataFallbackRuleName, rule.RuleName)
	assert.Equal(t, "diagnosis/data/", rule.Prefix)
	assert.Equal(t, int32(180), rule.ExpirationDays)
}

func TestBuildOnePrefixFallbackRule(t *testing.T) {
	rule := BuildOnePrefixFallbackRule(180)
	assert.Equal(t, onePrefixFallbackRuleName, rule.RuleName)
	assert.Equal(t, "1/", rule.Prefix)
	assert.Equal(t, int32(180), rule.ExpirationDays)
}

func TestBuildPrefixExpirationFallbackRules(t *testing.T) {
	rules := BuildPrefixExpirationFallbackRules(180)
	assert.Len(t, rules, 3)
	assert.Equal(t, "0/", rules[0].Prefix)
	assert.Equal(t, "diagnosis/data/", rules[1].Prefix)
	assert.Equal(t, "1/", rules[2].Prefix)
}

func TestShouldEnableMetricsGlacierRule(t *testing.T) {
	assert.True(t, shouldEnableMetricsGlacierRule("aws"))
	assert.True(t, shouldEnableMetricsGlacierRule("gcp"))
	assert.True(t, shouldEnableMetricsGlacierRule("azure"))
	assert.False(t, shouldEnableMetricsGlacierRule("alicloud"))
}

func TestBuildUnusedLifecycleRulePrefixes(t *testing.T) {
	awsPrefixes := buildUnusedLifecycleRulePrefixes("aws")
	assert.Contains(t, awsPrefixes, "profile_deletion_")
	assert.NotContains(t, awsPrefixes, metricsGlacierRuleName)

	alicloudPrefixes := buildUnusedLifecycleRulePrefixes("alicloud")
	assert.Contains(t, alicloudPrefixes, metricsGlacierRuleName)
}

func TestResolveZeroPrefixExpirationDays(t *testing.T) {
	t.Setenv("ZERO_PREFIX_EXPIRATION_DAYS", "")
	assert.Equal(t, int32(90), ResolveZeroPrefixExpirationDays(90))
	assert.Equal(t, DefaultZeroPrefixExpirationDays, ResolveZeroPrefixExpirationDays(-1))

	t.Setenv("ZERO_PREFIX_EXPIRATION_DAYS", "120")
	assert.Equal(t, int32(120), ResolveZeroPrefixExpirationDays(90))

	t.Setenv("ZERO_PREFIX_EXPIRATION_DAYS", "invalid")
	assert.Equal(t, int32(90), ResolveZeroPrefixExpirationDays(90))
}

func TestResolveProfileExpirationDays(t *testing.T) {
	t.Setenv("PROFILE_EXPIRATION_DAYS", "")
	assert.Equal(t, int32(45), ResolveProfileExpirationDays(45))
	assert.Equal(t, DefaultProfileExpirationDays, ResolveProfileExpirationDays(-1))

	t.Setenv("PROFILE_EXPIRATION_DAYS", "60")
	assert.Equal(t, int32(60), ResolveProfileExpirationDays(45))
}

func TestResolveOOMRecordExpirationDays(t *testing.T) {
	t.Setenv("OOM_RECORD_EXPIRATION_DAYS", "")
	assert.Equal(t, int32(45), ResolveOOMRecordExpirationDays(45))
	assert.Equal(t, DefaultOOMRecordExpirationDays, ResolveOOMRecordExpirationDays(-1))

	t.Setenv("OOM_RECORD_EXPIRATION_DAYS", "60")
	assert.Equal(t, int32(60), ResolveOOMRecordExpirationDays(45))
}

func TestResolveKeyvizRecordExpirationDays(t *testing.T) {
	t.Setenv("KEYVIZ_RECORD_EXPIRATION_DAYS", "")
	assert.Equal(t, int32(45), ResolveKeyvizRecordExpirationDays(45))
	assert.Equal(t, DefaultKeyvizRecordExpirationDays, ResolveKeyvizRecordExpirationDays(-1))

	t.Setenv("KEYVIZ_RECORD_EXPIRATION_DAYS", "60")
	assert.Equal(t, int32(60), ResolveKeyvizRecordExpirationDays(45))
}

func TestBuildDeltalakeRule(t *testing.T) {
	rule := BuildDeltalakeRule(180)
	assert.Equal(t, deltalakeExpirationRuleName, rule.RuleName)
	assert.Equal(t, deltalakePrefix, rule.Prefix)
	assert.Equal(t, int32(180), rule.ExpirationDays)
	assert.Equal(t, int32(0), rule.ColdDays)
}

func TestBuildAthenaResultsRule(t *testing.T) {
	rule := BuildAthenaResultsRule(1)
	assert.Equal(t, athenaResultsExpirationRuleName, rule.RuleName)
	assert.Equal(t, "athena-results/logs-query/", rule.Prefix)
	assert.Equal(t, int32(1), rule.ExpirationDays)
	assert.Equal(t, int32(0), rule.ColdDays)
}

func TestResolveDeltalakeExpirationDays(t *testing.T) {
	t.Setenv("DELTALAKE_EXPIRATION_DAYS", "")
	assert.Equal(t, int32(90), ResolveDeltalakeExpirationDays(90))
	assert.Equal(t, DefaultDeltalakeExpirationDays, ResolveDeltalakeExpirationDays(-1))

	t.Setenv("DELTALAKE_EXPIRATION_DAYS", "120")
	assert.Equal(t, int32(120), ResolveDeltalakeExpirationDays(90))
}

func TestResolveAthenaResultsExpirationDays(t *testing.T) {
	t.Setenv("ATHENA_RESULTS_EXPIRATION_DAYS", "")
	assert.Equal(t, int32(2), ResolveAthenaResultsExpirationDays(2))
	assert.Equal(t, DefaultAthenaResultsExpirationDays, ResolveAthenaResultsExpirationDays(-1))

	t.Setenv("ATHENA_RESULTS_EXPIRATION_DAYS", "3")
	assert.Equal(t, int32(3), ResolveAthenaResultsExpirationDays(2))
}

func TestBuildLifecycleRulesIncludesDeltalakeAndAthenaRules(t *testing.T) {
	rules := []message.LifecycleRuleParam{BuildDiagnosisTaskRule(), BuildHistoryTaskRule(), BuildMetricsGlacierRule("o11y-prod-shared-us-east-1")}
	rules = dedupeLifecycleRuleParams(rules)
	rules = append(rules, BuildDeltalakeRule(180))
	rules = append(rules, BuildAthenaResultsRule(1))
	rules = append(rules, BuildPrefixExpirationFallbackRules(180)...)

	assert.Len(t, rules, 8)
	assert.Equal(t, deltalakeExpirationRuleName, rules[3].RuleName)
	assert.Equal(t, athenaResultsExpirationRuleName, rules[4].RuleName)
}

func TestBuildLifecycleRulesKeepsProfileRulesWithZeroPrefixFallback(t *testing.T) {
	profileRule := message.LifecycleRuleParam{
		RuleName:       "profile_deletion_app-1",
		Prefix:         "0/tenant/project/app-1/profiles/",
		ExpirationDays: 32,
	}
	rules := dedupeLifecycleRuleParams([]message.LifecycleRuleParam{
		BuildDiagnosisTaskRule(),
		BuildHistoryTaskRule(),
		BuildMetricsGlacierRule("o11y-prod-shared-us-west-2"),
		profileRule,
	})
	rules = append(rules, BuildPrefixExpirationFallbackRules(180)...)

	assert.Len(t, rules, 7)
	assert.Equal(t, historyTaskExpirationRuleName, rules[1].RuleName)
	assert.Equal(t, profileRule, rules[3])
	assert.Equal(t, zeroPrefixFallbackRuleName, rules[4].RuleName)
	assert.Equal(t, diagnosisDataFallbackRuleName, rules[5].RuleName)
	assert.Equal(t, onePrefixFallbackRuleName, rules[6].RuleName)
}

func TestBuildLifecycleRulesSkipsMetricsGlacierForAliCloud(t *testing.T) {
	rules := []message.LifecycleRuleParam{BuildDiagnosisTaskRule(), BuildHistoryTaskRule()}
	if shouldEnableMetricsGlacierRule("alicloud") {
		rules = append(rules, BuildMetricsGlacierRule("o11y-prod-shared-ap-southeast-1"))
	}
	rules = dedupeLifecycleRuleParams(rules)
	rules = append(rules, BuildPrefixExpirationFallbackRules(180)...)

	assert.Len(t, rules, 5)
	assert.Equal(t, historyTaskExpirationRuleName, rules[1].RuleName)
	assert.Equal(t, zeroPrefixFallbackRuleName, rules[2].RuleName)
	assert.Equal(t, diagnosisDataFallbackRuleName, rules[3].RuleName)
	assert.Equal(t, onePrefixFallbackRuleName, rules[4].RuleName)
	for _, rule := range rules {
		assert.NotEqual(t, metricsGlacierRuleName, rule.RuleName)
	}
}

func TestBuildDiagnosisTaskRule(t *testing.T) {
	rule := BuildDiagnosisTaskRule()
	assert.Equal(t, "diagnosis_task_expiration", rule.RuleName)
	assert.Equal(t, diagnosisTaskPrefix, rule.Prefix)
	assert.Equal(t, int32(2), rule.ExpirationDays)
}

func TestBuildHistoryTaskRule(t *testing.T) {
	rule := BuildHistoryTaskRule()
	assert.Equal(t, historyTaskExpirationRuleName, rule.RuleName)
	assert.Equal(t, diagnosisHistoryTaskPrefix, rule.Prefix)
	assert.Equal(t, DefaultHistoryTaskExpirationDays, rule.ExpirationDays)
	assert.Equal(t, int32(32), rule.ExpirationDays)
}

func TestAppendClusterDataRetentionRulesFromHistoryTask(t *testing.T) {
	var rules []message.LifecycleRuleParam
	seenRuleNames := map[string]struct{}{}
	seenPrefixes := map[string]struct{}{}
	appendUniqueRule := func(rule message.LifecycleRuleParam) {
		if _, exists := seenRuleNames[rule.RuleName]; exists {
			return
		}
		if _, exists := seenPrefixes[rule.Prefix]; exists {
			return
		}
		seenRuleNames[rule.RuleName] = struct{}{}
		seenPrefixes[rule.Prefix] = struct{}{}
		rules = append(rules, rule)
	}

	identity, ok := lifecycleIdentityFromHistoryTask(O11yDiagnosisTask{
		Meta: Meta{ClusterID: "app-1", TenantID: "t1", ProjectID: "p1", ClusterType: "dedicated"},
	})
	assert.True(t, ok)
	appendClusterDataRetentionRules(appendUniqueRule, identity, 32, 32, 32)
	assert.Len(t, rules, 3)
	assert.Equal(t, "profile_deletion_app-1", rules[0].RuleName)
	assert.Equal(t, "0/t1/p1/app-1/profiles/", rules[0].Prefix)
	assert.Equal(t, "oom_record_deletion_app-1", rules[1].RuleName)
	assert.Equal(t, "keyviz_record_deletion_app-1", rules[2].RuleName)
	assert.Equal(t, "0/t1/p1/app-1/regions/", rules[2].Prefix)

	nextGen, ok := lifecycleIdentityFromHistoryTask(O11yDiagnosisTask{
		Meta: Meta{ClusterID: "app-ng", ClusterType: "nextgen"},
	})
	assert.True(t, ok)
	appendClusterDataRetentionRules(appendUniqueRule, nextGen, 32, 32, 32)
	assert.Len(t, rules, 4)
	assert.Equal(t, "profile_deletion_app-ng", rules[3].RuleName)
	assert.Equal(t, "0/shared_pool/app-ng/profiles/", rules[3].Prefix)
}

func TestLifecycleIdentityFromHistoryTaskSkipsInfra(t *testing.T) {
	_, ok := lifecycleIdentityFromHistoryTask(O11yDiagnosisTask{Meta: Meta{ClusterID: "seed"}})
	assert.False(t, ok)
}

func TestDedupeDropsZeroPrefixWhenAddedBeforeProfileRules(t *testing.T) {
	rules := dedupeLifecycleRuleParams([]message.LifecycleRuleParam{
		BuildZeroPrefixFallbackRule(180),
		{
			RuleName:       "profile_deletion_app-1",
			Prefix:         "0/tenant/project/app-1/profiles/",
			ExpirationDays: 32,
		},
	})

	assert.Len(t, rules, 1)
	assert.Equal(t, zeroPrefixFallbackRuleName, rules[0].RuleName)
}