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 hkey any
	switch addrlen {
	case net.IPv4len:
		hkey = binary.NativeEndian.Uint32(srcIp)
	case net.IPv6len:
		hkey = srcIp.To16().String()
	}

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

	_, 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 xprev any
	var loaded bool
	switch addrlen {
	case net.IPv4len:
		xprev, loaded = addr4.LoadOrStore(hkey, curr)
	case net.IPv6len:
		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")
	}

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

	return curr.SrcAddr
}