package testkits

import (
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	agentv2 "github.com/tidbcloud/observability/agent-control-server/api/v2"
)

// KubernetesAgentBuilder builds KubernetesAgent instances for testing.
type KubernetesAgentBuilder struct {
	agent *agentv2.KubernetesAgent
}

// NewKubernetesAgentBuilder creates a new KubernetesAgentBuilder with default values.
func NewKubernetesAgentBuilder() *KubernetesAgentBuilder {
	return &KubernetesAgentBuilder{
		agent: &agentv2.KubernetesAgent{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "test-agent",
				Namespace: "default",
			},
			Spec: agentv2.KubernetesAgentSpec{
				KubernetesClusterName: "test-cluster",
				O11y: agentv2.O11yConfigV2{
					TenantID:      "tenant-1",
					ApplicationID: "app-1",
					BizType:       "tidbcloud",
				},
			},
		},
	}
}

// WithAnnotations sets the annotations on the agent.
func (b *KubernetesAgentBuilder) WithAnnotations(annotations map[string]string) *KubernetesAgentBuilder {
	if b.agent.Annotations == nil {
		b.agent.Annotations = make(map[string]string)
	}
	for k, v := range annotations {
		b.agent.Annotations[k] = v
	}
	return b
}

// WithName sets the name of the agent.
func (b *KubernetesAgentBuilder) WithName(name string) *KubernetesAgentBuilder {
	b.agent.Name = name
	return b
}

// WithNamespace sets the namespace of the agent.
func (b *KubernetesAgentBuilder) WithNamespace(ns string) *KubernetesAgentBuilder {
	b.agent.Namespace = ns
	return b
}

// WithClusterName sets the Kubernetes cluster name.
func (b *KubernetesAgentBuilder) WithClusterName(name string) *KubernetesAgentBuilder {
	b.agent.Spec.KubernetesClusterName = name
	return b
}

// WithTenantID sets the tenant ID.
func (b *KubernetesAgentBuilder) WithTenantID(id string) *KubernetesAgentBuilder {
	b.agent.Spec.O11y.TenantID = id
	return b
}

// WithApplicationID sets the application ID.
func (b *KubernetesAgentBuilder) WithApplicationID(id string) *KubernetesAgentBuilder {
	b.agent.Spec.O11y.ApplicationID = id
	return b
}

// Build returns the built KubernetesAgent.
func (b *KubernetesAgentBuilder) Build() *agentv2.KubernetesAgent {
	return b.agent
}

// SharedPoolAgentBuilder builds SharedPoolAgent instances for testing.
type SharedPoolAgentBuilder struct {
	agent *agentv2.SharedPoolAgent
}

// NewSharedPoolAgentBuilder creates a new SharedPoolAgentBuilder with default values.
func NewSharedPoolAgentBuilder() *SharedPoolAgentBuilder {
	return &SharedPoolAgentBuilder{
		agent: &agentv2.SharedPoolAgent{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "test-agent",
				Namespace: "default",
			},
			Spec: agentv2.SharedPoolAgentSpec{
				SharedPoolID:          "pool-456",
				KubernetesClusterName: "test-cluster",
				O11y: agentv2.O11yConfigV2{
					TenantID:      "tenant-1",
					ApplicationID: "app-1",
					BizType:       "tidbcloud",
				},
			},
		},
	}
}

// WithAnnotations sets the annotations on the agent.
func (b *SharedPoolAgentBuilder) WithAnnotations(annotations map[string]string) *SharedPoolAgentBuilder {
	if b.agent.Annotations == nil {
		b.agent.Annotations = make(map[string]string)
	}
	for k, v := range annotations {
		b.agent.Annotations[k] = v
	}
	return b
}

// WithName sets the name of the agent.
func (b *SharedPoolAgentBuilder) WithName(name string) *SharedPoolAgentBuilder {
	b.agent.Name = name
	return b
}

// WithNamespace sets the namespace of the agent.
func (b *SharedPoolAgentBuilder) WithNamespace(ns string) *SharedPoolAgentBuilder {
	b.agent.Namespace = ns
	return b
}

// WithSharedPoolID sets the shared pool ID.
func (b *SharedPoolAgentBuilder) WithSharedPoolID(id string) *SharedPoolAgentBuilder {
	b.agent.Spec.SharedPoolID = id
	return b
}

// WithClusterName sets the Kubernetes cluster name.
func (b *SharedPoolAgentBuilder) WithClusterName(name string) *SharedPoolAgentBuilder {
	b.agent.Spec.KubernetesClusterName = name
	return b
}

// Build returns the built SharedPoolAgent.
func (b *SharedPoolAgentBuilder) Build() *agentv2.SharedPoolAgent {
	return b.agent
}

// TiDBGroupAgentBuilder builds TiDBGroupAgent instances for testing.
type TiDBGroupAgentBuilder struct {
	agent *agentv2.TiDBGroupAgent
}

// NewTiDBGroupAgentBuilder creates a new TiDBGroupAgentBuilder with default values.
func NewTiDBGroupAgentBuilder() *TiDBGroupAgentBuilder {
	return &TiDBGroupAgentBuilder{
		agent: &agentv2.TiDBGroupAgent{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "test-agent",
				Namespace: "default",
			},
			Spec: agentv2.TiDBGroupAgentSpec{
				TiDBGroup:             "group-123",
				SharedPoolID:          "pool-456",
				ClusterID:             "cluster-1",
				KubernetesClusterName: "test-cluster",
				O11y: agentv2.O11yConfigV2{
					TenantID:      "tenant-1",
					ApplicationID: "app-1",
					BizType:       "tidbcloud",
				},
			},
		},
	}
}

// WithAnnotations sets the annotations on the agent.
func (b *TiDBGroupAgentBuilder) WithAnnotations(annotations map[string]string) *TiDBGroupAgentBuilder {
	if b.agent.Annotations == nil {
		b.agent.Annotations = make(map[string]string)
	}
	for k, v := range annotations {
		b.agent.Annotations[k] = v
	}
	return b
}

// WithName sets the name of the agent.
func (b *TiDBGroupAgentBuilder) WithName(name string) *TiDBGroupAgentBuilder {
	b.agent.Name = name
	return b
}

// WithNamespace sets the namespace of the agent.
func (b *TiDBGroupAgentBuilder) WithNamespace(ns string) *TiDBGroupAgentBuilder {
	b.agent.Namespace = ns
	return b
}

// WithTiDBGroup sets the TiDB group.
func (b *TiDBGroupAgentBuilder) WithTiDBGroup(group string) *TiDBGroupAgentBuilder {
	b.agent.Spec.TiDBGroup = group
	return b
}

// WithClusterID sets the cluster ID.
func (b *TiDBGroupAgentBuilder) WithClusterID(id string) *TiDBGroupAgentBuilder {
	b.agent.Spec.ClusterID = id
	return b
}

// WithSharedPoolID sets the shared pool ID.
func (b *TiDBGroupAgentBuilder) WithSharedPoolID(id string) *TiDBGroupAgentBuilder {
	b.agent.Spec.SharedPoolID = id
	return b
}

// Build returns the built TiDBGroupAgent.
func (b *TiDBGroupAgentBuilder) Build() *agentv2.TiDBGroupAgent {
	return b.agent
}

// O11yTiDBAgentBuilder builds O11yTiDBAgent instances for testing.
type O11yTiDBAgentBuilder struct {
	agent *agentv2.O11yTiDBAgent
}

// NewO11yTiDBAgentBuilder creates a new O11yTiDBAgentBuilder with default values.
func NewO11yTiDBAgentBuilder() *O11yTiDBAgentBuilder {
	return &O11yTiDBAgentBuilder{
		agent: &agentv2.O11yTiDBAgent{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "test-agent",
				Namespace: "default",
			},
			Spec: agentv2.O11yTiDBAgentSpec{
				KubernetesClusterName: "test-cluster",
				O11y: agentv2.O11yConfig{
					BizType:       "tidbcloud",
					TenantID:      "tenant-1",
					ProjectID:     "project-1",
					ApplicationID: "app-1",
				},
			},
		},
	}
}

// WithAnnotations sets the annotations on the agent.
func (b *O11yTiDBAgentBuilder) WithAnnotations(annotations map[string]string) *O11yTiDBAgentBuilder {
	if b.agent.Annotations == nil {
		b.agent.Annotations = make(map[string]string)
	}
	for k, v := range annotations {
		b.agent.Annotations[k] = v
	}
	return b
}

// WithName sets the name of the agent.
func (b *O11yTiDBAgentBuilder) WithName(name string) *O11yTiDBAgentBuilder {
	b.agent.Name = name
	return b
}

// WithNamespace sets the namespace of the agent.
func (b *O11yTiDBAgentBuilder) WithNamespace(ns string) *O11yTiDBAgentBuilder {
	b.agent.Namespace = ns
	return b
}

// WithClusterName sets the Kubernetes cluster name.
func (b *O11yTiDBAgentBuilder) WithClusterName(name string) *O11yTiDBAgentBuilder {
	b.agent.Spec.KubernetesClusterName = name
	return b
}

// WithPause sets the pause flag.
func (b *O11yTiDBAgentBuilder) WithPause(pause bool) *O11yTiDBAgentBuilder {
	b.agent.Spec.Pause = pause
	return b
}

// Build returns the built O11yTiDBAgent.
func (b *O11yTiDBAgentBuilder) Build() *agentv2.O11yTiDBAgent {
	return b.agent
}

// O11yK8sAgentBuilder builds O11yK8sAgent instances for testing.
type O11yK8sAgentBuilder struct {
	agent *agentv2.O11yK8sAgent
}

// NewO11yK8sAgentBuilder creates a new O11yK8sAgentBuilder with default values.
func NewO11yK8sAgentBuilder() *O11yK8sAgentBuilder {
	return &O11yK8sAgentBuilder{
		agent: &agentv2.O11yK8sAgent{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "test-agent",
				Namespace: "default",
			},
			Spec: agentv2.O11yK8sAgentSpec{
				KubernetesClusterName: "test-cluster",
				O11y: agentv2.O11yConfig{
					BizType:       "tidbcloud",
					TenantID:      "tenant-1",
					ProjectID:     "project-1",
					ApplicationID: "app-1",
				},
			},
		},
	}
}

// WithAnnotations sets the annotations on the agent.
func (b *O11yK8sAgentBuilder) WithAnnotations(annotations map[string]string) *O11yK8sAgentBuilder {
	if b.agent.Annotations == nil {
		b.agent.Annotations = make(map[string]string)
	}
	for k, v := range annotations {
		b.agent.Annotations[k] = v
	}
	return b
}

// WithName sets the name of the agent.
func (b *O11yK8sAgentBuilder) WithName(name string) *O11yK8sAgentBuilder {
	b.agent.Name = name
	return b
}

// WithNamespace sets the namespace of the agent.
func (b *O11yK8sAgentBuilder) WithNamespace(ns string) *O11yK8sAgentBuilder {
	b.agent.Namespace = ns
	return b
}

// WithClusterName sets the Kubernetes cluster name.
func (b *O11yK8sAgentBuilder) WithClusterName(name string) *O11yK8sAgentBuilder {
	b.agent.Spec.KubernetesClusterName = name
	return b
}

// WithPause sets the pause flag.
func (b *O11yK8sAgentBuilder) WithPause(pause bool) *O11yK8sAgentBuilder {
	b.agent.Spec.Pause = pause
	return b
}

// Build returns the built O11yK8sAgent.
func (b *O11yK8sAgentBuilder) Build() *agentv2.O11yK8sAgent {
	return b.agent
}