package aerotun

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"net"
	"net/netip"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pierrec/lz4/v4"
)

const peerTimeout = 65 * time.Second

type TunnelRunner struct {
	secure                                                               *secureProtocol
	allowed                                                              []netip.Prefix
	effectiveMTU                                                         atomic.Int64
	configuredMTU                                                        int
	mu                                                                   sync.RWMutex
	cfg                                                                  TunnelInstance
	dev                                                                  VirtualDevice
	udpConn                                                              *net.UDPConn
	protocol                                                             *protocolState
	remoteAddr                                                           *net.UDPAddr
	stopChan                                                             chan struct{}
	closed                                                               bool
	wg                                                                   sync.WaitGroup
	routes                                                               routeManager
	lastHandshake                                                        time.Time
	lastError                                                            string
	rxBytes, txBytes, rawBytesIn, rawBytesOut, wireBytesIn, wireBytesOut uint64
	savedBytes, packetsRx, packetsTx, compressedPackets, rawPackets      uint64
	droppedPackets, authFailures, replayDrops                            uint64
	dropInvalidOutgoing, dropMissingSession, dropSendFailure             uint64
	dropInvalidIncoming, dropSourcePolicy, dropInterfaceWrite            uint64
}

func (t *TunnelRunner) recordDrop(kind *uint64, count uint64) {
	atomic.AddUint64(&t.droppedPackets, count)
	atomic.AddUint64(kind, count)
}

func NewTunnelRunner(cfg TunnelInstance) (*TunnelRunner, error) {
	if err := validateConfig(&cfg); err != nil {
		return nil, err
	}
	configuredMTU := cfg.MTU
	var remote *net.UDPAddr
	var err error
	if cfg.RemoteEndpoint != "" {
		remote, err = net.ResolveUDPAddr("udp", cfg.RemoteEndpoint)
		if err != nil {
			return nil, fmt.Errorf("resolve remote endpoint: %w", err)
		}
		if remote.IP == nil || remote.IP.IsUnspecified() || remote.IP.IsMulticast() {
			return nil, fmt.Errorf("remote endpoint must be a unicast IP")
		}
	}
	var p *protocolState
	var secure *secureProtocol
	if cfg.Protocol == 3 {
		secure, err = newSecureProtocol(cfg)
	} else {
		p, err = newProtocol(cfg.PSK)
	}
	if err != nil {
		return nil, err
	}

	if remote != nil {
		cfg.MTU, err = underlayMTUForHeader(cfg.MTU, remote.IP, packetHeaderSize(cfg.Protocol))
		if err != nil {
			return nil, err
		}
	}
	if local, _, _ := net.ParseCIDR(cfg.LocalAddress); local.To4() == nil && cfg.MTU < 1280 {
		return nil, fmt.Errorf("underlay cannot carry IPv6 tunnel MTU 1280")
	}
	dev, err := CreateDevice(cfg.Name, cfg.MTU, cfg.LocalAddress)
	if err != nil {
		return nil, fmt.Errorf("create interface: %w", err)
	}
	if err = configurePeer(dev, cfg); err != nil {
		_ = dev.Close()
		return nil, err
	}
	conn, err := net.ListenUDP("udp", &net.UDPAddr{Port: cfg.ListenPort})
	if err != nil {
		_ = dev.Close()
		return nil, fmt.Errorf("bind UDP: %w", err)
	}
	bufferSize := 4 << 20
	if cfg.Protocol == 3 {
		bufferSize = 1 << 20
	}
	_ = conn.SetReadBuffer(bufferSize)
	_ = conn.SetWriteBuffer(bufferSize)
	runner := &TunnelRunner{secure: secure, cfg: cfg, configuredMTU: configuredMTU, dev: dev, udpConn: conn, protocol: p, remoteAddr: remote, stopChan: make(chan struct{}), routes: routeManager{run: runIP}}
	for _, p := range cfg.AllowedIPs {
		prefix, _ := netip.ParsePrefix(p)
		runner.allowed = append(runner.allowed, prefix.Masked())
	}
	runner.effectiveMTU.Store(int64(cfg.MTU))
	return runner, nil
}

func (t *TunnelRunner) Start() error {
	if t.cfg.DefaultRoute {
		if err := t.routes.install(t.cfg.Name, t.remoteAddr.IP); err != nil {
			return err
		}
	}
	t.wg.Add(3)
	go func() { defer t.wg.Done(); t.txLoop() }()
	go func() { defer t.wg.Done(); t.rxLoop() }()
	go func() { defer t.wg.Done(); t.keepaliveLoop() }()
	log.Printf("[AeroTUN] %s started: protocol v%d, UDP %s, MTU %d", t.cfg.Name, t.cfg.Protocol, t.udpConn.LocalAddr(), t.cfg.MTU)
	return nil
}
func (t *TunnelRunner) IsRunning() bool {
	if t == nil {
		return false
	}
	t.mu.RLock()
	defer t.mu.RUnlock()
	return !t.closed
}

func (t *TunnelRunner) Stop() {
	t.mu.Lock()
	if t.closed {
		t.mu.Unlock()
		return
	}
	t.closed = true
	close(t.stopChan)
	t.mu.Unlock()
	_ = t.udpConn.Close()
	_ = t.dev.Close()
	t.wg.Wait()
	t.routes.cleanup()
}
func (t *TunnelRunner) setError(err error) {
	t.mu.Lock()
	t.lastError = err.Error()
	t.mu.Unlock()
}
func (t *TunnelRunner) clearError() {
	t.mu.Lock()
	t.lastError = ""
	t.mu.Unlock()
}
func (t *TunnelRunner) target() *net.UDPAddr {
	t.mu.RLock()
	defer t.mu.RUnlock()
	return t.remoteAddr
}
func (t *TunnelRunner) send(flags byte, payload []byte, target *net.UDPAddr, session sessionID) error {
	if t.secure != nil {
		var packet [64]byte
		var nonce [24]byte
		if target == nil {
			return errSession
		}
		b, err := t.secure.sealInto(packet[:0], &nonce, flags, payload)
		if err != nil {
			return err
		}
		n, err := t.udpConn.WriteToUDP(b, target)
		if err == nil {
			atomic.AddUint64(&t.wireBytesOut, uint64(n))
		}
		return err
	}
	if target == nil {
		return errSession
	}
	b, err := t.protocol.seal(flags, payload, session)
	if err != nil {
		return err
	}
	n, err := t.udpConn.WriteToUDP(b, target)
	if err == nil {
		atomic.AddUint64(&t.wireBytesOut, uint64(n))
	}
	return err
}
func (t *TunnelRunner) sendControl(flags byte, target *net.UDPAddr, session sessionID) error {
	payload := make([]byte, 2)
	binary.BigEndian.PutUint16(payload, uint16(t.effectiveMTU.Load()))
	return t.send(flags, payload, target, session)
}
func (t *TunnelRunner) txLoop() {
	if t.secure != nil {
		t.txLoopSecure()
		return
	}
	buf := make([]byte, maxPacketSize)
	var comp *adaptiveCompressor
	if t.cfg.Compression {
		comp = newAdaptiveCompressor(maxPacketSize)
	}
	for {
		n, err := t.dev.Read(buf)
		if err != nil {
			select {
			case <-t.stopChan:
				return
			default:
			}
			t.setError(fmt.Errorf("interface read: %w", err))
			time.Sleep(10 * time.Millisecond)
			continue
		}
		if n == 0 {
			continue
		}
		if n > int(t.effectiveMTU.Load()) || !validIPPacket(buf[:n]) {
			atomic.AddUint64(&t.droppedPackets, 1)
			t.setError(fmt.Errorf("invalid or oversized interface packet (%d bytes, MTU %d)", n, t.effectiveMTU.Load()))
			continue
		}
		payload := buf[:n]
		flags := byte(0)
		if comp != nil {
			if compressed, ok := comp.Compress(payload); ok {
				payload = compressed
				flags = FlagCompressedLZ4
			}
		}
		if err := t.send(flags, payload, t.target(), zeroSession); err != nil {
			atomic.AddUint64(&t.droppedPackets, 1)
			if !errors.Is(err, errSession) {
				t.setError(fmt.Errorf("UDP send: %w", err))
			}
			continue
		}
		atomic.AddUint64(&t.txBytes, uint64(n))
		atomic.AddUint64(&t.rawBytesOut, uint64(n))
		atomic.AddUint64(&t.packetsTx, 1)
		if flags == FlagCompressedLZ4 {
			atomic.AddUint64(&t.compressedPackets, 1)
			atomic.AddUint64(&t.savedBytes, uint64(n-len(payload)))
		} else {
			atomic.AddUint64(&t.rawPackets, 1)
		}
	}
}
func (t *TunnelRunner) rxLoop() {
	if t.secure != nil {
		t.rxLoopSecure()
		return
	}
	buf := make([]byte, maxPacketSize)
	decomp := make([]byte, maxPacketSize)
	for {
		n, src, err := t.udpConn.ReadFromUDP(buf)
		if err != nil {
			select {
			case <-t.stopChan:
				return
			default:
			}
			t.setError(err)
			continue
		}
		// A client pins its configured endpoint; servers support authenticated roaming.
		if t.cfg.Mode == "client" && !udpAddrsEqual(t.target(), src) {
			continue
		}
		payload, flags, sender, _, err := t.protocol.open(buf[:n])
		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(n))
		if flags == flagHello {
			if err := t.sendControl(flagAck, src, sender); err != nil {
				t.setError(err)
			}
			continue
		}
		// A replayable HELLO must not change MTU or connection state. Only
		// packets bound to our current boot/session may change the dataplane.
		previousMTU := t.effectiveMTU.Load()
		if flags == flagAck || flags == flagConfirm || flags == FlagKeepalive {
			if err := t.adjustMTU(src.IP, int(binary.BigEndian.Uint16(payload))); err != nil {
				t.setError(err)
				continue
			}
		}
		t.mu.Lock()
		t.remoteAddr = src
		t.lastHandshake = time.Now()
		t.mu.Unlock()
		if flags == flagAck {
			if err := t.sendControl(flagConfirm, src, sender); err != nil {
				t.setError(err)
			}
			continue
		}
		if flags == FlagKeepalive || flags == flagConfirm {
			// Propagate a lower underlay/path MTU immediately, including when
			// it was discovered only after receiving the final confirmation.
			if t.effectiveMTU.Load() < previousMTU {
				if err := t.sendControl(FlagKeepalive, src, sender); err != nil {
					t.setError(err)
				}
			}
			continue
		}
		if flags == FlagCompressedLZ4 {
			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) {
			atomic.AddUint64(&t.droppedPackets, 1)
			continue
		}
		wrote, err := t.dev.Write(payload)
		if err != nil || wrote != len(payload) {
			atomic.AddUint64(&t.droppedPackets, 1)
			t.setError(fmt.Errorf("interface write: %d/%d bytes, %v", wrote, len(payload), err))
			continue
		}
		atomic.AddUint64(&t.rxBytes, uint64(wrote))
		atomic.AddUint64(&t.rawBytesIn, uint64(wrote))
		atomic.AddUint64(&t.packetsRx, 1)
	}
}
func (t *TunnelRunner) keepaliveLoop() {
	if t.secure != nil {
		t.keepaliveSecure()
		return
	}
	// Initiate immediately, then retry quickly until the peer is authenticated.
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()
	lastSent := time.Time{}
	for {
		target := t.target()
		t.mu.RLock()
		connected := !t.lastHandshake.IsZero() && time.Since(t.lastHandshake) < peerTimeout
		t.mu.RUnlock()
		if target != nil && (!connected || time.Since(lastSent) >= 20*time.Second) {
			flags := byte(flagHello)
			if connected {
				flags = FlagKeepalive
			}
			if err := t.sendControl(flags, target, zeroSession); err != nil {
				t.setError(err)
			}
			lastSent = time.Now()
		}
		select {
		case <-t.stopChan:
			return
		case <-ticker.C:
		}
	}
}
func validIPPacket(b []byte) bool {
	if len(b) < 20 {
		return false
	}
	if b[0]>>4 == 4 {
		ihl := int(b[0]&15) * 4
		return ihl >= 20 && len(b) >= ihl && int(b[2])<<8|int(b[3]) == len(b)
	}
	if b[0]>>4 == 6 {
		return len(b) >= 40 && (int(b[4])<<8|int(b[5]))+40 == len(b)
	}
	return false
}
func (t *TunnelRunner) GetTelemetry() TunnelInstance {
	t.mu.RLock()
	inst := t.cfg
	if _, ok := t.dev.(*nativeTun); ok {
		inst.Backend = "tun"
	} else {
		inst.Backend = "veth"
	}
	if t.secure != nil {
		inst.Handshakes = t.secure.handshakes.Load()
	}
	inst.MTU = int(t.effectiveMTU.Load())
	inst.ConfiguredMTU = t.configuredMTU
	inst.Running = !t.closed
	inst.Connected = inst.Running && !t.lastHandshake.IsZero() && time.Since(t.lastHandshake) < peerTimeout
	inst.ActivePeerEndpoint = ""
	if inst.Connected && t.remoteAddr != nil {
		inst.ActivePeerEndpoint = t.remoteAddr.String()
	}
	inst.LastHandshake = "Never"
	if !t.lastHandshake.IsZero() {
		inst.LastHandshake = time.Since(t.lastHandshake).Truncate(time.Second).String() + " ago"
	}
	if t.secure != nil {
		session := t.secure.sessions.Load().current
		if session == nil || time.Since(session.born) > secureRejectAfter {
			inst.Connected = false
			inst.ActivePeerEndpoint = ""
		}
	}
	inst.LastError = t.lastError
	t.mu.RUnlock()
	inst.RxBytes = atomic.LoadUint64(&t.rxBytes)
	inst.TxBytes = atomic.LoadUint64(&t.txBytes)
	inst.RawBytesIn = atomic.LoadUint64(&t.rawBytesIn)
	inst.RawBytesOut = atomic.LoadUint64(&t.rawBytesOut)
	inst.WireBytesIn = atomic.LoadUint64(&t.wireBytesIn)
	inst.WireBytesOut = atomic.LoadUint64(&t.wireBytesOut)
	inst.SavedBytes = atomic.LoadUint64(&t.savedBytes)
	inst.PacketsRx = atomic.LoadUint64(&t.packetsRx)
	inst.PacketsTx = atomic.LoadUint64(&t.packetsTx)
	inst.CompressedPackets = atomic.LoadUint64(&t.compressedPackets)
	inst.RawPackets = atomic.LoadUint64(&t.rawPackets)
	inst.DroppedPackets = atomic.LoadUint64(&t.droppedPackets)
	inst.AuthFailures = atomic.LoadUint64(&t.authFailures)
	inst.ReplayDrops = atomic.LoadUint64(&t.replayDrops)
	if t.secure != nil {
		inst.DropBreakdown = &DropBreakdown{InvalidOutgoing: atomic.LoadUint64(&t.dropInvalidOutgoing), MissingSession: atomic.LoadUint64(&t.dropMissingSession), SendFailure: atomic.LoadUint64(&t.dropSendFailure), InvalidIncoming: atomic.LoadUint64(&t.dropInvalidIncoming), SourcePolicy: atomic.LoadUint64(&t.dropSourcePolicy), InterfaceWrite: atomic.LoadUint64(&t.dropInterfaceWrite)}
	}
	inst.CompressionRatio = "1.00x (0% saved)"
	if inst.RawBytesOut > inst.WireBytesOut && inst.WireBytesOut > 0 {
		inst.CompressionRatio = fmt.Sprintf("%.2fx (%.1f%% saved)", float64(inst.RawBytesOut)/float64(inst.WireBytesOut), (1-float64(inst.WireBytesOut)/float64(inst.RawBytesOut))*100)
	}
	return inst
}
func GenerateSecureKey() (string, error) {
	key := make([]byte, 32)
	if _, err := rand.Read(key); err != nil {
		return "", err
	}
	return fmt.Sprintf("%x", key), nil
}
func udpAddrsEqual(a, b *net.UDPAddr) bool {
	return a != nil && b != nil && a.IP.Equal(b.IP) && a.Port == b.Port && a.Zone == b.Zone
}