package engine

import (
	"bufio"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/mdlayher/netlink"
	"golang.org/x/sys/unix"
)

const (
	ifaceCacheTTL = 30 * time.Second
	arpCacheTTL   = 10 * time.Second
	ndpCacheTTL   = 10 * time.Second
)

type sourceMACResolver struct {
	mu sync.RWMutex

	lastIfaceRefresh time.Time
	ifaceByIP        map[string]net.HardwareAddr

	lastARPRefresh time.Time
	arpByIP        map[string]net.HardwareAddr

	lastNDPRefresh time.Time
	ndpByIP        map[string]net.HardwareAddr
}

func newSourceMACResolver() *sourceMACResolver {
	return &sourceMACResolver{
		ifaceByIP: make(map[string]net.HardwareAddr),
		arpByIP:   make(map[string]net.HardwareAddr),
		ndpByIP:   make(map[string]net.HardwareAddr),
	}
}

func (r *sourceMACResolver) Resolve(ip net.IP) net.HardwareAddr {
	if ip == nil {
		return nil
	}
	ipKey := ip.String()
	if ipKey == "" {
		return nil
	}

	now := time.Now()
	r.mu.RLock()
	ifaceRefreshDue := now.Sub(r.lastIfaceRefresh) > ifaceCacheTTL
	arpRefreshDue := now.Sub(r.lastARPRefresh) > arpCacheTTL
	ndpRefreshDue := now.Sub(r.lastNDPRefresh) > ndpCacheTTL
	if mac := r.ifaceByIP[ipKey]; len(mac) != 0 {
		out := append(net.HardwareAddr(nil), mac...)
		r.mu.RUnlock()
		return out
	}
	if mac := r.arpByIP[ipKey]; len(mac) != 0 && !arpRefreshDue {
		out := append(net.HardwareAddr(nil), mac...)
		r.mu.RUnlock()
		return out
	}
	if mac := r.ndpByIP[ipKey]; len(mac) != 0 && !ndpRefreshDue {
		out := append(net.HardwareAddr(nil), mac...)
		r.mu.RUnlock()
		return out
	}
	r.mu.RUnlock()

	if ifaceRefreshDue {
		r.refreshIfaceCache(now)
	}
	if arpRefreshDue {
		r.refreshARPCache(now)
	}
	if ndpRefreshDue {
		r.refreshNDPCache(now)
	}

	r.mu.RLock()
	defer r.mu.RUnlock()
	if mac := r.ifaceByIP[ipKey]; len(mac) != 0 {
		return append(net.HardwareAddr(nil), mac...)
	}
	if mac := r.arpByIP[ipKey]; len(mac) != 0 {
		return append(net.HardwareAddr(nil), mac...)
	}
	if mac := r.ndpByIP[ipKey]; len(mac) != 0 {
		return append(net.HardwareAddr(nil), mac...)
	}
	return nil
}

func (r *sourceMACResolver) refreshIfaceCache(now time.Time) {
	interfaces, err := net.Interfaces()
	if err != nil {
		return
	}

	m := make(map[string]net.HardwareAddr)
	for _, iface := range interfaces {
		if len(iface.HardwareAddr) == 0 {
			continue
		}
		addrs, err := iface.Addrs()
		if err != nil {
			continue
		}
		for _, addr := range addrs {
			ipNet, ok := addr.(*net.IPNet)
			if !ok || ipNet.IP == nil {
				continue
			}
			m[ipNet.IP.String()] = append(net.HardwareAddr(nil), iface.HardwareAddr...)
		}
	}

	r.mu.Lock()
	r.ifaceByIP = m
	r.lastIfaceRefresh = now
	r.mu.Unlock()
}

func (r *sourceMACResolver) refreshARPCache(now time.Time) {
	f, err := os.Open("/proc/net/arp")
	if err != nil {
		return
	}
	defer f.Close()

	m := make(map[string]net.HardwareAddr)
	scanner := bufio.NewScanner(f)
	lineNo := 0
	for scanner.Scan() {
		lineNo++
		if lineNo == 1 {
			continue // header
		}
		fields := strings.Fields(scanner.Text())
		if len(fields) < 4 {
			continue
		}
		ipStr := fields[0]
		hwAddr := fields[3]
		if hwAddr == "00:00:00:00:00:00" {
			continue
		}
		mac, err := net.ParseMAC(hwAddr)
		if err != nil {
			continue
		}
		m[ipStr] = append(net.HardwareAddr(nil), mac...)
	}
	if err := scanner.Err(); err != nil {
		return
	}

	r.mu.Lock()
	r.arpByIP = m
	r.lastARPRefresh = now
	r.mu.Unlock()
}

func (r *sourceMACResolver) refreshNDPCache(now time.Time) {
	m, ok := readNeighborCacheFile("/proc/net/ndisc_cache")
	if !ok || len(m) == 0 {
		// Fallback for environments without /proc/net/ndisc_cache.
		m = readIPv6NeighNetlink()
	}
	r.mu.Lock()
	r.ndpByIP = m
	r.lastNDPRefresh = now
	r.mu.Unlock()
}

func readNeighborCacheFile(path string) (map[string]net.HardwareAddr, bool) {
	f, err := os.Open(path)
	if err != nil {
		return nil, false
	}
	defer f.Close()

	m := make(map[string]net.HardwareAddr)
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		ip, mac, ok := parseNeighborLine(scanner.Text())
		if !ok {
			continue
		}
		m[ip] = mac
	}
	if err := scanner.Err(); err != nil {
		return nil, false
	}
	return m, true
}

func readIPv6NeighNetlink() map[string]net.HardwareAddr {
	const (
		ndMsgLen  = 12
		ndaDst    = 1
		ndaLLAddr = 2
	)

	m := make(map[string]net.HardwareAddr)
	conn, err := netlink.Dial(unix.NETLINK_ROUTE, nil)
	if err != nil {
		return m
	}
	defer conn.Close()

	req := make([]byte, ndMsgLen)
	req[0] = unix.AF_INET6
	msgs, err := conn.Execute(netlink.Message{
		Header: netlink.Header{
			Type:  unix.RTM_GETNEIGH,
			Flags: netlink.Request | netlink.Dump,
		},
		Data: req,
	})
	if err != nil {
		return m
	}

	for _, msg := range msgs {
		if msg.Header.Type != unix.RTM_NEWNEIGH || len(msg.Data) < ndMsgLen || msg.Data[0] != unix.AF_INET6 {
			continue
		}
		attrs, err := netlink.UnmarshalAttributes(msg.Data[ndMsgLen:])
		if err != nil {
			continue
		}
		var ipStr string
		var mac net.HardwareAddr
		for _, a := range attrs {
			switch a.Type {
			case ndaDst:
				if len(a.Data) == net.IPv6len {
					ipStr = net.IP(append([]byte(nil), a.Data...)).String()
				}
			case ndaLLAddr:
				if len(a.Data) >= 6 {
					candidate := append(net.HardwareAddr(nil), a.Data...)
					if candidate.String() != "00:00:00:00:00:00" {
						mac = candidate
					}
				}
			}
		}
		if ipStr != "" && len(mac) != 0 {
			m[ipStr] = mac
		}
	}

	return m
}

func parseNeighborLine(line string) (string, net.HardwareAddr, bool) {
	fields := strings.Fields(line)
	if len(fields) < 2 {
		return "", nil, false
	}
	var ipStr string
	var mac net.HardwareAddr
	for _, f := range fields {
		if ip := net.ParseIP(f); ip != nil {
			ipStr = ip.String()
			continue
		}
		if parsedMAC, err := net.ParseMAC(f); err == nil {
			// Ignore unresolved/bogus entries.
			if parsedMAC.String() != "00:00:00:00:00:00" {
				mac = append(net.HardwareAddr(nil), parsedMAC...)
			}
		}
	}
	if ipStr == "" || len(mac) == 0 {
		return "", nil, false
	}
	return ipStr, mac, true
}