package vpn

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/netip"

	corev1 "k8s.io/api/core/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"edge-infra.dev/pkg/sds/remoteaccess/constants"
	v1vpnconfig "edge-infra.dev/pkg/sds/remoteaccess/k8s/apis/vpnconfigs/v1"
)

var ErrNoIPAddressesAvailable = errors.New("no IP addresses available in subnet")
var ErrSubnetNotConfigured = errors.New("subnet has not been configured yet, run `ConfigureSubnet()`")

// map of IP addresses to ClusterEdgeID
type IPAddressPool map[string]string

func (v *VPN) ConfigureSubnet(ctx context.Context, c client.Client) error {
	subnetCIDR, err := getSubnetCIDRFromConfigMap(ctx, c)
	if err != nil {
		return err
	}

	// do nothing if address pool is populated and subnet has not changed
	if v.AvailableIPAddressPool != nil {
		if subnetCIDR == v.GetSubnetCIDR() {
			return nil
		}
	}

	// update subnet CIDR and repopulate IP address pool from current VPNConfig IP addresses
	v.SubnetCIDR = subnetCIDR
	return v.populateAvailableIPAddressPool(ctx, c)
}

func getSubnetCIDRFromConfigMap(ctx context.Context, c client.Client) (string, error) {
	configMap := &corev1.ConfigMap{}
	if err := c.Get(ctx, vpnConfigMapKey, configMap); err != nil {
		return "", err
	}

	subnetCIDR, ok := configMap.Data[constants.VPNConfigMapCIDRKey]
	if !ok {
		return "", fmt.Errorf("expected key '%s' in ConfigMap %s/%s", constants.VPNConfigMapCIDRKey, constants.VPNNamespace, vpnConfigMapKey.String())
	}
	return subnetCIDR, nil
}

func (v *VPN) populateAvailableIPAddressPool(ctx context.Context, c client.Client) error {
	if err := v.createInitialIPAddressPool(ctx, c); err != nil {
		return err
	}

	vpnConfigs, err := getVPNConfigList(ctx, c)
	if err != nil {
		return err
	}

	// set (in-subnet) IP addresses from all VPNConfigs as unavailable
	for _, vpnConfig := range vpnConfigs.Items {
		isInSubnet, err := v.IPAddressIsInSubnet(vpnConfig.IP())
		if err != nil {
			return err
		}
		if isInSubnet && v.CheckIPAddressIsAvailable(vpnConfig.IP()) {
			v.setIPAddressUnavailable(vpnConfig.IP(), vpnConfig.ClusterEdgeID())
		}
	}

	return nil
}

func (v *VPN) createInitialIPAddressPool(ctx context.Context, c client.Client) error {
	v.AvailableIPAddressPool = IPAddressPool{}

	subnet, err := netip.ParsePrefix(v.GetSubnetCIDR())
	if err != nil {
		return err
	}

	// set all IP addresses in pool to available, but reserve first two addresses for relay and client
	idx := 0
	for nextIP := subnet.Addr(); subnet.Contains(nextIP); nextIP = nextIP.Next() {
		ip := net.IP(nextIP.AsSlice())
		v.setIPAddressAvailable(ip)
		if err := v.reserveIPAddresses(ctx, c, ip, idx); err != nil {
			return err
		}
		idx++
	}
	return nil
}

// Reserves the first two IP addresses for the relay and client
func (v *VPN) reserveIPAddresses(ctx context.Context, c client.Client, ip net.IP, idx int) error {
	if idx == 0 {
		if err := v.Relay().UpdateIPAddress(ctx, c, constants.RelayName, "cluster-infra", ip); err != nil {
			return err
		}
		v.setIPAddressUnavailable(ip, "relay")
	} else if idx == 1 {
		if err := v.Client().UpdateIPAddress(ctx, c, constants.ClientName, "cluster-infra", ip); err != nil {
			return err
		}
		v.setIPAddressUnavailable(ip, "client")
	}
	return nil
}

// Returns a set of IP addresses assigned to the current VPNConfig objects
func getVPNConfigList(ctx context.Context, c client.Client) (*v1vpnconfig.VPNConfigList, error) {
	vpnConfigs := &v1vpnconfig.VPNConfigList{}
	if err := c.List(ctx, vpnConfigs, &client.ListOptions{Namespace: constants.VPNNamespace}); err != nil {
		return nil, err
	}
	return vpnConfigs, nil
}

// Check if a given IP address is in the subnet
func (v *VPN) IPAddressIsInSubnet(ip net.IP) (bool, error) {
	if ip == nil {
		return false, nil
	}

	subnetCIDR := v.GetSubnetCIDR()
	subnet, err := netip.ParsePrefix(subnetCIDR)
	if err != nil {
		return false, err
	}

	address, err := netip.ParseAddr(ip.String())
	if err != nil {
		return false, err
	}

	return subnet.Contains(address), nil
}

// Return an available IP address from the address pool and set as unavailable
func (v *VPN) RequestAvailableIPAddress(clusterEdgeID string) (net.IP, error) {
	if v.AvailableIPAddressPool == nil {
		return nil, ErrSubnetNotConfigured
	}

	for nextIP := range v.AvailableIPAddressPool {
		ip := net.ParseIP(nextIP)
		if v.CheckIPAddressIsAvailable(ip) {
			v.setIPAddressUnavailable(ip, clusterEdgeID)
			return ip, nil
		}
	}
	return nil, ErrNoIPAddressesAvailable
}

func (v *VPN) setIPAddressUnavailable(ip net.IP, clusterEdgeID string) {
	if v.AvailableIPAddressPool != nil {
		v.AvailableIPAddressPool[ip.String()] = clusterEdgeID
	}
}

func (v *VPN) setIPAddressAvailable(ip net.IP) {
	if v.AvailableIPAddressPool != nil {
		v.AvailableIPAddressPool[ip.String()] = ""
	}
}

func (v *VPN) CheckIPAddressIsAvailable(ip net.IP) bool {
	if v.AvailableIPAddressPool == nil {
		return false
	}
	return v.AvailableIPAddressPool[ip.String()] == ""
}

func (v *VPN) CheckIPAddressIsAssignedToStore(ip net.IP, clusterEdgeID string) bool {
	if v.AvailableIPAddressPool == nil {
		return false
	}
	return v.AvailableIPAddressPool[ip.String()] == clusterEdgeID
}