package networkpolicy

import (
	corev1 "k8s.io/api/core/v1"
	netv1 "k8s.io/api/networking/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"

	"github.com/tidbcloud/infra-provider/network-policy-controller/pkg/types"
)

const (
	LabelKeyNamespace = "kubernetes.io/metadata.name"
)

const (
	PolicyDefaultDeny         = "default-deny"
	PolicyAllowAllInNamespace = "allow-all-in-namespace"

	PolicyAllowToDNS           = "allow-to-dns"
	PolicyAllowFromDBaaSProxy  = "allow-from-dbaas-proxy"
	PolicyAllowFromSliReporter = "allow-from-sli-reporter"

	PolicyForO11yVectorAgent = "for-o11y-vector-agent"
	PolicyForO11yVMAgent     = "for-o11y-vm-agent"
	PolicyForTiDBDiscovery   = "for-tidb-discovery"
	PolicyForPD              = "for-pd"
	PolicyForTiDB            = "for-tidb"
	PolicyForTiKV            = "for-tikv"
	PolicyForTiFlash         = "for-tiflash"
	PolicyForBR              = "for-br"
	PolicyForTiDBLightning   = "for-tidb-lightning"
	PolicyForTiDBMonitor     = "for-tidb-monitor"
)

var AllPolicyNames = []string{
	PolicyDefaultDeny,
	PolicyAllowAllInNamespace,

	PolicyAllowToDNS,
	PolicyAllowFromDBaaSProxy,
	PolicyAllowFromSliReporter,

	PolicyForO11yVectorAgent,
	PolicyForO11yVMAgent,
	PolicyForTiDBDiscovery,
	PolicyForPD,
	PolicyForTiDB,
	PolicyForTiKV,
	PolicyForTiFlash,
	PolicyForBR,
	PolicyForTiDBLightning,
	PolicyForTiDBMonitor,
}

var AllowToS3PolicyNames = []string{
	PolicyForTiKV,
	PolicyForTiFlash,
	PolicyForBR,
	PolicyForTiDBLightning,
}

func (r *Reconciler) NewPolicyDefaultDeny(ns string) *netv1.NetworkPolicy {
	return &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyDefaultDeny,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeIngress,
				netv1.PolicyTypeEgress,
			},
			PodSelector: metav1.LabelSelector{},
		},
	}
}

func (r *Reconciler) NewPolicyAllowAllInNamespce(ns string) *netv1.NetworkPolicy {
	return &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyAllowAllInNamespace,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeEgress,
				netv1.PolicyTypeIngress,
			},
			PodSelector: metav1.LabelSelector{},
			Egress: []netv1.NetworkPolicyEgressRule{
				{
					To: []netv1.NetworkPolicyPeer{
						{
							PodSelector: &metav1.LabelSelector{},
						},
					},
				},
			},
			Ingress: []netv1.NetworkPolicyIngressRule{
				{
					From: []netv1.NetworkPolicyPeer{
						{
							PodSelector: &metav1.LabelSelector{},
						},
					},
				},
			},
		},
	}
}

func (r *Reconciler) NewPolicyAllowToDNS(ns string) *netv1.NetworkPolicy {
	tcpProtocol := corev1.ProtocolTCP
	udpProtocol := corev1.ProtocolUDP
	dnsPort := intstr.FromInt(53)
	return &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyAllowToDNS,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeEgress,
			},
			PodSelector: metav1.LabelSelector{},
			Egress: []netv1.NetworkPolicyEgressRule{
				{
					To: []netv1.NetworkPolicyPeer{
						{
							NamespaceSelector: &metav1.LabelSelector{
								MatchLabels: map[string]string{
									LabelKeyNamespace: r.DNSNamespace,
								},
							},
							PodSelector: r.DNSPodSelector,
						},
					},
					Ports: []netv1.NetworkPolicyPort{
						{
							Protocol: &tcpProtocol,
							Port:     &dnsPort,
						},
						{
							Protocol: &udpProtocol,
							Port:     &dnsPort,
						},
					},
				},
			},
		},
	}
}

// nolint: dupl // nolint
func (r *Reconciler) NewPolicyAllowFromDBaaSProxy(ns string) *netv1.NetworkPolicy {
	return &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyAllowFromDBaaSProxy,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeIngress,
			},
			// monitor & shell pod create by proxy should be enough? but allow all now.
			PodSelector: metav1.LabelSelector{},
			Ingress: []netv1.NetworkPolicyIngressRule{
				{
					From: []netv1.NetworkPolicyPeer{
						{
							NamespaceSelector: &metav1.LabelSelector{
								MatchLabels: map[string]string{
									LabelKeyNamespace: "tidb-admin",
								},
							},
							PodSelector: &metav1.LabelSelector{
								MatchLabels: map[string]string{
									"app": "dbaas-proxy",
								},
							},
						},
					},
				},
			},
		},
	}
}

// nolint: dupl // nolint
func (r *Reconciler) NewPolicyAllowFromSliReporter(ns string) *netv1.NetworkPolicy {
	return &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyAllowFromSliReporter,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeIngress,
			},
			PodSelector: metav1.LabelSelector{},
			Ingress: []netv1.NetworkPolicyIngressRule{
				{
					From: []netv1.NetworkPolicyPeer{
						{
							NamespaceSelector: &metav1.LabelSelector{
								MatchLabels: map[string]string{
									LabelKeyNamespace: "tidb-admin",
								},
							},
							PodSelector: &metav1.LabelSelector{
								MatchLabels: map[string]string{
									"app": "sli-reporter",
								},
							},
						},
					},
				},
			},
		},
	}
}

func (r *Reconciler) NewPolicyForO11yVectorAgent(ns string) *netv1.NetworkPolicy {
	p := &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyForO11yVectorAgent,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeEgress,
			},

			PodSelector: metav1.LabelSelector{
				MatchLabels: map[string]string{
					"app": "o11y-vector-agent",
				},
			},
		},
	}

	p.Spec.Egress = append(p.Spec.Egress, r.NewEgressToO11y()...)
	return p
}

func (r *Reconciler) NewPolicyForO11yVMAgent(ns string) *netv1.NetworkPolicy {
	p := &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyForO11yVMAgent,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeEgress,
			},

			PodSelector: metav1.LabelSelector{
				MatchLabels: map[string]string{
					"app.kubernetes.io/name": "victoria-metrics-agent",
				},
			},
		},
	}

	p.Spec.Egress = append(p.Spec.Egress, r.NewEgressToO11y()...)
	p.Spec.Egress = append(p.Spec.Egress, r.NewEgressToControlPlane()...)
	return p
}

func (r *Reconciler) NewPolicyForTiDBDiscovery(ns string) *netv1.NetworkPolicy {
	p := &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyForTiDBDiscovery,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeEgress,
			},

			PodSelector: metav1.LabelSelector{
				MatchLabels: map[string]string{
					types.LabelKeyComponent: types.LabelValComponentDiscovery,
				},
			},
		},
	}

	p.Spec.Egress = append(p.Spec.Egress, r.NewEgressToControlPlane()...)
	return p
}

func (r *Reconciler) NewPolicyForPD(ns string) *netv1.NetworkPolicy {
	p := &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyForPD,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeIngress,
			},

			PodSelector: metav1.LabelSelector{
				MatchLabels: map[string]string{
					types.LabelKeyComponent: types.LabelValComponentPD,
				},
			},
		},
	}

	p.Spec.Ingress = append(p.Spec.Ingress, r.NewIngressFromTiDBOperator()...)
	p.Spec.Ingress = append(p.Spec.Ingress, r.NewIngressFromTiflashReadNode()...)
	return p
}

func (r *Reconciler) NewPolicyForTiDB(ns string) *netv1.NetworkPolicy {
	p := &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyForTiDB,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeIngress,
				netv1.PolicyTypeEgress,
			},

			PodSelector: metav1.LabelSelector{
				MatchLabels: map[string]string{
					types.LabelKeyComponent: types.LabelValComponentTiDB,
				},
			},
		},
	}

	p.Spec.Ingress = append(p.Spec.Ingress, r.NewIngressFromLoadBalancer()...)
	p.Spec.Egress = append(p.Spec.Egress, r.NewEgressToTiflashReadNode()...)
	return p
}

func (r *Reconciler) NewPolicyForTiKV(ns string) *netv1.NetworkPolicy {
	p := &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyForTiKV,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeEgress,
			},

			PodSelector: metav1.LabelSelector{
				MatchLabels: map[string]string{
					types.LabelKeyComponent: types.LabelValComponentTiKV,
				},
			},
		},
	}

	p.Spec.Egress = append(p.Spec.Egress, r.NewEgressToS3()...)
	return p
}

func (r *Reconciler) NewPolicyForTiFlash(ns string) *netv1.NetworkPolicy {
	p := &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyForTiFlash,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeIngress,
				netv1.PolicyTypeEgress,
			},

			PodSelector: metav1.LabelSelector{
				MatchLabels: map[string]string{
					types.LabelKeyComponent: types.LabelValComponentTiFlash,
				},
			},
		},
	}

	p.Spec.Ingress = append(p.Spec.Ingress, r.NewIngressFromTiflashReadNode()...)
	p.Spec.Egress = append(p.Spec.Egress, r.NewEgressToS3()...)
	return p
}

func (r *Reconciler) NewPolicyForBR(ns string) *netv1.NetworkPolicy {
	p := &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyForBR,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeEgress,
			},

			PodSelector: metav1.LabelSelector{
				MatchExpressions: []metav1.LabelSelectorRequirement{
					{
						Key:      types.LabelKeyComponent,
						Operator: metav1.LabelSelectorOpIn,
						Values: []string{
							types.LabelValComponentBackup,
							types.LabelValComponentRestore,
							types.LabelValComponentClean,
						},
					},
				},
			},
		},
	}

	p.Spec.Egress = append(p.Spec.Egress, r.NewEgressToS3()...)
	p.Spec.Egress = append(p.Spec.Egress, r.NewEgressToControlPlane()...)
	return p
}

func (r *Reconciler) NewPolicyForTiDBLightning(ns string) *netv1.NetworkPolicy {
	p := &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyForTiDBLightning,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeEgress,
			},

			PodSelector: metav1.LabelSelector{
				MatchLabels: map[string]string{
					types.LabelKeyComponent: types.LabelValComponentTiDBLightning,
				},
			},
		},
	}

	p.Spec.Egress = append(p.Spec.Egress, r.NewEgressToS3()...)
	return p
}

func (r *Reconciler) NewPolicyForTiDBMonitor(ns string) *netv1.NetworkPolicy {
	p := &netv1.NetworkPolicy{
		ObjectMeta: metav1.ObjectMeta{
			Name:      PolicyForTiDBMonitor,
			Namespace: ns,
		},
		Spec: netv1.NetworkPolicySpec{
			PolicyTypes: []netv1.PolicyType{
				netv1.PolicyTypeEgress,
			},

			PodSelector: metav1.LabelSelector{
				MatchLabels: map[string]string{
					types.LabelKeyComponent: types.LabelValComponentMonitor,
				},
			},
		},
	}

	p.Spec.Egress = append(p.Spec.Egress, r.NewEgressToControlPlane()...)
	return p
}

func (r *Reconciler) NewEgressToO11y() []netv1.NetworkPolicyEgressRule {
	httpPort := intstr.FromInt(80)

	var peers []netv1.NetworkPolicyPeer
	for _, cidr := range r.O11yCIDRs {
		peers = append(peers, netv1.NetworkPolicyPeer{
			IPBlock: &netv1.IPBlock{
				CIDR: cidr,
			},
		})
	}

	return []netv1.NetworkPolicyEgressRule{
		{
			To: peers,
			Ports: []netv1.NetworkPolicyPort{
				{
					Port: &httpPort,
				},
			},
		},
	}
}

func (r *Reconciler) NewEgressToControlPlane() []netv1.NetworkPolicyEgressRule {
	tcpProtocol := corev1.ProtocolTCP
	httpsPort := intstr.FromInt(443)

	var peers []netv1.NetworkPolicyPeer
	for _, cidr := range r.ControlPlaneCIDRs {
		peers = append(peers, netv1.NetworkPolicyPeer{
			IPBlock: &netv1.IPBlock{
				CIDR: cidr,
			},
		})
	}

	return []netv1.NetworkPolicyEgressRule{
		{
			To: peers,
			Ports: []netv1.NetworkPolicyPort{
				{
					Protocol: &tcpProtocol,
					Port:     &httpsPort,
				},
			},
		},
	}
}

func (r *Reconciler) NewEgressToS3() []netv1.NetworkPolicyEgressRule {
	httpPort := intstr.FromInt(80)
	httpsPort := intstr.FromInt(443)
	var peers []netv1.NetworkPolicyPeer
	for _, cidr := range r.S3CIDRPoller.CIDRs() {
		peers = append(peers, netv1.NetworkPolicyPeer{
			IPBlock: &netv1.IPBlock{
				CIDR: cidr,
			},
		})
	}
	return []netv1.NetworkPolicyEgressRule{
		{
			To: peers,
			Ports: []netv1.NetworkPolicyPort{
				{
					Port: &httpPort,
				},
				{
					Port: &httpsPort,
				},
			},
		},
	}
}

func (r *Reconciler) NewIngressFromTiDBOperator() []netv1.NetworkPolicyIngressRule {
	return []netv1.NetworkPolicyIngressRule{
		{
			From: []netv1.NetworkPolicyPeer{
				{
					NamespaceSelector: &metav1.LabelSelector{
						MatchLabels: map[string]string{
							LabelKeyNamespace: "tidb-admin",
						},
					},
					PodSelector: &metav1.LabelSelector{
						MatchLabels: map[string]string{
							types.LabelKeyComponent: types.LabelValComponentControllerManager,
						},
					},
				},
			},
		},
	}
}

func (r *Reconciler) NewIngressFromLoadBalancer() []netv1.NetworkPolicyIngressRule {
	tcpProtocol := corev1.ProtocolTCP
	tidbPort := intstr.FromInt(4000)

	peers := []netv1.NetworkPolicyPeer{}
	for _, cidr := range r.LoadBalancerCIDRs {
		peers = append(peers, netv1.NetworkPolicyPeer{
			IPBlock: &netv1.IPBlock{CIDR: cidr},
		})
	}
	return []netv1.NetworkPolicyIngressRule{
		{
			From: peers,
			Ports: []netv1.NetworkPolicyPort{
				{
					Port:     &tidbPort,
					Protocol: &tcpProtocol,
				},
			},
		},
	}
}

func (r *Reconciler) NewEgressToTiflashReadNode() []netv1.NetworkPolicyEgressRule {
	return []netv1.NetworkPolicyEgressRule{
		{
			To: []netv1.NetworkPolicyPeer{
				{
					NamespaceSelector: &metav1.LabelSelector{
						MatchLabels: map[string]string{
							LabelKeyNamespace: "tidb-admin",
						},
					},
					PodSelector: &metav1.LabelSelector{
						MatchLabels: map[string]string{
							"app": "readnode",
						},
					},
				},
			},
		},
	}
}

func (r *Reconciler) NewIngressFromTiflashReadNode() []netv1.NetworkPolicyIngressRule {
	return []netv1.NetworkPolicyIngressRule{
		{
			From: []netv1.NetworkPolicyPeer{
				{
					NamespaceSelector: &metav1.LabelSelector{
						MatchLabels: map[string]string{
							LabelKeyNamespace: "tidb-admin",
						},
					},
					PodSelector: &metav1.LabelSelector{
						MatchLabels: map[string]string{
							"app": "readnode",
						},
					},
				},
			},
		},
	}
}