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
	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
}

func NewTunnelRunner(cfg TunnelInstance) (*TunnelRunner, error) {
	if err := validateConfig(&cfg); err != nil {
		return nil, err
	}
	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, 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) 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) 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)
	comp := make([]byte, lz4.CompressBlockBound(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 t.cfg.Compression && n > 40 {
			size, err := lz4.CompressBlock(payload, comp, nil)
			if err == nil && size > 0 && size < n-16 {
				payload = comp[:size]
				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.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)
	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
}