package aerotun

import (
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"os/exec"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pierrec/lz4/v4"
	"golang.org/x/crypto/chacha20poly1305"
)

type TunnelRunner struct {
	mu           sync.RWMutex
	cfg          TunnelInstance
	dev          VirtualDevice
	udpConn      *net.UDPConn
	aead         cipher.AEAD
	remoteAddr   *net.UDPAddr
	txCounter    uint64
	rxMaxCounter uint64
	replayWindow uint64
	stopChan     chan struct{}
	closed       bool

	// Metrics
	rxBytes           uint64
	txBytes           uint64
	rawBytesIn        uint64
	rawBytesOut       uint64
	wireBytesIn       uint64
	wireBytesOut      uint64
	savedBytes        uint64
	packetsRx         uint64
	packetsTx         uint64
	compressedPackets uint64
	rawPackets        uint64
	lastHandshake     time.Time
}

func NewTunnelRunner(cfg TunnelInstance) (*TunnelRunner, error) {
	// 1. Derive 32-byte ChaCha20-Poly1305 key from PSK
	if cfg.PSK == "" {
		return nil, fmt.Errorf("pre-shared key (PSK) is required")
	}
	key := sha256.Sum256([]byte(cfg.PSK))

	aead, err := chacha20poly1305.New(key[:])
	if err != nil {
		return nil, fmt.Errorf("failed to initialize ChaCha20-Poly1305: %w", err)
	}

	// 2. Create Virtual Network Device (TUN or VETH fallback)
	dev, err := CreateDevice(cfg.Name, cfg.MTU, cfg.LocalAddress)
	if err != nil {
		return nil, fmt.Errorf("failed to create virtual network device %s: %w", cfg.Name, err)
	}

	// 3. Resolve UDP socket / remote address
	var localUDPAddr *net.UDPAddr
	if cfg.ListenPort > 0 {
		localUDPAddr = &net.UDPAddr{Port: cfg.ListenPort}
	} else if stringsEqual(cfg.Mode, "server") {
		localUDPAddr = &net.UDPAddr{Port: DefaultPort}
	}

	udpConn, err := net.ListenUDP("udp", localUDPAddr)
	if err != nil {
		_ = dev.Close()
		return nil, fmt.Errorf("failed to bind UDP socket on port %d: %w", cfg.ListenPort, err)
	}

	var remoteAddr *net.UDPAddr
	if cfg.RemoteEndpoint != "" {
		resolved, err := net.ResolveUDPAddr("udp", cfg.RemoteEndpoint)
		if err == nil {
			remoteAddr = resolved
		}
	}

	runner := &TunnelRunner{
		cfg:        cfg,
		dev:        dev,
		udpConn:    udpConn,
		aead:       aead,
		remoteAddr: remoteAddr,
		stopChan:   make(chan struct{}),
	}

	return runner, nil
}

func (t *TunnelRunner) Start() {
	go t.txLoop()
	go t.rxLoop()
	go t.keepaliveLoop()

	if t.cfg.DefaultRoute {
		_ = exec.Command("ip", "route", "replace", "default", "dev", t.cfg.Name).Run()
	}
	log.Printf("[AeroTUN] Tunnel %s started in %s mode (listen port %d, MTU %d)", t.cfg.Name, t.cfg.Mode, t.cfg.ListenPort, t.cfg.MTU)
}

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()
	log.Printf("[AeroTUN] Tunnel %s stopped", t.cfg.Name)
}

// txLoop: Reads raw IP packets from TUN, compresses with LZ4, encrypts with ChaCha20-Poly1305, sends via UDP
func (t *TunnelRunner) txLoop() {
	tunBuf := make([]byte, 2048)
	lz4Buf := make([]byte, 2048)
	udpPacket := make([]byte, 2048)

	for {
		select {
		case <-t.stopChan:
			return
		default:
		}

		n, err := t.dev.Read(tunBuf)
		if err != nil {
			select {
			case <-t.stopChan:
				return
			default:
				time.Sleep(10 * time.Millisecond)
				continue
			}
		}
		if n == 0 {
			continue
		}

		rawIP := tunBuf[:n]
		atomic.AddUint64(&t.rawBytesOut, uint64(n))
		atomic.AddUint64(&t.txBytes, uint64(n))
		atomic.AddUint64(&t.packetsTx, 1)

		// 1. Adaptive LZ4 Compression
		var flags byte = 0
		payload := rawIP

		if t.cfg.Compression && n > 40 {
			compressedSize, err := lz4.CompressBlock(rawIP, lz4Buf, nil)
			// Only compress if saving at least 16 bytes
			if err == nil && compressedSize > 0 && compressedSize < n-16 {
				flags |= FlagCompressedLZ4
				payload = lz4Buf[:compressedSize]
				saved := uint64(n - compressedSize)
				atomic.AddUint64(&t.savedBytes, saved)
				atomic.AddUint64(&t.compressedPackets, 1)
			} else {
				atomic.AddUint64(&t.rawPackets, 1)
			}
		} else {
			atomic.AddUint64(&t.rawPackets, 1)
		}

		// 2. Build 12-byte Authenticated Header:
		// [Magic (2B)][Version (1B)][Flags (1B)][Counter (8B)]
		counter := atomic.AddUint64(&t.txCounter, 1)
		binary.BigEndian.PutUint16(udpPacket[0:2], MagicHeader)
		udpPacket[2] = ProtocolVersion
		udpPacket[3] = flags
		binary.BigEndian.PutUint64(udpPacket[4:12], counter)

		// 3. Construct 12-byte Nonce
		var nonce [12]byte
		copy(nonce[0:4], udpPacket[0:4])
		binary.BigEndian.PutUint64(nonce[4:12], counter)

		// 4. AEAD Seal: authenticate header [0:12] as AAD, encrypt payload
		encrypted := t.aead.Seal(udpPacket[:12], nonce[:], payload, udpPacket[:12])

		// 5. Transmit UDP packet
		t.mu.RLock()
		target := t.remoteAddr
		t.mu.RUnlock()

		if target != nil {
			_, _ = t.udpConn.WriteToUDP(encrypted, target)
			atomic.AddUint64(&t.wireBytesOut, uint64(len(encrypted)))
		}
	}
}

// rxLoop: Reads UDP datagrams, verifies AEAD & anti-replay, decompresses LZ4, writes raw IP to TUN
func (t *TunnelRunner) rxLoop() {
	udpBuf := make([]byte, 2048)
	decompBuf := make([]byte, 2048)

	for {
		select {
		case <-t.stopChan:
			return
		default:
		}

		n, srcAddr, err := t.udpConn.ReadFromUDP(udpBuf)
		if err != nil {
			select {
			case <-t.stopChan:
				return
			default:
				time.Sleep(10 * time.Millisecond)
				continue
			}
		}
		if n < 28 { // 12-byte header + 16-byte Poly1305 tag minimum
			continue
		}

		packet := udpBuf[:n]
		atomic.AddUint64(&t.wireBytesIn, uint64(n))

		// 1. Verify Magic Header & Protocol Version
		magic := binary.BigEndian.Uint16(packet[0:2])
		if magic != MagicHeader || packet[2] != ProtocolVersion {
			continue
		}

		flags := packet[3]
		counter := binary.BigEndian.Uint64(packet[4:12])

		// 2. Anti-Replay Sliding Window Check
		if !t.checkAndUpdateReplay(counter) {
			continue // Drop replayed packet
		}

		// 3. Construct 12-byte Nonce & Verify AEAD
		var nonce [12]byte
		copy(nonce[0:4], packet[0:4])
		binary.BigEndian.PutUint64(nonce[4:12], counter)

		decrypted, err := t.aead.Open(nil, nonce[:], packet[12:], packet[:12])
		if err != nil {
			continue // Authentication failed
		}

		// 4. Dynamic Roaming: Learn remote peer address in server mode
		t.mu.Lock()
		if t.remoteAddr == nil || (t.cfg.Mode == "server" && !udpAddrsEqual(t.remoteAddr, srcAddr)) {
			t.remoteAddr = srcAddr
			t.lastHandshake = time.Now()
		}
		t.lastHandshake = time.Now()
		t.mu.Unlock()

		// 5. Handle Keepalive
		if flags&FlagKeepalive != 0 {
			continue
		}

		// 6. LZ4 Decompress if flagged
		var ipPayload []byte
		if flags&FlagCompressedLZ4 != 0 {
			uncompressedSize, err := lz4.UncompressBlock(decrypted, decompBuf)
			if err != nil || uncompressedSize == 0 {
				continue
			}
			ipPayload = decompBuf[:uncompressedSize]
			atomic.AddUint64(&t.rawBytesIn, uint64(uncompressedSize))
		} else {
			ipPayload = decrypted
			atomic.AddUint64(&t.rawBytesIn, uint64(len(decrypted)))
		}

		// 7. Inject IP packet into TUN device
		_, _ = t.dev.Write(ipPayload)
		atomic.AddUint64(&t.rxBytes, uint64(len(ipPayload)))
		atomic.AddUint64(&t.packetsRx, 1)
	}
}

func (t *TunnelRunner) keepaliveLoop() {
	ticker := time.NewTicker(20 * time.Second)
	defer ticker.Stop()

	keepaliveBuf := make([]byte, 32)
	for {
		select {
		case <-t.stopChan:
			return
		case <-ticker.C:
			t.mu.RLock()
			target := t.remoteAddr
			t.mu.RUnlock()

			if target == nil {
				continue
			}

			counter := atomic.AddUint64(&t.txCounter, 1)
			binary.BigEndian.PutUint16(keepaliveBuf[0:2], MagicHeader)
			keepaliveBuf[2] = ProtocolVersion
			keepaliveBuf[3] = FlagKeepalive
			binary.BigEndian.PutUint64(keepaliveBuf[4:12], counter)

			var nonce [12]byte
			copy(nonce[0:4], keepaliveBuf[0:4])
			binary.BigEndian.PutUint64(nonce[4:12], counter)

			encrypted := t.aead.Seal(keepaliveBuf[:12], nonce[:], nil, keepaliveBuf[:12])
			_, _ = t.udpConn.WriteToUDP(encrypted, target)
		}
	}
}

// 64-packet Sliding Window Anti-Replay algorithm
func (t *TunnelRunner) checkAndUpdateReplay(counter uint64) bool {
	t.mu.Lock()
	defer t.mu.Unlock()

	if counter == 0 {
		return false
	}
	if counter > t.rxMaxCounter {
		diff := counter - t.rxMaxCounter
		if diff >= 64 {
			t.replayWindow = 1
		} else {
			t.replayWindow = (t.replayWindow << diff) | 1
		}
		t.rxMaxCounter = counter
		return true
	}

	diff := t.rxMaxCounter - counter
	if diff >= 64 {
		return false // Packet is too old
	}

	mask := uint64(1) << diff
	if (t.replayWindow & mask) != 0 {
		return false // Duplicate packet!
	}

	t.replayWindow |= mask
	return true
}

func (t *TunnelRunner) GetTelemetry() TunnelInstance {
	t.mu.RLock()
	peerStr := ""
	if t.remoteAddr != nil {
		peerStr = t.remoteAddr.String()
	}
	lastHandshakeStr := "Never"
	if !t.lastHandshake.IsZero() {
		lastHandshakeStr = time.Since(t.lastHandshake).Truncate(time.Second).String() + " ago"
	}
	t.mu.RUnlock()

	inst := t.cfg
	inst.Running = !t.closed
	inst.ActivePeerEndpoint = peerStr
	inst.LastHandshake = lastHandshakeStr
	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)

	// Calculate compression ratio
	if inst.RawBytesOut > 0 && inst.WireBytesOut > 0 {
		ratio := float64(inst.RawBytesOut) / float64(inst.WireBytesOut)
		if ratio > 1.0 {
			savedPct := (1.0 - (float64(inst.WireBytesOut) / float64(inst.RawBytesOut))) * 100.0
			inst.CompressionRatio = fmt.Sprintf("%.2fx (%.1f%% saved)", ratio, savedPct)
		} else {
			inst.CompressionRatio = "1.00x (0% saved)"
		}
	} else {
		inst.CompressionRatio = "1.00x (0% saved)"
	}

	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 {
	if a == nil || b == nil {
		return a == b
	}
	return a.IP.Equal(b.IP) && a.Port == b.Port
}

func stringsEqual(a, b string) bool {
	return len(a) == len(b) && (a == b || (len(a) > 0 && (a[0]|0x20) == (b[0]|0x20)))
}