package engine

import (
	"bufio"
	"net"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"
)

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 = readIPv6NeighCommand()
	}
	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 readIPv6NeighCommand() map[string]net.HardwareAddr {
	m := make(map[string]net.HardwareAddr)
	out, err := exec.Command("ip", "-6", "neigh", "show").Output()
	if err != nil {
		return m
	}
	for _, line := range strings.Split(string(out), "\n") {
		ip, mac, ok := parseNeighborLine(line)
		if !ok {
			continue
		}
		m[ip] = 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
}