package forwarder

import (
	"net"
	"strconv"
	"sync/atomic"

	wgdevice "golang.zx2c4.com/wireguard/device"
	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/pkg/tcpip/header"
	"gvisor.dev/gvisor/pkg/tcpip/stack"

	nblog "github.com/netbirdio/netbird/client/firewall/uspfilter/log"
)

// PacketCapture captures raw packets for debugging. Implementations must be
// safe for concurrent use and must not block.
type PacketCapture interface {
	Offer(data []byte, outbound bool)
}

// endpoint implements stack.LinkEndpoint and handles integration with the wireguard device
type endpoint struct {
	logger     *nblog.Logger
	dispatcher stack.NetworkDispatcher
	device     *wgdevice.Device
	mtu        atomic.Uint32
	capture    atomic.Pointer[PacketCapture]
}

func (e *endpoint) Attach(dispatcher stack.NetworkDispatcher) {
	e.dispatcher = dispatcher
}

func (e *endpoint) IsAttached() bool {
	return e.dispatcher != nil
}

func (e *endpoint) MTU() uint32 {
	return e.mtu.Load()
}

func (e *endpoint) Capabilities() stack.LinkEndpointCapabilities {
	return stack.CapabilityNone
}

func (e *endpoint) MaxHeaderLength() uint16 {
	return 0
}

func (e *endpoint) LinkAddress() tcpip.LinkAddress {
	return ""
}

func (e *endpoint) WritePackets(pkts stack.PacketBufferList) (int, tcpip.Error) {
	var written int
	for _, pkt := range pkts.AsSlice() {
		data := stack.PayloadSince(pkt.NetworkHeader())
		if data == nil {
			continue
		}

		raw := pkt.NetworkHeader().View().AsSlice()
		if len(raw) == 0 {
			continue
		}
		var address tcpip.Address
		if raw[0]>>4 == 6 {
			address = header.IPv6(raw).DestinationAddress()
		} else {
			address = header.IPv4(raw).DestinationAddress()
		}

		pktBytes := data.AsSlice()
		if err := e.device.CreateOutboundPacket(pktBytes, address.AsSlice()); err != nil {
			e.logger.Error1("CreateOutboundPacket: %v", err)
			continue
		}

		if pc := e.capture.Load(); pc != nil {
			(*pc).Offer(pktBytes, true)
		}
		written++
	}

	return written, nil
}

func (e *endpoint) Wait() {
	// not required
}

func (e *endpoint) ARPHardwareType() header.ARPHardwareType {
	return header.ARPHardwareNone
}

func (e *endpoint) AddHeader(*stack.PacketBuffer) {
	// not required
}

func (e *endpoint) ParseHeader(*stack.PacketBuffer) bool {
	return true
}

func (e *endpoint) Close() {
	// Endpoint cleanup - nothing to do as device is managed externally
}

func (e *endpoint) SetLinkAddress(tcpip.LinkAddress) {
	// Link address is not used for this endpoint type
}

func (e *endpoint) SetMTU(mtu uint32) {
	e.mtu.Store(mtu)
}

func (e *endpoint) SetOnCloseAction(func()) {
	// No action needed on close
}

type epID stack.TransportEndpointID

func (i epID) String() string {
	// src and remote is swapped
	return net.JoinHostPort(i.RemoteAddress.String(), strconv.Itoa(int(i.RemotePort))) +
		" → " +
		net.JoinHostPort(i.LocalAddress.String(), strconv.Itoa(int(i.LocalPort)))
}
