package task

import (
	"context"
	"fmt"

	rbacv1 "k8s.io/api/rbac/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	agentv2 "github.com/tidbcloud/observability/agent-control-server/api/v2"
	"github.com/tidbcloud/observability/agent-control-server/controllers/o11ytidbagent/metadata"
	"github.com/tidbcloud/observability/agent-control-server/pkg/agentbuild"
	"github.com/tidbcloud/observability/agent-control-server/pkg/cloud/aws/awscrossplane"
	"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/o11y"
	"github.com/tidbcloud/observability/agent-control-server/pkg/reconciler"
	"github.com/tidbcloud/observability/agent-control-server/pkg/util"
)

type VectorTask struct {
	*reconciler.BaseTask[*agentv2.O11yTiDBAgent]
}

func NewVectorTask(ctx *reconciler.Context[*agentv2.O11yTiDBAgent]) *VectorTask {
	return &VectorTask{
		BaseTask: &reconciler.BaseTask[*agentv2.O11yTiDBAgent]{Context: ctx},
	}
}

func (t *VectorTask) Name() string {
	return "Vector"
}

func (t *VectorTask) ClusterMeta() deploy.ClusterMeta {
	return metadata.ClusterMeta(t.Agent)
}

func (t *VectorTask) Create(ctx context.Context) deploy.Result {
	cli, err := t.Remote.GetClient(ctx)
	if err != nil {
		return deploy.Panic(err)
	}
	if err := agentbuild.CheckNamespaceExists(ctx, cli, t.buildNamespace()); err != nil {
		return deploy.Panic(err) // this might happen when the tidb namespace is not created
	}
	cfg := t.Config.O11yConfig.CustomerConfigsV2.Dedicated.TiDBCluster.BuildConfig().Vector
	if err := agentbuild.CreateOrUpdateRole(ctx, cli, t.buildRole); err != nil {
		return deploy.Panic(err)
	}
	if err := agentbuild.CreateOrUpdateRoleBinding(ctx, cli, t.buildRoleBinding); err != nil {
		return deploy.Panic(err)
	}
	if err := agentbuild.InstallOrUpgradeHelm(ctx, t.Remote, t.build, cfg.Chart.Name, cfg.Chart.Version, consts.O11yNameClusterVector, t.buildNamespace(), nil); err != nil {
		return deploy.Panic(err)
	}
	return deploy.Done()
}

func (t *VectorTask) PostCreate(ctx context.Context) deploy.Result {
	cli, err := t.Remote.GetClient(ctx)
	if err != nil {
		return deploy.Panic(err)
	}
	if err := agentbuild.CheckStatefulSetReady(ctx, cli, consts.O11yNameClusterVector, t.buildNamespace()); err != nil {
		return deploy.Panic(err)
	}
	switch util.GetAnnotation(t.Agent, consts.AnnotationKeyO11ySpecScaleStatus) {
	case consts.SpecScaleStatusNewNodeGroupSynced:
		// Assume the vector task comes after vmagent, so we only update the status in the vector task.
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11ySpecScaleStatus, consts.SpecScaleStatusResourceUpdated)
	}
	return deploy.Done()
}

func (t *VectorTask) Delete(ctx context.Context) (res deploy.Result) {
	cli, err := t.Remote.GetClient(ctx)
	if err != nil {
		return deploy.Panic(err)
	}
	name := consts.O11yNameClusterVector
	if err := agentbuild.DeleteRoleBinding(ctx, cli, name, t.buildNamespace()); err != nil {
		return deploy.Panic(err)
	}
	if err := agentbuild.DeleteRole(ctx, cli, name, t.buildNamespace()); err != nil {
		return deploy.Panic(err)
	}
	if err := agentbuild.DeleteHelm(ctx, t.Remote, name, t.buildNamespace()); err != nil {
		return deploy.Panic(err)
	}
	if err := agentbuild.DeletePVC(ctx, cli, fmt.Sprintf("data-%s-0", name), t.buildNamespace()); err != nil {
		return deploy.Panic(err)
	}
	return deploy.Done()
}

func (t *VectorTask) PostDelete(ctx context.Context) (res deploy.Result) {
	cli, err := t.Remote.GetClient(ctx)
	if err != nil {
		return deploy.Panic(err)
	}
	return deploy.PanicOrDone(agentbuild.CheckStatefulSetDeleted(ctx, cli, consts.O11yNameClusterVector, t.buildNamespace()))
}

func (t *VectorTask) build() (string, error) {
	cfg := t.Config.O11yConfig.CustomerConfigsV2.Dedicated.TiDBCluster.BuildConfig().Vector
	imageRepo, imageTag, err := agentbuild.BuildVectorImage(t.Agent, cfg.Image)
	if err != nil {
		return "", err
	}

	desiredSpec := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yDesiredSpec)
	switch util.GetAnnotation(t.Agent, consts.AnnotationKeyO11ySpecScaleStatus) {
	case consts.SpecScaleStatusPendingScale:
		desiredSpec = util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yPreviousSpec)
	}
	if desiredSpec == "" {
		desiredSpec = "default"
	}

	spec := t.Config.O11yConfig.CustomerConfigsV2.Dedicated.Specs.TiDBCluster.Select(desiredSpec)
	if spec == nil {
		return "", fmt.Errorf("desired spec not found: %s", desiredSpec)
	}

	resourcesJson, err := agentbuild.BuildVectorResources(t.Agent, spec.Vector.DeepCopy())
	if err != nil {
		return "", err
	}

	configTemplate := fmt.Sprintf("dedicated/tidbcluster/vector-values-%s.yaml", t.Config.O11yConfig.Vendor)
	tmpl, err := agentbuild.BuildVectorTemplate(t.Agent, t.Templates, configTemplate)
	if err != nil {
		return "", err
	}

	replicas := 1
	if t.Agent.Paused() {
		replicas = 0
	}

	vmImportURL, err := t.buildVMImportURL()
	if err != nil {
		return "", err
	}

	return agentbuild.BuildVectorValues(tmpl, map[string]string{
		"ImageRepo":         imageRepo,
		"ImageTag":          imageTag,
		"Resources":         resourcesJson,
		"PriorityClassName": consts.PriorityClassNameClusterCritical,
		"Vendor":            t.Config.O11yConfig.Vendor,
		"Region":            t.Config.O11yConfig.Region,
		"Bucket":            util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yLogBucketName),
		"TenantID":          t.Agent.Spec.O11y.TenantID,
		"ProjectID":         t.Agent.Spec.O11y.ProjectID,
		"ClusterID":         t.Agent.Spec.O11y.ApplicationID,
		"ClusterName":       t.Agent.Spec.KubernetesClusterName,
		"VMImportURL":       vmImportURL,
		"RoleARN":           util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yAgentRoleArn),
		"Replicas":          fmt.Sprintf("%d", replicas),
		"AzureClientID":     util.GetAnnotation(t.Agent, consts.AnnotationKeyAKSKubeletClientID),
		"ContainerName":     util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yLogBucketName),
		"StorageAccount":    util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yLogStorageAccount),
	})
}

func (t *VectorTask) buildVMImportURL() (string, error) {
	o11yID := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yID)

	// For AWS, get DNS from VPC endpoint if private DNS is disabled
	if t.Config.O11yConfig.Vendor == consts.VendorTypeAWS {
		if util.GetAnnotation(t.Agent, consts.AnnotationKeyDisablePrivateEndpointPrivateDNS) != "false" {
			endpointName := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11ySvcInterfaceEndpointName)
			if endpointName == "" {
				return "", fmt.Errorf("endpoint name is required when private DNS is disabled")
			}
			dns, err := awscrossplane.GetVPCEndpointDNSByEndpointName(context.Background(), t.Local, endpointName)
			if err != nil {
				return "", fmt.Errorf("failed to get VPC endpoint DNS for %s: %w", endpointName, err)
			}
			builder := o11y.NewWriteURLBuilder(dns)
			return builder.ForAWS().BuildVMImportV2(o11yID), nil
		}
	}

	// For GCP, get IP from PSC endpoint if private DNS is disabled
	if t.Config.O11yConfig.Vendor == consts.VendorTypeGCP {
		if util.GetAnnotation(t.Agent, consts.AnnotationKeyDisablePrivateEndpointPrivateDNS) != "false" {
			ip := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11ySvcPscEndpointIP)
			if ip == "" {
				return "", fmt.Errorf("o11y service psc endpoint ip annotation is empty, waiting for privatelink task to populate it")
			}
			// Use IP address directly for VMImport URL
			return fmt.Sprintf("http://%s/%s/api/v1/import", ip, o11yID), nil
		}
	}

	// Standard build method
	dnsName := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yServiceDNSName)
	builder := o11y.NewWriteURLBuilder(dnsName)
	switch t.Config.O11yConfig.Vendor {
	case consts.VendorTypeAWS:
		return builder.ForAWS().BuildVMImportV2(o11yID), nil
	case consts.VendorTypeGCP:
		endpointName := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11ySvcInterfaceEndpointName)
		return builder.ForGCP(endpointName).BuildVMImportV2(o11yID), nil
	case consts.VendorTypeAzure:
		return builder.ForAzure().BuildVMImportV2(o11yID), nil
	default:
		return "", fmt.Errorf("unsupported vendor: %s", t.Config.O11yConfig.Vendor)
	}
}

func (t *VectorTask) buildRole() (*rbacv1.Role, error) {
	return &rbacv1.Role{
		ObjectMeta: metav1.ObjectMeta{
			Name:      consts.O11yNameClusterVector,
			Namespace: t.buildNamespace(),
		},
		Rules: []rbacv1.PolicyRule{{
			APIGroups: []string{""},
			Resources: []string{"pods"},
			Verbs:     []string{"get", "list", "watch"},
		}},
	}, nil
}

func (t *VectorTask) buildRoleBinding() (*rbacv1.RoleBinding, error) {
	name := consts.O11yNameClusterVector
	namespace := t.buildNamespace()
	return &rbacv1.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
		Subjects: []rbacv1.Subject{{
			Kind:      "ServiceAccount",
			Name:      name,
			Namespace: namespace,
		}},
		RoleRef: rbacv1.RoleRef{
			Kind:     "Role",
			Name:     name,
			APIGroup: "rbac.authorization.k8s.io",
		},
	}, nil
}

func (t *VectorTask) buildNamespace() string {
	return fmt.Sprintf("tidb%s", t.Agent.Spec.O11y.ApplicationID)
}