package o11ytidbagent

import (
	"context"
	"errors"
	"testing"

	"github.com/go-logr/logr"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	agentv2 "github.com/tidbcloud/observability/agent-control-server/api/v2"
	"github.com/tidbcloud/observability/agent-control-server/pkg/config"
	"github.com/tidbcloud/observability/agent-control-server/pkg/consts"
	"github.com/tidbcloud/observability/agent-control-server/pkg/deploy"
	"github.com/tidbcloud/observability/agent-control-server/pkg/reconciler"
)

func TestBuildTasks(t *testing.T) {
	tests := []struct {
		name            string
		vendor          string
		remoteAvailable bool
		agentConfig     func(*agentv2.O11yTiDBAgent)
		build           func(*AgentReconciler, *reconciler.Context[*agentv2.O11yTiDBAgent], bool) []deploy.Task
		want            []string
	}{
		{
			name:            "aws includes storage class only when remote is available",
			vendor:          consts.VendorTypeAWS,
			remoteAvailable: true,
			agentConfig: func(agent *agentv2.O11yTiDBAgent) {
				agent.Spec.EKS = &agentv2.EKSConfig{}
			},
			build: (*AgentReconciler).buildTasks,
			want: []string{
				"Guardian", "Backend", "AutoResourcing", "StorageClass",
				"PrivateLink", "NodeGroup",
				"PriorityClass", "VMAgent", "Vector", "Restart",
				"DrainNode", "UpdateVersion",
			},
		},
		{
			name:            "aws skips remote dependent tasks when remote is unavailable",
			vendor:          consts.VendorTypeAWS,
			remoteAvailable: false,
			agentConfig: func(agent *agentv2.O11yTiDBAgent) {
				agent.Spec.EKS = &agentv2.EKSConfig{}
			},
			build: (*AgentReconciler).buildTasks,
			want: []string{
				"Guardian", "Backend", "AutoResourcing", "PrivateLink",
				"NodeGroup", "UpdateVersion",
			},
		},
		{
			name:            "gcp includes storage class before private link when remote is available",
			vendor:          consts.VendorTypeGCP,
			remoteAvailable: true,
			agentConfig: func(agent *agentv2.O11yTiDBAgent) {
				agent.Spec.GKE = &agentv2.GKEConfig{}
			},
			build: (*AgentReconciler).buildTasks,
			want: []string{
				"Guardian", "Backend", "AutoResourcing", "StorageClass",
				"PrivateLink", "NodePool",
				"PriorityClass", "VMAgent", "Vector", "Restart",
				"DrainNode", "UpdateVersion",
			},
		},
		{
			name:            "gcp skips storage class and workloads when remote is unavailable",
			vendor:          consts.VendorTypeGCP,
			remoteAvailable: false,
			agentConfig: func(agent *agentv2.O11yTiDBAgent) {
				agent.Spec.GKE = &agentv2.GKEConfig{}
			},
			build: (*AgentReconciler).buildTasks,
			want: []string{
				"Guardian", "Backend", "AutoResourcing", "PrivateLink",
				"NodePool", "UpdateVersion",
			},
		},
		{
			name:            "azure fetches kubelet identity before backend",
			vendor:          consts.VendorTypeAzure,
			remoteAvailable: true,
			agentConfig: func(agent *agentv2.O11yTiDBAgent) {
				agent.Spec.AKS = &agentv2.AKSConfig{
					SubscriptionID: "subscription-1",
					ResourceGroup:  "rg-1",
					SubnetID:       "subnet-1",
					PodSubnetID:    "pod-subnet-1",
				}
			},
			build: (*AgentReconciler).buildTasks,
			want: []string{
				"Guardian", "KubeletIdentity", "Backend", "AutoResourcing",
				"StorageClass", "PrivateLinkIP", "NodePool",
				"PriorityClass", "VMAgent", "Vector",
				"Restart", "DrainNode",
				"UpdateVersion",
			},
		},
		{
			name:            "azure skips storage class and workloads when remote is unavailable",
			vendor:          consts.VendorTypeAzure,
			remoteAvailable: false,
			agentConfig: func(agent *agentv2.O11yTiDBAgent) {
				agent.Spec.AKS = &agentv2.AKSConfig{
					SubscriptionID: "subscription-1",
					ResourceGroup:  "rg-1",
					SubnetID:       "subnet-1",
					PodSubnetID:    "pod-subnet-1",
				}
			},
			build: (*AgentReconciler).buildTasks,
			want: []string{
				"Guardian", "KubeletIdentity", "Backend", "AutoResourcing",
				"PrivateLinkIP", "NodePool", "UpdateVersion",
			},
		},
		{
			name:            "unknown vendor with remote includes only shared remote tasks",
			vendor:          consts.VendorTypeAliCloud,
			remoteAvailable: true,
			build:           (*AgentReconciler).buildTasks,
			want: []string{
				"Guardian", "Backend", "AutoResourcing",
				"PriorityClass", "VMAgent", "Vector",
				"Restart", "DrainNode", "UpdateVersion",
			},
		},
		{
			name:            "unknown vendor without remote keeps only control-plane tasks",
			vendor:          consts.VendorTypeAliCloud,
			remoteAvailable: false,
			build:           (*AgentReconciler).buildTasks,
			want:            []string{"Guardian", "Backend", "AutoResourcing", "UpdateVersion"},
		},
		{
			name:            "pause deploy scales workloads only when remote is available",
			vendor:          consts.VendorTypeAWS,
			remoteAvailable: true,
			build:           (*AgentReconciler).buildPauseDeployTasks,
			want:            []string{"Guardian", "VMAgent", "Vector"},
		},
		{
			name:            "pause deploy keeps only guardian when remote is unavailable",
			vendor:          consts.VendorTypeAWS,
			remoteAvailable: false,
			build:           (*AgentReconciler).buildPauseDeployTasks,
			want:            []string{"Guardian"},
		},
		{
			name:   "pause cleanup removes aws node group",
			vendor: consts.VendorTypeAWS,
			agentConfig: func(agent *agentv2.O11yTiDBAgent) {
				agent.Spec.EKS = &agentv2.EKSConfig{}
			},
			build: func(r *AgentReconciler, ctx *reconciler.Context[*agentv2.O11yTiDBAgent], _ bool) []deploy.Task {
				return r.buildPauseCleanupTasks(ctx)
			},
			want: []string{"NodeGroup"},
		},
		{
			name:   "pause cleanup removes gcp node pool",
			vendor: consts.VendorTypeGCP,
			agentConfig: func(agent *agentv2.O11yTiDBAgent) {
				agent.Spec.GKE = &agentv2.GKEConfig{}
			},
			build: func(r *AgentReconciler, ctx *reconciler.Context[*agentv2.O11yTiDBAgent], _ bool) []deploy.Task {
				return r.buildPauseCleanupTasks(ctx)
			},
			want: []string{"NodePool"},
		},
		{
			name:   "pause cleanup removes azure node pool",
			vendor: consts.VendorTypeAzure,
			agentConfig: func(agent *agentv2.O11yTiDBAgent) {
				agent.Spec.AKS = &agentv2.AKSConfig{
					SubscriptionID: "subscription-1",
					ResourceGroup:  "rg-1",
					SubnetID:       "subnet-1",
					PodSubnetID:    "pod-subnet-1",
				}
			},
			build: func(r *AgentReconciler, ctx *reconciler.Context[*agentv2.O11yTiDBAgent], _ bool) []deploy.Task {
				return r.buildPauseCleanupTasks(ctx)
			},
			want: []string{"NodePool"},
		},
		{
			name:   "pause cleanup has no task for unsupported vendor",
			vendor: consts.VendorTypeAliCloud,
			build: func(r *AgentReconciler, ctx *reconciler.Context[*agentv2.O11yTiDBAgent], _ bool) []deploy.Task {
				return r.buildPauseCleanupTasks(ctx)
			},
			want: []string{},
		},
		{
			name:            "hibernated keeps only guardian task",
			vendor:          consts.VendorTypeAWS,
			remoteAvailable: true,
			build: func(reconciler *AgentReconciler, ctx *reconciler.Context[*agentv2.O11yTiDBAgent], _ bool) []deploy.Task {
				return reconciler.buildHibernatedTasks(ctx)
			},
			want: []string{"Guardian"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			agent := newO11yTiDBAgent()
			if tt.agentConfig != nil {
				tt.agentConfig(agent)
			}
			ar := &AgentReconciler{Reconciler: reconciler.Reconciler{
				Config: newO11yTiDBRootConfig(tt.vendor),
			}}
			ctx := &reconciler.Context[*agentv2.O11yTiDBAgent]{
				Agent:  agent,
				Config: ar.Config,
			}

			got := tt.build(ar, ctx, tt.remoteAvailable)

			assert.Len(t, got, len(tt.want))
			for i, want := range tt.want {
				assert.Equal(t, want, got[i].Name())
			}
		})
	}
}

func newO11yTiDBAgent() *agentv2.O11yTiDBAgent {
	return &agentv2.O11yTiDBAgent{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "agent-1",
			Namespace: "default",
		},
		Spec: agentv2.O11yTiDBAgentSpec{
			KubernetesClusterName: "cluster-1",
			O11y: agentv2.O11yConfig{
				BizType:       "tidbcloud",
				TenantID:      "tenant-1",
				ProjectID:     "project-1",
				ApplicationID: "app-1",
			},
		},
	}
}

func TestUpdateStatusToNormalPanicPhase(t *testing.T) {
	tests := []struct {
		name       string
		startPhase string
		results    []deploy.Result
		wantPhase  string
	}{
		{
			name:       "panic in results sets phase to Panic",
			startPhase: agentv2.O11yAgentPhaseNormal,
			results: []deploy.Result{
				deploy.Done(),
				deploy.Panic(errors.New("vmpodscrape not operational")),
			},
			wantPhase: agentv2.O11yAgentPhasePanic,
		},
		{
			name:       "partial recovery with Error clears Panic",
			startPhase: agentv2.O11yAgentPhasePanic,
			results: []deploy.Result{
				deploy.Done(),
				deploy.Error(errors.New("transient")),
			},
			wantPhase: agentv2.O11yAgentPhaseNormal,
		},
		{
			name:       "full recovery returns phase to Normal",
			startPhase: agentv2.O11yAgentPhasePanic,
			results: []deploy.Result{
				deploy.Done(),
				deploy.Done(),
			},
			wantPhase: agentv2.O11yAgentPhaseNormal,
		},
		{
			name:       "Error from Normal leaves phase Normal",
			startPhase: agentv2.O11yAgentPhaseNormal,
			results: []deploy.Result{
				deploy.Error(errors.New("e")),
			},
			wantPhase: agentv2.O11yAgentPhaseNormal,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			agent := newO11yTiDBAgent()
			agent.Status.Phase = tt.startPhase

			scheme := runtime.NewScheme()
			require.NoError(t, agentv2.AddToScheme(scheme))
			cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(agent).Build()

			ctx := &reconciler.Context[*agentv2.O11yTiDBAgent]{
				Agent:  agent,
				Local:  cli,
				Config: newO11yTiDBRootConfig(consts.VendorTypeAWS),
			}

			_, _ = updateStatusToNormal(context.Background(), logr.Discard(), ctx, tt.results)
			assert.Equal(t, tt.wantPhase, agent.Status.Phase)
		})
	}
}

func newO11yTiDBRootConfig(vendor string) *config.RootConfig {
	return &config.RootConfig{
		O11yConfig: config.ObservabilityInfra{
			Env:     "dev",
			Vendor:  vendor,
			Version: 1,
			CustomerConfigsV2: config.CustomerConfigsV2{
				Dedicated: config.CustomerDedicated{
					Repos: map[string]string{"vm": "https://charts.example.com"},
					Specs: config.CustomerDedicatedSpecs{
						TiDBCluster: config.CustomerAgentSpecs{{Name: "default", MaxNodes: 10}},
					},
				},
			},
		},
	}
}