package task

import (
	"context"
	"fmt"
	"strings"

	agentv2 "github.com/tidbcloud/observability/agent-control-server/api/v2"
	"github.com/tidbcloud/observability/agent-control-server/controllers/o11yk8sagent/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/reconciler"
	"github.com/tidbcloud/observability/agent-control-server/pkg/util"
)

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

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

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

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

const (
	// EKSRoleAnnotation is the annotation key for EKS IAM role
	EKSRoleAnnotation = "eks.amazonaws.com/role-arn"
)

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, consts.O11ySystemNamespace); err != nil {
		return deploy.Panic(err) // this might happen when the tidb namespace is not created
	}
	// Remove eks.amazonaws.com/role-arn annotation from the ServiceAccount if it exists
	// This is needed because we no longer use IRSA for Vector
	// If we set the inheritedIAM
	if strings.Contains(util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yAgentRoleArn), "writer") {
		_, err := agentbuild.CheckAndUpdateServiceAccountAnnotations(ctx, cli, consts.O11yNameVector, consts.O11ySystemNamespace, EKSRoleAnnotation)
		if err != nil {
			return deploy.Panic(err)
		}
	}
	cfg := t.Config.O11yConfig.CustomerConfigsV2.Dedicated.K8s.BuildConfig().Vector
	return deploy.PanicOrDone(agentbuild.InstallOrUpgradeHelm(ctx, t.Remote, t.build, cfg.Chart.Name, cfg.Chart.Version, consts.O11yNameVector, consts.O11ySystemNamespace, nil))
}

func (t *VectorTask) PostCreate(ctx context.Context) deploy.Result {
	cli, err := t.Remote.GetClient(ctx)
	if err != nil {
		return deploy.Panic(err)
	}
	return deploy.PanicOrDone(agentbuild.CheckDaemonSetReady(ctx, cli, consts.O11yNameVector, consts.O11ySystemNamespace))
}

func (t *VectorTask) Delete(ctx context.Context) deploy.Result {
	return deploy.PanicOrDone(agentbuild.DeleteHelm(ctx, t.Remote, consts.O11yNameVector, consts.O11ySystemNamespace))
}

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

func (t *VectorTask) build() (string, error) {
	cfg := t.Config.O11yConfig.CustomerConfigsV2.Dedicated.K8s.BuildConfig().Vector

	imageRepo, imageTag, err := agentbuild.BuildVectorImage(t.Agent, cfg.Image)
	if err != nil {
		return "", err
	}
	resourcesJson, err := agentbuild.BuildVectorResources(t.Agent, cfg.Resources.DeepCopy())
	if err != nil {
		return "", err
	}

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

	values := map[string]string{
		"ImageRepo":         imageRepo,
		"ImageTag":          imageTag,
		"Resources":         resourcesJson,
		"PriorityClassName": consts.PriorityClassNameNodeCritical,
		"Env":               t.Config.O11yConfig.Env,
		"Provider":          consts.VendorToProvider(t.Config.O11yConfig.Vendor), // eks, gke, aks
		"Region":            t.Config.O11yConfig.Region,
		"Bucket":            util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yLogBucketName),
		"ClusterName":       t.Agent.Spec.KubernetesClusterName,
		"TenantID":          t.Agent.Spec.O11y.TenantID,
		"ProjectID":         t.Agent.Spec.O11y.ProjectID,
		"LokiURL":           lokiurl,
		"RoleARN":           util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yAgentRoleArn),
		"AzureClientID":     util.GetAnnotation(t.Agent, consts.AnnotationKeyAKSKubeletClientID),
		"ContainerName":     util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yLogBucketName),
		"StorageAccount":    util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yLogStorageAccount),
	}
	return agentbuild.BuildVectorValues(tmpl, values)
}

func (t *VectorTask) buildLokiURL() (string, error) {
	// 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("o11y service interface endpoint name annotation is empty")
			}
			dns, err := awscrossplane.GetVPCEndpointDNSByEndpointName(context.Background(), t.Local, endpointName)
			if err != nil {
				return "", err
			}
			return fmt.Sprintf("http://%s/loki", dns), 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")
			}
			return fmt.Sprintf("http://%s/loki", ip), nil
		}
	}
	// Other Vendor use service private DNS name
	return fmt.Sprintf("http://%s/loki", util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yServiceDNSName)), nil
}