package aerotun

import (
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"net/netip"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/pierrec/lz4/v4"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

const secureBatchSize = 16

type batchSocket interface {
	ReadBatch([]ipv4.Message, int) (int, error)
	WriteBatch([]ipv4.Message, int) (int, error)
}

func newBatchSocket(c *net.UDPConn) batchSocket {
	if c.LocalAddr().(*net.UDPAddr).IP.To4() != nil {
		return ipv4.NewPacketConn(c)
	}
	return ipv6.NewPacketConn(c)
}

type securePacket struct {
	buf        []byte
	n, wire    int
	compressed bool
}

func (t *TunnelRunner) txLoopSecure() {
	size := t.cfg.MTU + secureHeaderSize + 16
	pool := make(chan *securePacket, secureBatchSize*2)
	queue := make(chan *securePacket, secureBatchSize)
	for range cap(pool) {
		pool <- &securePacket{buf: make([]byte, size)}
	}
	done := make(chan struct{})
	go func() {
		defer close(done)
		defer close(queue)
		for {
			var p *securePacket
			select {
			case <-t.stopChan:
				return
			case p = <-pool:
			}
			n, err := t.dev.Read(p.buf[secureHeaderSize : secureHeaderSize+t.cfg.MTU])
			if err != nil {
				select {
				case <-t.stopChan:
					return
				default:
				}
				t.setError(fmt.Errorf("secure interface read: %w", err))
				pool <- p
				time.Sleep(10 * time.Millisecond)
				continue
			}
			if n == 0 || n > int(t.effectiveMTU.Load()) || !validIPPacket(p.buf[secureHeaderSize:secureHeaderSize+n]) {
				atomic.AddUint64(&t.droppedPackets, 1)
				pool <- p
				continue
			}
			p.n = n
			select {
			case <-t.stopChan:
				return
			case queue <- p:
			}
		}
	}()
	defer func() { <-done }()
	conn := newBatchSocket(t.udpConn)
	messages := make([]ipv4.Message, secureBatchSize)
	var buffers [secureBatchSize][1][]byte
	for i := range messages {
		messages[i].Buffers = buffers[i][:]
	}
	var packets [secureBatchSize]*securePacket
	var nonce [24]byte
	var comp []byte
	if t.cfg.Compression {
		comp = make([]byte, lz4.CompressBlockBound(t.cfg.MTU))
	}
	for {
		var first *securePacket
		select {
		case <-t.stopChan:
			return
		case p, ok := <-queue:
			if !ok {
				return
			}
			first = p
		}
		count := 1
		packets[0] = first
	drain:
		for count < secureBatchSize {
			select {
			case p, ok := <-queue:
				if !ok {
					break drain
				}
				packets[count] = p
				count++
			default:
				break drain
			}
		}
		target := t.target()
		ready := 0
		for i := 0; i < count; i++ {
			p := packets[i]
			payload := p.buf[secureHeaderSize : secureHeaderSize+p.n]
			flags := byte(0)
			p.compressed = false
			p.wire = p.n
			if target != nil && t.secure.compressionEnabled() && len(comp) > 0 && p.n > 40 {
				if n, err := lz4.CompressBlock(payload, comp, nil); err == nil && n > 0 && n < p.n-16 {
					payload = comp[:n]
					flags = FlagCompressedLZ4
					p.compressed = true
					p.wire = n
				}
			}
			if target == nil {
				atomic.AddUint64(&t.droppedPackets, 1)
				pool <- p
				continue
			}
			out, err := t.secure.sealInto(p.buf[:0], &nonce, flags, payload)
			if err != nil {
				atomic.AddUint64(&t.droppedPackets, 1)
				pool <- p
				continue
			}
			packets[ready] = p
			messages[ready].Buffers[0] = out
			messages[ready].Addr = target
			ready++
		}
		sent := 0
		for sent < ready {
			n, err := conn.WriteBatch(messages[sent:ready], 0)
			for j := sent; j < sent+n; j++ {
				p := packets[j]
				atomic.AddUint64(&t.txBytes, uint64(p.n))
				atomic.AddUint64(&t.rawBytesOut, uint64(p.n))
				atomic.AddUint64(&t.packetsTx, 1)
				atomic.AddUint64(&t.wireBytesOut, uint64(len(messages[j].Buffers[0])))
				if p.compressed {
					atomic.AddUint64(&t.compressedPackets, 1)
					atomic.AddUint64(&t.savedBytes, uint64(p.n-p.wire))
				} else {
					atomic.AddUint64(&t.rawPackets, 1)
				}
			}
			sent += n
			if err != nil || n == 0 {
				if err != nil {
					t.setError(err)
				}
				atomic.AddUint64(&t.droppedPackets, uint64(ready-sent))
				break
			}
		}
		for i := 0; i < ready; i++ {
			pool <- packets[i]
		}
	}
}
func (t *TunnelRunner) allowedSource(b []byte) bool {
	var src netip.Addr
	if b[0]>>4 == 4 {
		var ip [4]byte
		copy(ip[:], b[12:16])
		src = netip.AddrFrom4(ip)
	} else {
		var ip [16]byte
		copy(ip[:], b[8:24])
		src = netip.AddrFrom16(ip)
	}
	for _, prefix := range t.allowed {
		if prefix.Contains(src) {
			return true
		}
	}
	return false
}
func (t *TunnelRunner) rxLoopSecure() {
	conn := newBatchSocket(t.udpConn)
	size := max(t.cfg.MTU+secureHeaderSize+16, handshakeMaxSize)
	messages := make([]ipv4.Message, secureBatchSize)
	buffers := make([][]byte, secureBatchSize)
	var vectors [secureBatchSize][1][]byte
	for i := range messages {
		buffers[i] = make([]byte, size)
		vectors[i][0] = buffers[i]
		messages[i].Buffers = vectors[i][:]
	}
	plain := make([]byte, t.cfg.MTU)
	var decomp []byte
	if t.cfg.Compression {
		decomp = make([]byte, t.cfg.MTU)
	}
	var nonce [24]byte
	for {
		n, err := conn.ReadBatch(messages, 0)
		if err != nil {
			select {
			case <-t.stopChan:
				return
			default:
			}
			t.setError(err)
			continue
		}
		for i := 0; i < n; i++ {
			m := &messages[i]
			src, ok := m.Addr.(*net.UDPAddr)
			if !ok || m.Flags&syscall.MSG_TRUNC != 0 || m.N < secureHeaderSize+16 {
				atomic.AddUint64(&t.droppedPackets, 1)
				continue
			}
			if t.cfg.Mode == "client" && !udpAddrsEqual(t.target(), src) {
				continue
			}
			b := m.Buffers[0][:m.N]
			if b[2] != 3 {
				atomic.AddUint64(&t.authFailures, 1)
				continue
			}
			if b[3] == secureInit || b[3] == secureResponse || b[3] == secureCookie {
				response, peerMTU, err := t.secure.handleHandshake(b, int(t.effectiveMTU.Load()), src.AddrPort())
				if err != nil {
					if !errors.Is(err, errHandshakeRate) {
						atomic.AddUint64(&t.authFailures, 1)
					}
					continue
				}
				atomic.AddUint64(&t.wireBytesIn, uint64(m.N))
				if response != nil {
					if n, err := t.udpConn.WriteToUDP(response, src); err == nil {
						atomic.AddUint64(&t.wireBytesOut, uint64(n))
					} else {
						t.setError(err)
					}
				}
				if peerMTU > 0 {
					if err := t.adjustMTU(src.IP, peerMTU); err != nil {
						t.setError(err)
						continue
					}
					t.markSecurePeer(src)
					if err := t.sendControl(FlagKeepalive, src, zeroSession); err != nil {
						t.setError(err)
					}
				}
				continue
			}
			payload, flags, peerMTU, err := t.secure.openInto(plain, &nonce, b)
			if err != nil {
				if errors.Is(err, errReplay) {
					atomic.AddUint64(&t.replayDrops, 1)
				} else {
					atomic.AddUint64(&t.authFailures, 1)
				}
				continue
			}
			atomic.AddUint64(&t.wireBytesIn, uint64(m.N))
			oldMTU := t.effectiveMTU.Load()
			if flags == FlagKeepalive {
				peerMTU = int(binary.BigEndian.Uint16(payload))
			}
			if peerMTU > 0 {
				if err := t.adjustMTU(src.IP, peerMTU); err != nil {
					t.setError(err)
					continue
				}
			}
			t.markSecurePeer(src)
			if flags == FlagKeepalive {
				if t.effectiveMTU.Load() < oldMTU {
					if err := t.sendControl(FlagKeepalive, src, zeroSession); err != nil {
						t.setError(err)
					}
				}
				continue
			}
			if flags == FlagCompressedLZ4 {
				if len(decomp) == 0 {
					atomic.AddUint64(&t.droppedPackets, 1)
					continue
				}
				size, err := lz4.UncompressBlock(payload, decomp[:int(t.effectiveMTU.Load())])
				if err != nil || size == 0 {
					atomic.AddUint64(&t.droppedPackets, 1)
					continue
				}
				payload = decomp[:size]
			}
			if len(payload) > int(t.effectiveMTU.Load()) || !validIPPacket(payload) || !t.allowedSource(payload) {
				atomic.AddUint64(&t.droppedPackets, 1)
				continue
			}
			wrote, err := t.dev.Write(payload)
			if err != nil || wrote != len(payload) {
				atomic.AddUint64(&t.droppedPackets, 1)
				if err != nil {
					t.setError(err)
				}
				continue
			}
			atomic.AddUint64(&t.rxBytes, uint64(wrote))
			atomic.AddUint64(&t.rawBytesIn, uint64(wrote))
			atomic.AddUint64(&t.packetsRx, 1)
		}
	}
}
func (t *TunnelRunner) markSecurePeer(src *net.UDPAddr) {
	t.mu.Lock()
	if !udpAddrsEqual(t.remoteAddr, src) {
		t.remoteAddr = src
	}
	t.lastHandshake = time.Now()
	t.mu.Unlock()
}
func (t *TunnelRunner) keepaliveSecure() {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	lastSent := time.Time{}
	for {
		t.secure.expire()
		target := t.target()
		t.mu.RLock()
		connected := !t.lastHandshake.IsZero() && time.Since(t.lastHandshake) < peerTimeout
		t.mu.RUnlock()
		if target != nil {
			b, err := t.secure.initiation(int(t.effectiveMTU.Load()), !connected)
			if err != nil {
				t.setError(err)
			} else if b != nil {
				if n, err := t.udpConn.WriteToUDP(b, target); err == nil {
					atomic.AddUint64(&t.wireBytesOut, uint64(n))
				} else {
					t.setError(err)
				}
			}
			if connected && time.Since(lastSent) >= 20*time.Second {
				if err := t.sendControl(FlagKeepalive, target, zeroSession); err != nil {
					t.setError(err)
				}
				lastSent = time.Now()
			}
		}
		select {
		case <-t.stopChan:
			return
		case <-ticker.C:
		}
	}
}