package throttler

import (
	"fmt"
	"math"
	"net"
	"reflect"
	"strconv"
	"strings"
	"syscall"

	"github.com/dustin/go-humanize"
	"github.com/vishvananda/netlink"
	"golang.org/x/sys/unix"
	corev1 "k8s.io/api/core/v1"

	"github.com/tidbcloud/shared-provider/pkg/types"
	"github.com/tidbcloud/shared-provider/pkg/utils/closer"
	"github.com/tidbcloud/shared-provider/pkg/utils/netns"
)

const (
	// number of bits to shift when convert between byte and bit
	shiftBits    = 3
	maxBandwidth = math.MaxUint64 >> shiftBits
)

type ThrottleConfig struct {
	Ingress *Rule
	Egress  *Rule
}

type Rule struct {
	Bandwidth uint64
	Port      uint16
}

func parseRule(anno map[string]string, bandwidthKey, portKey string) (*Rule, error) {
	strBandwidth, ok := anno[bandwidthKey]
	if !ok {
		return nil, nil
	}
	// TODO: extract a common utils to convert uint64 and string
	bandwidth, err := humanize.ParseBytes(strings.TrimSuffix(strBandwidth, "/s"))
	if err != nil {
		return nil, err
	}
	if bandwidth >= maxBandwidth {
		return nil, fmt.Errorf("too big bandwidth limit, the maximum supported now is %v", humanize.Bytes(maxBandwidth)+"/s")
	}
	r := Rule{
		Bandwidth: bandwidth,
	}
	// TODO: now port must be specified
	strPort, ok := anno[portKey]
	if !ok {
		return nil, fmt.Errorf("%s in annoataions is needed", portKey)
	}

	port, err := strconv.ParseUint(strPort, 10, 16)
	if err != nil {
		return nil, err
	}
	r.Port = uint16(port)

	return &r, nil
}

func NewThrottleConfig(anno map[string]string) (*ThrottleConfig, error) {
	cfg := ThrottleConfig{}
	if anno == nil {
		return &cfg, nil
	}

	ingress, err := parseRule(anno, types.AnnoKeyIngress, types.AnnoKeyIngressPort)
	if err != nil {
		return nil, err
	}

	cfg.Ingress = ingress

	egress, err := parseRule(anno, types.AnnoKeyEgress, types.AnnoKeyEgressPort)
	if err != nil {
		return nil, err
	}

	cfg.Egress = egress

	return &cfg, nil
}

type rules struct {
	ingress *ingressRule
	egress  *egressRule
}

type ingressRule struct {
	qdisc  *netlink.Ingress
	filter *netlink.U32

	common *commonRule
}

type egressRule struct {
	common *commonRule
}

type commonRule struct {
	qdisc  *netlink.Htb
	class  *netlink.HtbClass
	filter *netlink.U32
}

type Throttler interface {
	Throttle(p *corev1.Pod, cfg *ThrottleConfig) error
}

func NewThrottler(r netns.Runtime, ifName string) Throttler {
	return &throttler{
		r:           r,
		ifName:      ifName,
		ifbName:     "ifb0",
		qdiscHandle: netlink.MakeHandle(1, 0),
		classHandle: netlink.MakeHandle(1, 1),
		// nolint: gomnd // this is initialization
		ingressHandle: netlink.MakeHandle(0xffff, 0),
	}
}

type throttler struct {
	r netns.Runtime

	ifName  string
	ifbName string

	qdiscHandle   uint32
	classHandle   uint32
	ingressHandle uint32
}

func (t *throttler) Throttle(p *corev1.Pod, cfg *ThrottleConfig) error {
	ns, err := t.r.GetSandboxNetNS(string(p.UID))
	if err != nil {
		return fmt.Errorf("can't get netns of pod %s/%s, uid %s: %w", p.Namespace, p.Name, p.UID, err)
	}
	defer closer.Close(ns)

	if err := ns.Do(func() error {
		link, err := netlink.LinkByName(t.ifName)
		if err != nil {
			return fmt.Errorf("can't get link named %s: %w", t.ifName, err)
		}
		ifb, err := getOrCreateIfb(t.ifbName, link.Attrs().MTU)
		if err != nil {
			return fmt.Errorf("can't get or create ifb %s: %w", t.ifbName, err)
		}
		cur, err := t.findCurrentRules(link, ifb)
		if err != nil {
			return err
		}
		expected := t.getExpectedRules(link, ifb, cfg)
		if err := t.createOrUpdate(cur, expected); err != nil {
			return err
		}

		return nil
	}); err != nil {
		return fmt.Errorf("throttle for %s/%s failed, uid %s: %w", p.Namespace, p.Name, p.UID, err)
	}

	return nil
}

func getOrCreateIfb(name string, mtu int) (netlink.Link, error) {
	ifbLink, err := netlink.LinkByName(name)
	if err != nil {
		if _, ok := err.(netlink.LinkNotFoundError); !ok {
			return nil, fmt.Errorf("can't get ifb link: %w", err)
		}
		if err := netlink.LinkAdd(newIfb(name, mtu)); err != nil {
			return nil, err
		}
		ifb, err := netlink.LinkByName(name)
		if err != nil {
			return nil, err
		}
		return ifb, nil
	}

	return ifbLink, nil
}

func (t *throttler) getExpectedRules(link netlink.Link, ifb netlink.Link, cfg *ThrottleConfig) *rules {
	r := rules{}

	if cfg.Ingress != nil {
		ingress := newIngress(link, t.ingressHandle)
		ingressFilter := newIngressFilter(link, ifb, t.ingressHandle)

		sel := WithDstPort(nil, cfg.Ingress.Port)
		common := t.newCommonRule(ifb, cfg.Ingress.Bandwidth, sel)

		r.ingress = &ingressRule{
			qdisc:  ingress,
			filter: ingressFilter,

			common: common,
		}
	}

	if cfg.Egress != nil {
		sel := WithSrcPort(nil, cfg.Egress.Port)
		common := t.newCommonRule(link, cfg.Egress.Bandwidth, sel)
		r.egress = &egressRule{
			common: common,
		}
	}

	return &r
}

func (t *throttler) newCommonRule(link netlink.Link, bandwidth uint64, sel *netlink.TcU32Sel) *commonRule {
	qdisc := newCommonQdisc(link, t.qdiscHandle)
	class := newCommonClass(link, t.qdiscHandle, t.classHandle, bandwidth)
	filter := newCommonFilter(link, t.qdiscHandle, t.classHandle, sel)

	return &commonRule{
		qdisc:  qdisc,
		class:  class,
		filter: filter,
	}
}

func newCommonQdisc(link netlink.Link, handle uint32) *netlink.Htb {
	htb := netlink.NewHtb(netlink.QdiscAttrs{
		LinkIndex: link.Attrs().Index,
		Handle:    handle,
		Parent:    netlink.HANDLE_ROOT,
	})
	htb.Defcls = 30

	return htb
}

func newCommonClass(link netlink.Link, parent uint32, handle uint32, rateBytes uint64) *netlink.HtbClass {
	class := netlink.NewHtbClass(netlink.ClassAttrs{
		LinkIndex: link.Attrs().Index,
		Parent:    parent,
		Handle:    handle,
	}, netlink.HtbClassAttrs{
		Rate: rateBytes << shiftBits,
	})

	return class
}

func newCommonFilter(link netlink.Link, parent uint32, flowid uint32, sel *netlink.TcU32Sel) *netlink.U32 {
	filter := &netlink.U32{
		FilterAttrs: netlink.FilterAttrs{
			LinkIndex: link.Attrs().Index,
			Parent:    parent,
			Protocol:  unix.ETH_P_IP,
			Priority:  1,
		},
		ClassId: flowid,
		Sel:     sel,
	}
	return filter
}

func newIfb(name string, mtu int) *netlink.Ifb {
	attrs := netlink.NewLinkAttrs()
	attrs.Name = name
	attrs.Flags = net.FlagUp
	attrs.MTU = mtu
	return &netlink.Ifb{
		LinkAttrs: attrs,
	}
}

func newIngress(link netlink.Link, handle uint32) *netlink.Ingress {
	return &netlink.Ingress{
		QdiscAttrs: netlink.QdiscAttrs{
			LinkIndex: link.Attrs().Index,
			Handle:    handle,
			Parent:    netlink.HANDLE_INGRESS,
		},
	}
}

func newIngressFilter(link netlink.Link, ifb netlink.Link, parent uint32) *netlink.U32 {
	return &netlink.U32{
		FilterAttrs: netlink.FilterAttrs{
			LinkIndex: link.Attrs().Index,
			Parent:    parent,
			Priority:  1,
			Protocol:  syscall.ETH_P_ALL,
		},
		RedirIndex: ifb.Attrs().Index,
		Actions: []netlink.Action{
			&netlink.MirredAction{
				ActionAttrs:  netlink.ActionAttrs{},
				MirredAction: netlink.TCA_EGRESS_REDIR,
				Ifindex:      ifb.Attrs().Index,
			},
		},
	}
}

func (t *throttler) findCurrentRules(link, ifb netlink.Link) (*rules, error) {
	r := rules{}

	egressCommon, err := t.findCommonRule(link)
	if err != nil {
		return nil, err
	}

	r.egress = &egressRule{
		common: egressCommon,
	}

	ingress, err := t.findIngress(link)
	if err != nil {
		return nil, err
	}
	ingressFilter, err := t.findIngressFilter(link)
	if err != nil {
		return nil, err
	}
	r.ingress = &ingressRule{
		qdisc:  ingress,
		filter: ingressFilter,
	}

	ingresCommon, err := t.findCommonRule(ifb)
	if err != nil {
		return nil, err
	}

	r.ingress.common = ingresCommon

	return &r, nil
}

func (t *throttler) findCommonRule(link netlink.Link) (*commonRule, error) {
	htb, err := t.findQdisc(link)
	if err != nil {
		return nil, err
	}
	class, err := t.findClass(link)
	if err != nil {
		return nil, err
	}
	filter, err := t.findFilter(link)
	if err != nil {
		return nil, err
	}

	return &commonRule{
		qdisc:  htb,
		class:  class,
		filter: filter,
	}, nil
}

func (t *throttler) findQdisc(link netlink.Link) (*netlink.Htb, error) {
	qdiscs, err := netlink.QdiscList(link)
	if err != nil {
		return nil, err
	}
	for _, qdisc := range qdiscs {
		if qdisc.Attrs().Parent != netlink.HANDLE_ROOT {
			continue
		}
		if qdisc.Attrs().Handle != t.qdiscHandle {
			continue
		}
		htb, ok := qdisc.(*netlink.Htb)
		if !ok {
			return nil, nil
		}
		return htb, nil
	}

	return nil, nil
}

func (t *throttler) findClass(link netlink.Link) (*netlink.HtbClass, error) {
	classes, err := netlink.ClassList(link, t.qdiscHandle)
	if err != nil {
		return nil, err
	}
	for _, cls := range classes {
		if cls.Attrs().Handle == t.classHandle {
			htbCls, ok := cls.(*netlink.HtbClass)
			if ok {
				return htbCls, nil
			}
		}
	}

	return nil, nil
}

func (t *throttler) findFilter(link netlink.Link) (*netlink.U32, error) {
	filters, err := netlink.FilterList(link, t.qdiscHandle)
	if err != nil {
		return nil, err
	}
	if len(filters) == 0 {
		return nil, nil
	}

	if len(filters) != 1 {
		return nil, nil
	}

	filter := filters[0]

	f, ok := filter.(*netlink.U32)
	if !ok {
		return nil, nil
	}
	if f.ClassId != t.classHandle {
		return nil, nil
	}

	return f, nil
}

func (t *throttler) findIngress(link netlink.Link) (*netlink.Ingress, error) {
	qdiscs, err := netlink.QdiscList(link)
	if err != nil {
		return nil, err
	}
	for _, qdisc := range qdiscs {
		if qdisc.Attrs().Parent != netlink.HANDLE_INGRESS {
			continue
		}
		if qdisc.Attrs().Handle != t.ingressHandle {
			continue
		}
		ingress, ok := qdisc.(*netlink.Ingress)
		if !ok {
			return nil, nil
		}
		return ingress, nil
	}

	return nil, nil
}

func (t *throttler) findIngressFilter(link netlink.Link) (*netlink.U32, error) {
	filters, err := netlink.FilterList(link, t.ingressHandle)
	if err != nil {
		return nil, err
	}
	if len(filters) == 0 {
		return nil, nil
	}
	if len(filters) != 1 {
		return nil, nil
	}

	filter := filters[0]

	f, ok := filter.(*netlink.U32)
	if !ok {
		return nil, nil
	}

	return f, nil
}

// ingress
// ip link set dev ifb0 up
// tc qdisc  add dev eth0 handle ffff: ingress
// tc qdisc  add dev eth0 handle ffff: ingress
// tc filter add dev eth0 parent ffff: protocol ip u32 match u32 0 0 action mirred egress redirect dev ifb0
// tc qdisc  add dev ifb0 root handle 1: htb default 30
// tc class  add dev ifb0 parent 1: classid 1:1 htb rate $rateBits
// tc filter add dev ifb0 protocol ip parent 1: prio 1 u32 match ip dport $port 0xffff flowid 1:1
//
// egress
// tc qdisc  add dev eth0 root handle 1: htb default 30
// tc class  add dev eth0 parent 1: classid 1:1 htb rate $rateBits
// tc filter add dev eth0 protocol ip parent 1: prio 1 u32 match ip sport $port 0xffff flowid 1:1
func (t *throttler) createOrUpdate(cur, expected *rules) error {
	if compare(cur, expected) {
		return nil
	}

	if cur.ingress != nil {
		if cur.ingress.qdisc != nil {
			if err := netlink.QdiscDel(cur.ingress.qdisc); err != nil {
				return fmt.Errorf("can't del ingress qdisc: %w", err)
			}
		}
		if cur.ingress.common.qdisc != nil {
			if err := netlink.QdiscDel(cur.ingress.common.qdisc); err != nil {
				return fmt.Errorf("can't del ingress common qidsc: %w", err)
			}
		}
	}
	if cur.egress != nil {
		if cur.egress.common.qdisc != nil {
			if err := netlink.QdiscDel(cur.egress.common.qdisc); err != nil {
				return fmt.Errorf("can't del egress common qidsc: %w", err)
			}
		}
	}
	if err := createIngressRule(expected.ingress); err != nil {
		return fmt.Errorf("can't create ingress rules: %w", err)
	}

	if err := createEgressRule(expected.egress); err != nil {
		return fmt.Errorf("can't create egress rules: %w", err)
	}

	return nil
}

func createIngressRule(ingress *ingressRule) error {
	if ingress == nil {
		return nil
	}
	if ingress.qdisc != nil {
		if err := netlink.QdiscAdd(ingress.qdisc); err != nil {
			return err
		}
	}
	if ingress.filter != nil {
		if err := netlink.FilterAdd(ingress.filter); err != nil {
			return err
		}
	}

	return createCommonRule(ingress.common)
}

func createEgressRule(egress *egressRule) error {
	if egress == nil {
		return nil
	}
	return createCommonRule(egress.common)
}

func createCommonRule(common *commonRule) error {
	if common.qdisc != nil {
		if err := netlink.QdiscAdd(common.qdisc); err != nil {
			return err
		}
	}

	if common.class != nil {
		if err := netlink.ClassAdd(common.class); err != nil {
			return err
		}
	}
	if common.filter != nil {
		if err := netlink.FilterAdd(common.filter); err != nil {
			return err
		}
	}

	return nil
}

func compare(a, b *rules) bool {
	if a == nil && b == nil {
		return true
	}

	if a == nil || b == nil {
		return false
	}

	if !compareIngressRule(a.ingress, b.ingress) {
		return false
	}

	if !compareEgressRule(a.egress, b.egress) {
		return false
	}

	return true
}

func compareIngressRule(a, b *ingressRule) bool {
	if a == nil && b == nil {
		return true
	}

	if a == nil || b == nil {
		return false
	}

	if !compareIngress(a.qdisc, b.qdisc) {
		return false
	}

	if !compareIngressFilter(a.filter, b.filter) {
		return false
	}

	if !compareCommonRule(a.common, b.common) {
		return false
	}

	return true
}

func compareEgressRule(a, b *egressRule) bool {
	if a == nil && b == nil {
		return true
	}

	if a == nil || b == nil {
		return false
	}

	if !compareCommonRule(a.common, b.common) {
		return false
	}

	return true
}

func compareIngress(a, b *netlink.Ingress) bool {
	if a == nil && b == nil {
		return true
	}

	if a == nil || b == nil {
		return false
	}

	return true
}

func compareIngressFilter(a, b *netlink.U32) bool {
	if a == nil && b == nil {
		return true
	}

	if a == nil || b == nil {
		return false
	}

	// TODO: real compare two filter

	return true
}

func compareCommonRule(a, b *commonRule) bool {
	if a == nil && b == nil {
		return true
	}

	if a == nil || b == nil {
		return false
	}

	if !compareQdisc(a.qdisc, b.qdisc) {
		return false
	}

	if !compareClass(a.class, b.class) {
		return false
	}

	if !compareFilter(a.filter, b.filter) {
		return false
	}

	return true
}

func compareQdisc(a, b *netlink.Htb) bool {
	if a == nil && b == nil {
		return true
	}

	if a == nil || b == nil {
		return false
	}

	if a.Defcls != b.Defcls {
		return false
	}

	return true
}

func compareClass(a, b *netlink.HtbClass) bool {
	if a == nil && b == nil {
		return true
	}

	if a == nil || b == nil {
		return false
	}

	if a.Rate != b.Rate {
		return false
	}

	return true
}

func compareFilter(a, b *netlink.U32) bool {
	if a == nil && b == nil {
		return true
	}

	if a == nil || b == nil {
		return false
	}

	if !reflect.DeepEqual(a.Sel, b.Sel) {
		return false
	}

	return true
}