package engine

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

const (
	ifaceCacheTTL = 30 * time.Second
	arpCacheTTL   = 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
}

func newSourceMACResolver() *sourceMACResolver {
	return &sourceMACResolver{
		ifaceByIP: make(map[string]net.HardwareAddr),
		arpByIP:   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
	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
	}
	r.mu.RUnlock()

	if ifaceRefreshDue {
		r.refreshIfaceCache(now)
	}
	if arpRefreshDue {
		r.refreshARPCache(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...)
	}
	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()
}