package vl2

import (
	"encoding/binary"
	"log/slog"
	"net"
	"sync"
	"time"
)

type ARPEntry struct {
	MAC      net.HardwareAddr
	LastSeen time.Time
}

type ARPProxy struct {
	cache map[string]*ARPEntry
	mu    sync.RWMutex
	log   *slog.Logger
}

func NewARPProxy(log *slog.Logger) *ARPProxy {
	return &ARPProxy{
		cache: make(map[string]*ARPEntry),
		log:   log.With("component", "arp"),
	}
}

func (a *ARPProxy) Learn(ip net.IP, mac net.HardwareAddr) {
	key := ip.String()
	a.mu.Lock()
	defer a.mu.Unlock()
	a.cache[key] = &ARPEntry{
		MAC:      mac,
		LastSeen: time.Now(),
	}
}

func (a *ARPProxy) Lookup(ip net.IP) net.HardwareAddr {
	key := ip.String()
	a.mu.RLock()
	defer a.mu.RUnlock()
	entry, ok := a.cache[key]
	if !ok {
		return nil
	}
	return entry.MAC
}

func (a *ARPProxy) HandleARP(frame *EthernetFrame) []byte {
	if len(frame.Payload) < 28 {
		return nil
	}
	pl := frame.Payload

	hrd := binary.BigEndian.Uint16(pl[0:2])
	pro := binary.BigEndian.Uint16(pl[2:4])
	hln := pl[4]
	pln := pl[5]
	op := binary.BigEndian.Uint16(pl[6:8])

	if hrd != 1 || pro != EtherTypeIPv4 || hln != 6 || pln != 4 {
		return nil
	}

	senderMAC := net.HardwareAddr(pl[8:14])
	senderIP := net.IP(pl[14:18])
	targetIP := net.IP(pl[24:28])

	a.Learn(senderIP, senderMAC)

	if op == 1 { // ARP request
		if localMAC := a.Lookup(targetIP); localMAC != nil {
			reply := make([]byte, 42)
			copy(reply[0:6], senderMAC)
			copy(reply[6:12], localMAC)
			binary.BigEndian.PutUint16(reply[12:14], EtherTypeARP)

			reply[14] = 0x00
			reply[15] = 0x01
			binary.BigEndian.PutUint16(reply[16:18], EtherTypeIPv4)
			reply[18] = 6
			reply[19] = 4
			binary.BigEndian.PutUint16(reply[20:22], 2)

			copy(reply[22:28], localMAC)
			copy(reply[28:32], targetIP)
			copy(reply[32:38], senderMAC)
			copy(reply[38:42], senderIP)

			return reply
		}
	}

	return nil
}

func (a *ARPProxy) PeerFromARP(frame *EthernetFrame) (net.IP, net.HardwareAddr) {
	if len(frame.Payload) < 28 {
		return nil, nil
	}
	pl := frame.Payload
	hrd := binary.BigEndian.Uint16(pl[0:2])
	pro := binary.BigEndian.Uint16(pl[2:4])
	if hrd != 1 || pro != EtherTypeIPv4 {
		return nil, nil
	}
	return net.IP(pl[14:18]), net.HardwareAddr(pl[8:14])
}

func (a *ARPProxy) CleanExpired() {
	a.mu.Lock()
	defer a.mu.Unlock()
	cutoff := time.Now().Add(-30 * time.Minute)
	for k, v := range a.cache {
		if v.LastSeen.Before(cutoff) {
			delete(a.cache, k)
		}
	}
}