package main

import (
	"crypto/rand"
	"encoding/binary"
	"log"
	"math"
	"net"
	"sync"
	"time"
)

const (
	addrMapHouseKeepInterval = time.Second / 2
)

var (
	addr4, addr6 sync.Map
)

type AddrMap struct {
	SrcAddr net.IP
	DstAddr net.IP
	Ttl     uint32
	Created time.Time
}

func (a *AddrMap) GetTtl() int32 {
	x := math.Trunc(time.Since(a.Created).Round(addrMapHouseKeepInterval).Seconds())
	return int32(a.Ttl) - int32(x)
}

func setupAddrMapHousekeeping() {
	go func() {
		for {
			time.Sleep(addrMapHouseKeepInterval)
			addr4.Range(func(key, value any) bool {
				a, ok := value.(AddrMap)
				if ok {
					if a.GetTtl() > 1 {
						return true
					}
				}
				// delete if value is bogus or if ttl is less than second
				addr4.Delete(key)
				return true
			})
		}
	}()
	go func() {
		for {
			time.Sleep(addrMapHouseKeepInterval)
			addr6.Range(func(key, value any) bool {
				a, ok := value.(AddrMap)
				if ok {
					if a.GetTtl() > 1 {
						return true
					}
				}
				// delete if value is bogus or if ttl is less than second
				addr6.Delete(key)
				return true
			})
		}
	}()
}

func addrMapGet(srcIp net.IP, dstCidr *net.IPNet, ttl uint32) net.IP {
	addrlen := len(srcIp)
	switch addrlen {
	case net.IPv4len, net.IPv6len:
	default:
		log.Fatalf("addrMapGet(): src size mismatch: %v", addrlen)
	}
	if addrlen != len(dstCidr.IP) {
		log.Fatalf("addrMapGet(): src/dst size mismatch: %v vs %v", addrlen, len(dstCidr.IP))
	}
	if addrlen != len(dstCidr.Mask) {
		log.Fatalf("addrMapGet(): src/dst size mismatch: %v vs %v", addrlen, len(dstCidr.IP))
	}

	var curr AddrMap
	curr.SrcAddr = make([]byte, addrlen)
	curr.DstAddr = make([]byte, addrlen)
	copy(curr.DstAddr, srcIp)
	curr.Ttl = ttl

	for {
		_, err := rand.Read(curr.SrcAddr)
		if err != nil {
			log.Fatalf("rand.Read(): error %v", err)
		}

		// adjust random bytes to dstCidr
		for i := range addrlen / 4 {
			a := binary.NativeEndian.Uint32(dstCidr.IP[i*4:])
			b := binary.NativeEndian.Uint32(curr.SrcAddr[i*4:])
			m := binary.NativeEndian.Uint32(dstCidr.Mask[i*4:])
			a += (b & ^m)
			binary.NativeEndian.PutUint32(curr.SrcAddr[i*4:], a)
		}

		curr.Created = time.Now()

		var hkey, xprev any
		var loaded bool
		switch addrlen {
		case net.IPv4len:
			hkey = binary.NativeEndian.Uint32(curr.SrcAddr)
			xprev, loaded = addr4.LoadOrStore(hkey, curr)
		case net.IPv6len:
			hkey = binary.NativeEndian.Uint64(curr.SrcAddr[net.IPv6len/2:])
			xprev, loaded = addr6.LoadOrStore(hkey, curr)
		}
		if !loaded {
			// early return
			return curr.SrcAddr
		}

		prev, ok := xprev.(AddrMap)
		if !ok {
			log.Fatalf("addrMapGet(): wrong value type from sync.Map")
		}

		if !net.IP.Equal(curr.SrcAddr, prev.SrcAddr) {
			// generate next random address
			continue
		}
		if !net.IP.Equal(curr.DstAddr, prev.DstAddr) {
			// generate next random address
			continue
		}

		if prev.GetTtl() < int32(curr.Ttl) {
			switch addrlen {
			case net.IPv4len:
				addr4.Store(hkey, curr)
			case net.IPv6len:
				addr6.Store(hkey, curr)
			}
		}

		break
	}

	return curr.SrcAddr
}