package task

import (
	"context"
	"fmt"
	"maps"
	"regexp"
	"sort"
	"strings"

	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"

	agentv2 "github.com/tidbcloud/observability/agent-control-server/api/v2"
	"github.com/tidbcloud/observability/agent-control-server/controllers/sharedpoolagent/metadata"
	"github.com/tidbcloud/observability/agent-control-server/pkg/agentbuild"
	"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 vectorTarget struct {
	InstanceID  string
	ReleaseName string
}

type VectorTask struct {
	*reconciler.BaseTask[*agentv2.SharedPoolAgent]
	targets []vectorTarget // for PostCreate check
}

func NewVectorTask(ctx *reconciler.Context[*agentv2.SharedPoolAgent]) *VectorTask {
	return &VectorTask{
		BaseTask: &reconciler.BaseTask[*agentv2.SharedPoolAgent]{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)
	}
	return deploy.PanicOrDone(t.create(ctx, cli))
}

func (t *VectorTask) PostCreate(ctx context.Context) deploy.Result {
	cli, err := t.Remote.GetClient(ctx)
	if err != nil {
		return deploy.Panic(err)
	}
	for _, target := range t.targets {
		if err := agentbuild.CheckStatefulSetReady(ctx, cli, target.ReleaseName, t.buildNamespace()); err != nil {
			return deploy.Panic(err)
		}
		if t.Config.IsServerlessMode() {
			continue
		}
		labels := map[string]string{
			consts.TagKeyComponent:    consts.TagValueComponentO11y,
			consts.TagKeyTenant:       t.Agent.Spec.O11y.TenantID,
			consts.TagKeyResourcePool: t.Agent.Spec.SharedPoolID,
		}
		if _, err := agentbuild.CheckAndUpdatePVCLabels(ctx, cli, t.buildVectorPVCName(target.ReleaseName), t.buildNamespace(), labels); err != nil {
			return deploy.Panic(err)
		}
	}
	return deploy.Done()
}

func (t *VectorTask) Delete(ctx context.Context) deploy.Result {
	cli, err := t.Remote.GetClient(ctx)
	if err != nil {
		return deploy.Error(err)
	}
	targets, err := t.buildTargets(ctx, cli)
	if err != nil {
		return deploy.Panic(err)
	}
	for _, target := range targets {
		name := target.ReleaseName
		namespace := t.buildNamespace()

		if err := agentbuild.DeleteHelm(ctx, t.Remote, name, namespace); err != nil {
			return deploy.Panic(err)
		}
		if err := agentbuild.DeleteRoleBinding(ctx, cli, name, namespace); err != nil {
			return deploy.Panic(err)
		}
		if err := agentbuild.DeleteRole(ctx, cli, name, namespace); err != nil {
			return deploy.Panic(err)
		}
		if err := t.deleteVectorPVCs(ctx, cli, target.ReleaseName); err != nil {
			return deploy.Panic(err)
		}
	}

	if t.Config.IsServerlessMode() {
		if err := t.deleteServerlessTLSSecret(ctx, cli); err != nil {
			return deploy.Panic(err)
		}
	}
	return deploy.Done()
}

func (t *VectorTask) PostDelete(ctx context.Context) deploy.Result {
	cli, err := t.Remote.GetClient(ctx)
	if err != nil {
		return deploy.Error(err)
	}
	targets, err := t.buildTargets(ctx, cli)
	if err != nil {
		return deploy.Panic(err)
	}
	for _, target := range targets {
		if err := agentbuild.CheckStatefulSetDeleted(ctx, cli, target.ReleaseName, t.buildNamespace()); err != nil {
			return deploy.Panic(err)
		}
	}
	return deploy.Done()
}

func (t *VectorTask) create(ctx context.Context, cli client.Client) error {
	// TODO: temporary use ak/sk on alicloud
	// When vector supports node role for oss, remove this
	var accessKey, secretKey string
	if t.Config.O11yConfig.Vendor == "alicloud" {
		secret := &corev1.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "o11y-oss-secret",
				Namespace: "observability",
			},
		}
		if err := t.Local.Get(context.Background(), client.ObjectKeyFromObject(secret), secret); err != nil {
			return err
		}
		accessKey = string(secret.Data["access_key"])
		secretKey = string(secret.Data["secret_key"])
	}

	dnsName := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yServiceDNSName)
	remoteWriteURL := fmt.Sprintf("http://%s/multitenant/api/v1/write", dnsName)

	cfg := t.Config.O11yConfig.CustomerConfigsV2.NextGen.SharedPool.BuildConfig().Vector
	imageRepo, imageTag, err := agentbuild.BuildVectorImage(t.Agent, cfg.Image)
	if err != nil {
		return err
	}

	desiredSpec := util.GetComponentDesiredSpec(t.Agent.GetAnnotations(), "default", consts.ComponentVector)
	spec := t.Config.O11yConfig.CustomerConfigsV2.NextGen.SharedPool.Specs.Select(desiredSpec)
	if spec == nil || spec.Vector == 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("sharedpool/vector-values-%s.yaml", t.Config.O11yConfig.Vendor)
	tmpl, err := agentbuild.BuildVectorTemplate(t.Agent, t.Templates, configTemplate)
	if err != nil {
		return err
	}

	targets, err := t.buildTargets(ctx, cli)
	if err != nil {
		return err
	}
	if len(targets) == 0 {
		return nil
	}
	keyspaceToVmTenants, err := t.buildKeyspaceToVmTenants(ctx)
	if err != nil {
		return err
	}

	var serverlessExtValues map[string]string
	if t.Config.IsServerlessMode() {
		if serverlessExtValues, err = t.buildServerlessExtValues(ctx, cli, cfg); err != nil {
			return err
		}
	}

	for _, target := range targets {
		name := target.ReleaseName
		namespace := t.buildNamespace()
		builder := func() (string, error) {
			values := map[string]string{
				"ImageRepo":           imageRepo,
				"ImageTag":            imageTag,
				"Resources":           resourcesJson,
				"InstanceID":          target.InstanceID,
				"KeyspaceToVMTenants": keyspaceToVmTenants,
				"RemoteWriteURL":      remoteWriteURL,
				"AccessKeyID":         accessKey,
				"SecretAccessKey":     secretKey,
				"SecretName":          name,
				"PriorityClassName":   consts.PriorityClassNameClusterCritical,
				"SharedPoolID":        t.Agent.Spec.SharedPoolID,
				"Region":              t.Config.O11yConfig.Region,
				"TenantID":            t.Agent.Spec.O11y.TenantID,
				"Bucket":              util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yLogBucketName),
				"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),
				"RefNamespace":        t.Agent.Namespace,
				"RefName":             t.Agent.Name,
			}
			maps.Copy(values, serverlessExtValues)
			return agentbuild.BuildVectorValues(tmpl, values)
		}
		if err := agentbuild.CreateOrUpdateRole(ctx, cli, t.buildRole(name, namespace)); err != nil {
			return err
		}
		if err := agentbuild.CreateOrUpdateRoleBinding(ctx, cli, t.buildRoleBinding(name, namespace)); err != nil {
			return err
		}
		if err := agentbuild.InstallOrUpgradeHelm(ctx, t.Remote, builder, cfg.Chart.Name, cfg.Chart.Version, name, namespace, nil); err != nil {
			return err
		}
	}
	// TODO: Remove unnecessary helm releases
	t.targets = targets // for later PostCreate check
	return nil
}

var regexpTiKVPodNamePrefix = regexp.MustCompile(`^db-(\w+)-tikv-`)

func (t *VectorTask) buildInstanceIDs(ctx context.Context, cli client.Client) (map[string]struct{}, error) {
	pods := &corev1.PodList{}
	if err := cli.List(ctx, pods, client.InNamespace(t.buildNamespace()), client.MatchingLabels{
		"pingcap.com/component": "tikv",
	}); err != nil {
		return nil, fmt.Errorf("failed to list pods: %w", err)
	}
	if len(pods.Items) == 0 {
		return nil, nil
	}
	instanceIDs := map[string]struct{}{}
	for _, pod := range pods.Items {
		matches := regexpTiKVPodNamePrefix.FindStringSubmatch(pod.Name)
		if len(matches) != 2 {
			continue
		}
		instanceIDs[matches[1]] = struct{}{}
	}
	return instanceIDs, nil
}

func (t *VectorTask) buildTargets(ctx context.Context, cli client.Client) ([]vectorTarget, error) {
	if t.Config.IsServerlessMode() {
		return []vectorTarget{{ReleaseName: serverlessVectorName}}, nil
	}

	instanceIDs, err := t.buildInstanceIDs(ctx, cli)
	if err != nil {
		return nil, err
	}
	return t.buildTargetsFromInstanceIDs(instanceIDs), nil
}

func (t *VectorTask) buildTargetsFromInstanceIDs(instanceIDs map[string]struct{}) []vectorTarget {
	if len(instanceIDs) == 0 {
		return nil
	}

	targets := make([]vectorTarget, 0, len(instanceIDs))
	for id := range instanceIDs {
		targets = append(targets, vectorTarget{
			InstanceID:  id,
			ReleaseName: t.buildName(id),
		})
	}
	sort.Slice(targets, func(i, j int) bool {
		return targets[i].InstanceID < targets[j].InstanceID
	})
	return targets
}

func (t *VectorTask) listVectorPVCNames(ctx context.Context, cli client.Client, releaseName, namespace string) ([]string, error) {
	pvcList := &corev1.PersistentVolumeClaimList{}
	if err := cli.List(ctx, pvcList, client.InNamespace(namespace)); err != nil {
		return nil, fmt.Errorf("failed to list vector pvcs in %s: %w", namespace, err)
	}

	prefix := t.buildVectorPVCPrefix(releaseName)
	names := make([]string, 0, len(pvcList.Items))
	for _, pvc := range pvcList.Items {
		if strings.HasPrefix(pvc.Name, prefix) {
			names = append(names, pvc.Name)
		}
	}
	sort.Strings(names)
	return names, nil
}

func (t *VectorTask) deleteVectorPVCs(ctx context.Context, cli client.Client, releaseName string) error {
	namespace := t.buildNamespace()
	pvcNames, err := t.listVectorPVCNames(ctx, cli, releaseName, namespace)
	if err != nil {
		return err
	}
	for _, pvcName := range pvcNames {
		if err := agentbuild.DeletePVC(ctx, cli, pvcName, namespace); err != nil {
			return err
		}
	}
	return nil
}

type keyspaceVmTenant struct {
	keyspace    string
	vmAccountID string
	vmProjectID string
}

func (t *VectorTask) buildKeyspaceToVmTenants(ctx context.Context) (string, error) {
	tidbgroupagents := &agentv2.TiDBGroupAgentList{}
	if err := t.Local.List(ctx, tidbgroupagents); err != nil {
		return "", fmt.Errorf("failed to list tidbgroupagents: %w", err)
	}
	sort.SliceStable(tidbgroupagents.Items, func(i, j int) bool {
		if tidbgroupagents.Items[i].Spec.ClusterID != tidbgroupagents.Items[j].Spec.ClusterID {
			return tidbgroupagents.Items[i].Spec.ClusterID < tidbgroupagents.Items[j].Spec.ClusterID
		}
		return tidbgroupagents.Items[i].Name < tidbgroupagents.Items[j].Name
	})
	keyspaceVmTenants := []keyspaceVmTenant{}
	for _, tga := range tidbgroupagents.Items {
		if tga.Spec.SharedPoolID != t.Agent.Spec.SharedPoolID {
			continue
		}
		keyspaceName := tga.Spec.ClusterID
		vmAccountID := util.GetAnnotation(&tga, consts.AnnotationKeyO11yVMAccountID)
		vmProjectID := util.GetAnnotation(&tga, consts.AnnotationKeyO11yVMProjectID)
		if keyspaceName == "" || vmAccountID == "" || vmProjectID == "" {
			continue
		}
		keyspaceVmTenants = append(keyspaceVmTenants, keyspaceVmTenant{
			keyspace:    keyspaceName,
			vmAccountID: vmAccountID,
			vmProjectID: vmProjectID,
		})
	}
	if len(keyspaceVmTenants) == 0 {
		return "", nil
	}
	var s string
	for _, k := range keyspaceVmTenants {
		s += fmt.Sprintf("%s:%s:%s,", k.keyspace, k.vmAccountID, k.vmProjectID)
	}
	s = strings.TrimSuffix(s, ",")
	return s, nil
}

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

func (t *VectorTask) buildRoleBinding(name, namespace string) func() (*rbacv1.RoleBinding, error) {
	return func() (*rbacv1.RoleBinding, error) {
		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) buildName(instanceID string) string {
	return fmt.Sprintf("o11y-vector-%s", instanceID)
}

func (t *VectorTask) buildNamespace() string {
	if t.Config.IsServerlessMode() {
		return consts.ServerlessObservabilityNamespace
	}
	return fmt.Sprintf("tidb%s", t.Agent.Spec.SharedPoolID)
}

func (t *VectorTask) buildVectorPVCName(releaseName string) string {
	return fmt.Sprintf("data-%s-0", releaseName)
}

func (t *VectorTask) buildVectorPVCPrefix(releaseName string) string {
	return fmt.Sprintf("data-%s-", releaseName)
}