package tasks

import (
	"bytes"
	"context"
	"io"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	v1 "github.com/tidbcloud/observability/api/regionalserver/resourceworker/v1"
	"github.com/tidbcloud/observability/pkg/common/message"
	"github.com/tidbcloud/observability/pkg/resourceworker/providers/types"
	"gopkg.in/yaml.v2"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func TestClusterIDFromDiagnosisTaskKey(t *testing.T) {
	clusterID, ok := clusterIDFromDiagnosisTaskKey("diagnosis/tasks/processor-dedicated/cluster-1.yaml")
	assert.True(t, ok)
	assert.Equal(t, "cluster-1", clusterID)

	_, ok = clusterIDFromDiagnosisTaskKey("diagnosis/tasks/processor-dedicated/")
	assert.False(t, ok)

	_, ok = clusterIDFromDiagnosisTaskKey("diagnosis/tasks/processor-dedicated/nested/cluster-1.yaml")
	assert.False(t, ok)

	_, ok = clusterIDFromDiagnosisTaskKey("other/prefix/cluster-1.yaml")
	assert.False(t, ok)
}

func TestClusterIDFromHistoryTaskKey(t *testing.T) {
	clusterID, ok := clusterIDFromTaskKey(diagnosisHistoryTaskPrefix, "diagnosis/tasks/history-tasks/cluster-1.yaml")
	assert.True(t, ok)
	assert.Equal(t, "cluster-1", clusterID)

	_, ok = clusterIDFromTaskKey(diagnosisHistoryTaskPrefix, "diagnosis/tasks/processor-dedicated/cluster-1.yaml")
	assert.False(t, ok)
}

func TestHistoryTaskKey(t *testing.T) {
	assert.Equal(t, "diagnosis/tasks/history-tasks/app-1.yaml", historyTaskKey("app-1"))
}

func TestBuildO11yDiagnosisTask(t *testing.T) {
	cluster := v1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name: "o11y-cluster-1",
			Labels: map[string]string{
				"observability.pingcap.com/bizType":       "tidbcloud",
				"observability.pingcap.com/applicationID": "app-1",
				"observability.pingcap.com/tenantID":      "tenant-1",
				"observability.pingcap.com/projectID":     "project-1",
				"observability.pingcap.com/vendor":        "aws",
				"observability.pingcap.com/region":        "us-west-2",
			},
			CreationTimestamp: metav1.NewTime(time.Unix(1700000000, 0)),
		},
	}

	task, ok := buildO11yDiagnosisTask(cluster)
	assert.True(t, ok)
	assert.Equal(t, "app-1", task.Meta.ClusterID)
	assert.Equal(t, "dedicated", task.Meta.ClusterType)
	assert.Equal(t, "project-1", task.Meta.ProjectID)

	cluster.ObjectMeta.Labels["observability.pingcap.com/bizType"] = "unknown"
	_, ok = buildO11yDiagnosisTask(cluster)
	assert.False(t, ok)
}

func TestDiagnosisTaskKey(t *testing.T) {
	assert.Equal(t, "diagnosis/tasks/processor-dedicated/app-1.yaml", diagnosisTaskKey("app-1"))
}

type memoryS3 struct {
	objects map[string][]byte
}

func newMemoryS3() *memoryS3 {
	return &memoryS3{objects: map[string][]byte{}}
}

func (m *memoryS3) CreateBucket(context.Context, string) error { return nil }
func (m *memoryS3) BucketExists(context.Context, string) (bool, error) {
	return true, nil
}
func (m *memoryS3) AddLifecycle(context.Context, string, []message.LifecycleRuleParam, []string) error {
	return nil
}
func (m *memoryS3) RemoveLifecycle(context.Context, string, string) error { return nil }

func (m *memoryS3) Get(_ context.Context, _ string, objectName string) (*io.ReadCloser, error) {
	data, ok := m.objects[objectName]
	if !ok {
		return nil, types.ObjectNotFoundError
	}
	rc := io.NopCloser(bytes.NewReader(data))
	return &rc, nil
}

func (m *memoryS3) Put(_ context.Context, _ string, objectName string, reader *bytes.Reader) error {
	data, err := io.ReadAll(reader)
	if err != nil {
		return err
	}
	m.objects[objectName] = data
	return nil
}

func (m *memoryS3) ListObjects(_ context.Context, _ string, prefix string, _ types.ListOptions) ([]string, error) {
	keys := make([]string, 0)
	for key := range m.objects {
		if len(key) >= len(prefix) && key[:len(prefix)] == prefix {
			keys = append(keys, key)
		}
	}
	return keys, nil
}

func (m *memoryS3) Delete(_ context.Context, _ string, objectNames ...string) error {
	for _, name := range objectNames {
		delete(m.objects, name)
	}
	return nil
}

func TestArchiveStaleDiagnosisTasksMovesToHistory(t *testing.T) {
	ctx := context.Background()
	store := newMemoryS3()
	stale := O11yDiagnosisTask{Meta: Meta{ClusterID: "gone-1", TenantID: "t1", ProjectID: "p1", ClusterType: "dedicated"}}
	data, err := yaml.Marshal(stale)
	assert.NoError(t, err)
	assert.NoError(t, store.Put(ctx, "bucket", diagnosisTaskKey("gone-1"), bytes.NewReader(data)))
	assert.NoError(t, store.Put(ctx, "bucket", diagnosisTaskKey("alive-1"), bytes.NewReader(data)))

	active := map[string]O11yDiagnosisTask{"alive-1": {Meta: Meta{ClusterID: "alive-1"}}}
	assert.NoError(t, archiveStaleDiagnosisTasks(ctx, store, "bucket", active))

	_, err = store.Get(ctx, "bucket", diagnosisTaskKey("gone-1"))
	assert.Error(t, err)
	rc, err := store.Get(ctx, "bucket", historyTaskKey("gone-1"))
	assert.NoError(t, err)
	_ = (*rc).Close()
	rc, err = store.Get(ctx, "bucket", diagnosisTaskKey("alive-1"))
	assert.NoError(t, err)
	_ = (*rc).Close()
}

func TestRemoveActiveHistoryTasks(t *testing.T) {
	ctx := context.Background()
	store := newMemoryS3()
	task := O11yDiagnosisTask{Meta: Meta{ClusterID: "alive-1"}}
	data, err := yaml.Marshal(task)
	assert.NoError(t, err)
	assert.NoError(t, store.Put(ctx, "bucket", historyTaskKey("alive-1"), bytes.NewReader(data)))
	assert.NoError(t, store.Put(ctx, "bucket", historyTaskKey("gone-1"), bytes.NewReader(data)))

	assert.NoError(t, removeActiveHistoryTasks(ctx, store, "bucket", map[string]O11yDiagnosisTask{
		"alive-1": task,
	}))

	_, err = store.Get(ctx, "bucket", historyTaskKey("alive-1"))
	assert.Error(t, err)
	rc, err := store.Get(ctx, "bucket", historyTaskKey("gone-1"))
	assert.NoError(t, err)
	_ = (*rc).Close()
}

func TestLoadHistoryDiagnosisTasks(t *testing.T) {
	ctx := context.Background()
	store := newMemoryS3()
	task := O11yDiagnosisTask{Meta: Meta{ClusterID: "gone-1", TenantID: "t1", ProjectID: "p1", ClusterType: "dedicated"}}
	data, err := yaml.Marshal(task)
	assert.NoError(t, err)
	assert.NoError(t, store.Put(ctx, "bucket", historyTaskKey("gone-1"), bytes.NewReader(data)))
	assert.NoError(t, store.Put(ctx, "bucket", diagnosisHistoryTaskPrefix+"ignore.txt", bytes.NewReader([]byte("nope"))))

	loaded, err := loadHistoryDiagnosisTasks(ctx, store, "bucket")
	assert.NoError(t, err)
	assert.Len(t, loaded, 1)
	assert.Equal(t, "gone-1", loaded[0].Meta.ClusterID)
}