// Copyright (c) 2026, Pim van Pelt <pim@ipng.ch>

package vpp

import (
	"fmt"
	"net"
	"time"

	lb "git.ipng.ch/ipng/vpp-maglev/internal/vpp/binapi/lb"
	lb_types "git.ipng.ch/ipng/vpp-maglev/internal/vpp/binapi/lb_types"
)

// LBConf mirrors VPP's lb_conf_get_reply: global LB plugin settings.
type LBConf struct {
	IP4SrcAddress        net.IP
	IP6SrcAddress        net.IP
	StickyBucketsPerCore uint32
	FlowTimeout          uint32
}

// LBVIP mirrors VPP's lb_vip_details plus the set of application servers
// attached to this VIP (from lb_as_v2_details).
type LBVIP struct {
	Prefix          *net.IPNet // VIP address + prefix length
	Protocol        uint8      // IP proto (6=TCP, 17=UDP, 255=any)
	Port            uint16     // 0 = all-port VIP
	Encap           string     // gre4|gre6|l3dsr|nat4|nat6
	SrvType         string     // clusterip|nodeport
	Dscp            uint8
	TargetPort      uint16
	FlowTableLength uint16
	// SrcIPSticky is scraped from `show lb vips verbose` via cli_inband;
	// VPP's lb_vip_details does not carry this flag. Populated by
	// GetLBStateAll and GetLBStateVIP; see queryLBSticky in lbsync.go.
	SrcIPSticky bool
	ASes        []LBAS
}

// LBAS mirrors VPP's lb_as_v2_details: one application server bound to a VIP.
type LBAS struct {
	Address    net.IP
	Weight     uint8
	Flags      uint8 // bit 0 = used (alive), bit 1 = flushed
	NumBuckets uint32
	InUseSince time.Time // from VPP seconds-since-epoch (0 = never)
}

// LBState is a snapshot of the VPP LB plugin state.
type LBState struct {
	Conf LBConf
	VIPs []LBVIP
}

// GetLBStateAll fetches a full snapshot of the LB plugin state (global config
// plus every VIP and its application servers).
// Returns an error if VPP is not connected.
func (c *Client) GetLBStateAll() (*LBState, error) {
	ch, err := c.apiChannel()
	if err != nil {
		return nil, err
	}
	defer ch.Close()

	state := &LBState{}

	conf, err := getLBConf(ch)
	if err != nil {
		return nil, err
	}
	state.Conf = conf

	vips, err := dumpAllVIPs(ch)
	if err != nil {
		return nil, err
	}
	stickyMap, err := queryLBSticky(ch)
	if err != nil {
		return nil, err
	}
	for i := range vips {
		ases, err := dumpASesForVIP(ch, vips[i].Protocol, vips[i].Port)
		if err != nil {
			return nil, err
		}
		vips[i].ASes = ases
		vips[i].SrcIPSticky = stickyMap[makeVIPKey(vips[i].Prefix, vips[i].Protocol, vips[i].Port)]
	}
	state.VIPs = vips
	return state, nil
}

// GetLBStateVIP fetches a single VIP from VPP. Returns (nil, nil) if the VIP
// does not exist in VPP (caller must treat absence as "needs to be added").
// Returns an error only on transport/VPP failures.
func (c *Client) GetLBStateVIP(prefix *net.IPNet, protocol uint8, port uint16) (*LBVIP, error) {
	ch, err := c.apiChannel()
	if err != nil {
		return nil, err
	}
	defer ch.Close()
	vip, err := lookupVIP(ch, prefix, protocol, port)
	if err != nil || vip == nil {
		return vip, err
	}
	stickyMap, err := queryLBSticky(ch)
	if err != nil {
		return nil, err
	}
	vip.SrcIPSticky = stickyMap[makeVIPKey(vip.Prefix, vip.Protocol, vip.Port)]
	return vip, nil
}

// ---- low-level helpers (used by both Get and Sync paths) -------------------

func getLBConf(ch *loggedChannel) (LBConf, error) {
	reply := &lb.LbConfGetReply{}
	if err := ch.SendRequest(&lb.LbConfGet{}).ReceiveReply(reply); err != nil {
		return LBConf{}, fmt.Errorf("lb_conf_get: %w", err)
	}
	return LBConf{
		IP4SrcAddress:        ip4ToNetIP(reply.IP4SrcAddress),
		IP6SrcAddress:        ip6ToNetIP(reply.IP6SrcAddress),
		StickyBucketsPerCore: reply.StickyBucketsPerCore,
		FlowTimeout:          reply.FlowTimeout,
	}, nil
}

// dumpAllVIPs returns every VIP known to VPP (metadata only — ASes not populated).
func dumpAllVIPs(ch *loggedChannel) ([]LBVIP, error) {
	reqCtx := ch.SendMultiRequest(&lb.LbVipDump{})
	var out []LBVIP
	for {
		reply := &lb.LbVipDetails{}
		stop, err := reqCtx.ReceiveReply(reply)
		if err != nil {
			return nil, fmt.Errorf("lb_vip_dump: %w", err)
		}
		if stop {
			break
		}
		out = append(out, vipFromDetails(reply))
	}
	return out, nil
}

// lookupVIP finds a single VIP by (prefix, protocol, port) and returns it
// populated with its application servers, or nil if the VIP does not exist.
func lookupVIP(ch *loggedChannel, prefix *net.IPNet, protocol uint8, port uint16) (*LBVIP, error) {
	all, err := dumpAllVIPs(ch)
	if err != nil {
		return nil, err
	}
	want := prefix.String()
	for i := range all {
		if all[i].Prefix.String() != want {
			continue
		}
		if all[i].Protocol != protocol || all[i].Port != port {
			continue
		}
		ases, err := dumpASesForVIP(ch, protocol, port)
		if err != nil {
			return nil, err
		}
		all[i].ASes = ases
		return &all[i], nil
	}
	return nil, nil
}

// dumpASesForVIP returns the application servers bound to the VIP identified
// by (protocol, port). VPP's lb_as_v2_dump filter is used; we also guard
// defensively against replies for other VIPs.
func dumpASesForVIP(ch *loggedChannel, protocol uint8, port uint16) ([]LBAS, error) {
	req := &lb.LbAsV2Dump{
		Protocol: protocol,
		Port:     port,
	}
	reqCtx := ch.SendMultiRequest(req)
	var out []LBAS
	for {
		reply := &lb.LbAsV2Details{}
		stop, err := reqCtx.ReceiveReply(reply)
		if err != nil {
			return nil, fmt.Errorf("lb_as_v2_dump: %w", err)
		}
		if stop {
			break
		}
		if reply.Vip.Port != port || uint8(reply.Vip.Protocol) != protocol {
			continue
		}
		var inUse time.Time
		if reply.InUseSince != 0 {
			inUse = time.Unix(int64(reply.InUseSince), 0)
		}
		out = append(out, LBAS{
			Address:    reply.AppSrv.ToIP(),
			Weight:     reply.Weight,
			Flags:      reply.Flags,
			NumBuckets: reply.NumBuckets,
			InUseSince: inUse,
		})
	}
	return out, nil
}

// vipFromDetails builds an LBVIP (without ASes) from a VPP lb_vip_details reply.
func vipFromDetails(reply *lb.LbVipDetails) LBVIP {
	return LBVIP{
		Prefix:          lbVipPrefix(reply.Vip),
		Protocol:        uint8(reply.Vip.Protocol),
		Port:            reply.Vip.Port,
		Encap:           encapString(reply.Encap),
		SrvType:         srvTypeString(reply.SrvType),
		Dscp:            uint8(reply.Dscp),
		TargetPort:      reply.TargetPort,
		FlowTableLength: reply.FlowTableLength,
	}
}

// lbVipPrefix converts a VPP lb_vip's address+prefix to a *net.IPNet.
func lbVipPrefix(v lb_types.LbVip) *net.IPNet {
	ip := v.Pfx.Address.ToIP()
	bits := 32
	if ip.To4() == nil {
		bits = 128
	}
	return &net.IPNet{
		IP:   ip,
		Mask: net.CIDRMask(int(v.Pfx.Len), bits),
	}
}

func ip4ToNetIP(a [4]byte) net.IP {
	// VPP reports 255.255.255.255 when no IPv4 src is configured.
	if a == [4]byte{0xff, 0xff, 0xff, 0xff} {
		return nil
	}
	return net.IPv4(a[0], a[1], a[2], a[3]).To4()
}

func ip6ToNetIP(a [16]byte) net.IP {
	// VPP reports all-ones when no IPv6 src is configured.
	allOnes := true
	for _, b := range a {
		if b != 0xff {
			allOnes = false
			break
		}
	}
	if allOnes {
		return nil
	}
	ip := make(net.IP, 16)
	copy(ip, a[:])
	return ip
}

func encapString(e lb_types.LbEncapType) string {
	switch e {
	case lb_types.LB_API_ENCAP_TYPE_GRE4:
		return "gre4"
	case lb_types.LB_API_ENCAP_TYPE_GRE6:
		return "gre6"
	case lb_types.LB_API_ENCAP_TYPE_L3DSR:
		return "l3dsr"
	case lb_types.LB_API_ENCAP_TYPE_NAT4:
		return "nat4"
	case lb_types.LB_API_ENCAP_TYPE_NAT6:
		return "nat6"
	}
	return fmt.Sprintf("unknown(%d)", e)
}

func srvTypeString(t lb_types.LbSrvType) string {
	switch t {
	case lb_types.LB_API_SRV_TYPE_CLUSTERIP:
		return "clusterip"
	case lb_types.LB_API_SRV_TYPE_NODEPORT:
		return "nodeport"
	}
	return fmt.Sprintf("unknown(%d)", t)
}