package task

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"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/testkits"
)

func makeServerlessTLSSecret(t *testing.T, name, namespace string, pkcs8 bool) *corev1.Secret {
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	assert.NoError(t, err)

	privateKeyBlock := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
	}
	if pkcs8 {
		pkcs8Bytes, marshalErr := x509.MarshalPKCS8PrivateKey(privateKey)
		assert.NoError(t, marshalErr)
		privateKeyBlock = &pem.Block{
			Type:  "PRIVATE KEY",
			Bytes: pkcs8Bytes,
		}
	}

	return &corev1.Secret{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
		Type: corev1.SecretTypeTLS,
		Data: map[string][]byte{
			"ca.crt":                []byte("ca"),
			corev1.TLSCertKey:       []byte("cert"),
			corev1.TLSPrivateKeyKey: pem.EncodeToMemory(privateKeyBlock),
		},
	}
}

func TestVectorTask_serverlessRenderTemplateValues(t *testing.T) {
	values := map[string]string{
		"ImageRepo":         "385595570414.dkr.ecr.us-west-2.amazonaws.com/tidbcloud/vector",
		"ImageTag":          "0.49-1124-ng-conprof-debian",
		"Replicas":          "3",
		"Resources":         `{"limits":{"cpu":"1","memory":"1Gi"},"requests":{"cpu":"100m","memory":"128Mi"}}`,
		"PriorityClassName": consts.PriorityClassNameClusterCritical,
		"SharedPoolID":      "pool-456",
		"PDAddress":         "https://" + serverlessVectorDefaultPDAddress,
		"TLSSecretName":     serverlessTLSSecretName,
		"ManagerServerURL":  serverlessVectorDefaultMgrAddr,
		"TiDBNamespace":     serverlessVectorDefaultNamespace,
		"EnableTiKVTopSQL":  "false",
		"Bucket":            "o11y-prod-shared-us-east-1-serverless",
		"Region":            "us-east-1",
		"RoleARN":           "arn:aws:iam::557537366020:role/o11yA-NiQoWCEZPLmsaqBp8g2jr4",
		"RefNamespace":      "observability",
		"RefName":           "dev-us-east-1-f01",
	}

	assertRenderTemplateValues(t, "serverless/sharedpool/vector-values-aws.yaml", values, `
role: Aggregator
replicas: 3
image:
  repository: "385595570414.dkr.ecr.us-west-2.amazonaws.com/tidbcloud/vector"
  tag: "0.49-1124-ng-conprof-debian"
resources: {"limits":{"cpu":"1","memory":"1Gi"},"requests":{"cpu":"100m","memory":"128Mi"}}
env:
- name: VECTOR_STS_REPLICA_COUNT
  value: "3"
- name: VECTOR_STS_ID
  valueFrom:
    fieldRef:
      fieldPath: metadata.labels['apps.kubernetes.io/pod-index']
podPriorityClassName: "tidb-cluster-o11y-critical"
podLabels:
  vector.dev/exclude: "false"
  controller/component: "vector"
  controller/agent: "sharedpool"
  controller/ref-namespace: "observability"
  controller/ref-name: "dev-us-east-1-f01"
podSecurityContext:
  runAsNonRoot: true
  runAsUser: 1000
  runAsGroup: 2000
  fsGroup: 2000
securityContext:
  readOnlyRootFilesystem: true
  allowPrivilegeEscalation: false
  privileged: false
  capabilities:
    drop:
    - NET_RAW
persistence:
  enabled: true
  size: 10Gi
  storageClassName: o11y-metric-agent
  hostPath:
    enabled: false
service:
  enabled: false
extraVolumes:
- name: tidb-tls
  secret:
    secretName: "o11y-vector-sharedpool-tidb-tls"
    items:
    - key: ca.crt
      path: ca.crt
    - key: tls.crt
      path: tls.crt
    - key: tls.key
      path: tls.key
extraVolumeMounts:
- name: tidb-tls
  mountPath: /etc/vector/tidb-tls
dataDir: /vector-data-dir
customConfig:
  data_dir: /vector-data-dir
  api:
    enabled: true
    address: 127.0.0.1:8686
    playground: false
  sources:
    internal_metrics:
      type: internal_metrics
    topsql_v2:
      type: topsql_v2
      downsampling_interval: 60
      top_n: 100
      pd_address: "https://serverless-cluster-pd.tidb-serverless.svc:2379"
      manager_server_address: "https://manager-server.tidb-admin.svc:8000"
      tidb_namespace: "super-vip-tidb-pool,canary-super-vip-tidb-pool"
      enable_tikv_topsql: false
      topology_fetch_interval_seconds: 10
      init_retry_delay_seconds: 5
      tls:
        ca_file: "/etc/vector/tidb-tls/ca.crt"
        crt_file: "/etc/vector/tidb-tls/tls.crt"
        key_file: "/etc/vector/tidb-tls/tls.key"
        verify_hostname: false
        verify_certificate: false
      topru:
        enable: true
        report_interval_seconds: 60
        item_interval_seconds: 60
  transforms:
    route_by_type:
      type: route
      inputs:
      - topsql_v2
      route:
        non_meta_logs: .source_table == "tidb_topsql" || .source_table == "tikv_topsql" || .source_table == "tikv_topregion" || .source_table == "topsql_topru"
        meta_logs: .source_table == "topsql_sql_meta" || .source_table == "topsql_plan_meta"
  sinks:
    prom_exporter:
      type: prometheus_exporter
      inputs:
      - internal_metrics
      address: 0.0.0.0:9090
    topsql_s3_data:
      type: topsql_data_deltalake
      inputs:
      - route_by_type.non_meta_logs
      base_path: "s3://o11y-prod-shared-us-east-1-serverless/deltalake/org=xxx/cluster=xxx/type=topsql"
      bucket: "o11y-prod-shared-us-east-1-serverless"
      region: "us-east-1"
      batch_size: 1000
      timeout_secs: 30
      max_delay_secs: 10
      enable_keyspace_cluster_mapping: true
      pd_address: "https://serverless-cluster-pd.tidb-serverless.svc:2379"
      pd_tls:
        ca_file: "/etc/vector/tidb-tls/ca.crt"
        crt_file: "/etc/vector/tidb-tls/tls.crt"
        key_file: "/etc/vector/tidb-tls/tls.key"
        verify_hostname: false
        verify_certificate: false
      auth:
        assume_role: "arn:aws:iam::557537366020:role/o11yA-NiQoWCEZPLmsaqBp8g2jr4"
    topsql_s3_meta:
      type: topsql_meta_deltalake
      inputs:
      - route_by_type.meta_logs
      base_path: "s3://o11y-prod-shared-us-east-1-serverless/deltalake/org=xxx/cluster=xxx/type=topsql"
      bucket: "o11y-prod-shared-us-east-1-serverless"
      region: "us-east-1"
      batch_size: 1000
      timeout_secs: 30
      max_delay_secs: 10
      meta_cache_capacity: 100000
      auth:
        assume_role: "arn:aws:iam::557537366020:role/o11yA-NiQoWCEZPLmsaqBp8g2jr4"
`)
}

func TestVectorTask_buildServerlessTLSSecret(t *testing.T) {
	cfg := baseVectorConfig()
	cfg.Mode = config.ModeServerless

	agent := testkits.NewSharedPoolAgentBuilder().Build()
	task := makeVectorTask(agent, testkits.TaskOptions{Config: cfg})
	sourceSecret := makeServerlessTLSSecret(t, serverlessVectorDefaultSecret, consts.ServerlessObservabilityNamespace, false)

	targetSecret, err := task.buildServerlessTLSSecret(sourceSecret)
	assert.NoError(t, err)
	assert.Equal(t, consts.ServerlessObservabilityNamespace, targetSecret.Namespace)
	assert.Equal(t, []byte("ca"), targetSecret.Data["ca.crt"])
	assert.Equal(t, []byte("cert"), targetSecret.Data[corev1.TLSCertKey])
	assert.True(t, strings.HasPrefix(string(targetSecret.Data[corev1.TLSPrivateKeyKey]), "-----BEGIN PRIVATE KEY-----"))
}

func TestVectorTask_deleteServerlessTLSSecret(t *testing.T) {
	tests := []struct {
		name    string
		config  config.RootConfig
		fakeCli client.Client
		wantErr bool
		verify  func(t *testing.T, cli client.Client)
	}{
		{
			name: "serverless mode deletes TLS secret",
			config: func() config.RootConfig {
				cfg := baseVectorConfig()
				cfg.Mode = config.ModeServerless
				return cfg
			}(),
			fakeCli: testkits.BuildClient(
				&corev1.Secret{
					ObjectMeta: metav1.ObjectMeta{
						Name:      serverlessTLSSecretName,
						Namespace: consts.ServerlessObservabilityNamespace,
					},
					Type: corev1.SecretTypeTLS,
				},
			),
			verify: func(t *testing.T, cli client.Client) {
				secret := &corev1.Secret{}
				err := cli.Get(context.Background(), types.NamespacedName{
					Name:      serverlessTLSSecretName,
					Namespace: consts.ServerlessObservabilityNamespace,
				}, secret)
				assert.Error(t, err)
			},
		},
		{
			name: "no error when TLS secret does not exist",
			config: func() config.RootConfig {
				cfg := baseVectorConfig()
				cfg.Mode = config.ModeServerless
				return cfg
			}(),
			fakeCli: testkits.BuildClient(),
		},
		{
			name: "returns error when delete fails",
			config: func() config.RootConfig {
				cfg := baseVectorConfig()
				cfg.Mode = config.ModeServerless
				return cfg
			}(),
			fakeCli: testkits.NewClientGetter(
				testkits.BuildClient(),
			).WithDeleteFunc(func(_ context.Context, _ client.Object, _ ...client.DeleteOption) error {
				return errors.New("delete failed")
			}),
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			agent := testkits.NewSharedPoolAgentBuilder().Build()
			task := makeVectorTask(agent, testkits.TaskOptions{Config: tt.config})
			err := task.deleteServerlessTLSSecret(context.Background(), tt.fakeCli)
			if tt.wantErr {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
			}
			if tt.verify != nil {
				tt.verify(t, tt.fakeCli)
			}
		})
	}
}

func TestNormalizePrivateKeyToPKCS8(t *testing.T) {
	pkcs1Secret := makeServerlessTLSSecret(t, "source-pkcs1", consts.ServerlessObservabilityNamespace, false)
	pkcs8Secret := makeServerlessTLSSecret(t, "source-pkcs8", consts.ServerlessObservabilityNamespace, true)

	tests := []struct {
		name       string
		privatePEM []byte
		wantErr    string
	}{
		{
			name:       "converts pkcs1 rsa key",
			privatePEM: pkcs1Secret.Data[corev1.TLSPrivateKeyKey],
		},
		{
			name:       "keeps pkcs8 key",
			privatePEM: pkcs8Secret.Data[corev1.TLSPrivateKeyKey],
		},
		{
			name:       "invalid pem data",
			privatePEM: []byte("not a pem block"),
			wantErr:    "failed to decode",
		},
		{
			name:       "unsupported pem type",
			privatePEM: pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: []byte("garbage")}),
			wantErr:    "unsupported",
		},
		{
			name:       "corrupted pkcs8 bytes",
			privatePEM: pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: []byte("garbage")}),
			wantErr:    "failed to parse PKCS#8",
		},
		{
			name:       "corrupted pkcs1 bytes",
			privatePEM: pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: []byte("garbage")}),
			wantErr:    "failed to parse PKCS#1",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			normalizedPEM, err := normalizePrivateKeyToPKCS8(tt.privatePEM)
			if tt.wantErr != "" {
				assert.Error(t, err)
				assert.Contains(t, err.Error(), tt.wantErr)
			} else {
				assert.NoError(t, err)
				assert.True(t, strings.HasPrefix(string(normalizedPEM), "-----BEGIN PRIVATE KEY-----"))
			}
		})
	}
}

func TestEnsureHTTPSAddress(t *testing.T) {
	tests := []struct {
		name    string
		address string
		want    string
	}{
		{name: "prepends https when no scheme", address: "example.com:2379", want: "https://example.com:2379"},
		{name: "keeps https address unchanged", address: "https://example.com:2379", want: "https://example.com:2379"},
		{name: "keeps http address unchanged", address: "http://example.com", want: "http://example.com"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			assert.Equal(t, tt.want, ensureHTTPSAddress(tt.address))
		})
	}
}