package aerotun

import (
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"io"
	"net/netip"
	"sync"
	"sync/atomic"
	"time"

	"github.com/flynn/noise"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/hkdf"
)

const (
	secureCookie        byte = 3
	secureHeaderSize         = 16
	secureVersion            = 3
	secureInit          byte = 1
	secureResponse      byte = 2
	secureRaw           byte = 16
	secureCompressed    byte = 17
	secureKeepalive     byte = 18
	secureRekeyAfter         = 90 * time.Second
	secureRejectAfter        = 180 * time.Second
	securePreviousGrace      = 30 * time.Second
	handshakeMaxSize         = 512
)

var errPacket = errors.New("invalid authenticated tunnel packet")
var errHandshakeRate = errors.New("handshake rate limited")

// A 1024-packet window permits bounded reordering using a fixed 128-byte bitmap.
type secureReplay struct {
	max  uint64
	bits [16]uint64
}

func (r *secureReplay) valid(n uint64) bool {
	if n == 0 {
		return false
	}
	if n > r.max {
		return true
	}
	if r.max-n >= 1024 {
		return false
	}
	pos := n % 1024
	return r.bits[pos/64]&(uint64(1)<<(pos%64)) == 0
}
func (r *secureReplay) accept(n uint64) bool {
	if !r.valid(n) {
		return false
	}
	if n > r.max {
		if n-r.max >= 1024 {
			clear(r.bits[:])
		} else {
			for missing := r.max + 1; missing < n; missing++ {
				pos := missing % 1024
				r.bits[pos/64] &^= uint64(1) << (pos % 64)
			}
		}
		r.max = n
	}
	pos := n % 1024
	r.bits[pos/64] |= uint64(1) << (pos % 64)
	return true
}

type secureSession struct {
	local, remote      uint32
	tx, rx             cipher.AEAD
	txPrefix, rxPrefix [16]byte
	counter            atomic.Uint64
	rxMu               sync.Mutex
	replay             secureReplay
	born               time.Time
	peerMTU            int
	compression        bool
}
type secureSessions struct {
	current, previous, pending *secureSession
	previousUntil              time.Time
}
type noiseAttempt struct {
	state             *noise.HandshakeState
	packet            []byte
	header            [16]byte
	started, lastSent time.Time
}
type secureProtocol struct {
	cookieSecret   [32]byte
	mu             sync.Mutex // Only handshakes/publication, never bulk encryption/decryption.
	cfg            TunnelInstance
	psk            []byte
	identity       noise.DHKey
	peer           []byte
	sessions       atomic.Pointer[secureSessions]
	attempt        *noiseAttempt
	lastRequest    time.Time
	cachedInit     [32]byte
	cachedResponse []byte
	cachedAt       time.Time
	handshakes     atomic.Uint64
}

func newSecureProtocol(cfg TunnelInstance) (*secureProtocol, error) {
	psk, err := decodeKey(cfg.PSK)
	if err != nil {
		return nil, err
	}
	priv, err := decodeKey(cfg.PrivateKey)
	if err != nil {
		return nil, err
	}
	pub, err := decodeKey(cfg.PublicKey)
	if err != nil {
		return nil, err
	}
	peer, err := decodeKey(cfg.PeerPublicKey)
	if err != nil {
		return nil, err
	}
	p := &secureProtocol{cfg: cfg, psk: psk, identity: noise.DHKey{Private: priv, Public: pub}, peer: peer}
	if _, err := rand.Read(p.cookieSecret[:]); err != nil {
		return nil, err
	}
	p.sessions.Store(&secureSessions{})
	return p, nil
}
func randomIndex() (uint32, error) {
	var b [4]byte
	for {
		if _, err := rand.Read(b[:]); err != nil {
			return 0, err
		}
		if n := binary.LittleEndian.Uint32(b[:]); n != 0 {
			return n, nil
		}
	}
}
func secureHeader(dst []byte, kind byte, index uint32, n uint64) {
	binary.BigEndian.PutUint16(dst, MagicHeader)
	dst[2] = 3
	dst[3] = kind
	binary.BigEndian.PutUint32(dst[4:8], index)
	binary.BigEndian.PutUint64(dst[8:16], n)
}
func secureMetadata(index uint32, mtu int, compression bool) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint32(b, index)
	binary.BigEndian.PutUint16(b[4:], uint16(mtu))
	if compression {
		b[6] = 1
	}
	return b
}
func parseSecureMetadata(b []byte) (index uint32, mtu int, comp bool, err error) {
	if len(b) != 8 || b[7] != 0 || b[6] > 1 {
		return 0, 0, false, errPacket
	}
	index = binary.BigEndian.Uint32(b)
	mtu = int(binary.BigEndian.Uint16(b[4:]))
	comp = b[6] == 1
	if index == 0 || mtu < 576 || mtu > 9000 {
		err = errPacket
	}
	return
}
func (p *secureProtocol) newHandshake(initiator bool, header []byte) (*noise.HandshakeState, error) {
	prologue := append([]byte("AeroTUN-v3/Noise-IKpsk2/XChaCha-data-v1\x00"), header...)
	cfg := noise.Config{CipherSuite: noise.NewCipherSuite(noise.DH25519, noise.CipherChaChaPoly, noise.HashSHA256), Pattern: noise.HandshakeIK, Initiator: initiator, Prologue: prologue, PresharedKey: p.psk, PresharedKeyPlacement: 2, StaticKeypair: p.identity}
	if initiator {
		cfg.PeerStatic = p.peer
	}
	return noise.NewHandshakeState(cfg)
}
func (p *secureProtocol) authenticateHandshake(b []byte) bool {
	if len(b) < 32 || len(b) > handshakeMaxSize {
		return false
	}
	mac := hmac.New(sha256.New, p.psk)
	mac.Write([]byte("AeroTUN-v3-handshake-MAC"))
	mac.Write(b[:len(b)-16])
	return hmac.Equal(mac.Sum(nil)[:16], b[len(b)-16:])
}
func (p *secureProtocol) finishHandshake(b []byte) []byte {
	mac := hmac.New(sha256.New, p.psk)
	mac.Write([]byte("AeroTUN-v3-handshake-MAC"))
	mac.Write(b)
	return append(b, mac.Sum(nil)[:16]...)
}
func (p *secureProtocol) initiation(mtu int, force bool) ([]byte, error) {
	if p.cfg.Mode != "client" {
		return nil, nil
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	now := time.Now()
	if p.attempt != nil && now.Sub(p.attempt.started) < 5*time.Second {
		if now.Sub(p.attempt.lastSent) < time.Second {
			return nil, nil
		}
		p.attempt.lastSent = now
		return p.attempt.packet, nil
	}
	if s := p.sessions.Load().current; s != nil && !force && now.Sub(s.born) < secureRekeyAfter && s.counter.Load() < 1<<60 {
		return nil, nil
	}
	index, err := randomIndex()
	if err != nil {
		return nil, err
	}
	var header [16]byte
	var nonce [8]byte
	if _, err = rand.Read(nonce[:]); err != nil {
		return nil, err
	}
	secureHeader(header[:], secureInit, index, binary.BigEndian.Uint64(nonce[:]))
	hs, err := p.newHandshake(true, header[:])
	if err != nil {
		return nil, err
	}
	b, _, _, err := hs.WriteMessage(append([]byte(nil), header[:]...), secureMetadata(index, mtu, p.cfg.Compression))
	if err != nil {
		return nil, err
	}
	b = append(b, make([]byte, 16)...) // stateless source-address cookie
	b = p.finishHandshake(b)
	p.attempt = &noiseAttempt{state: hs, packet: b, header: header, started: now, lastSent: now}
	return b, nil
}
func deriveDataState(c *noise.CipherState, binding []byte) (cipher.AEAD, [16]byte, error) {
	// Export Noise split keys only for domain-separated data-key derivation. Noise
	// CipherStates are never used for transport, so there are no competing nonces.
	key := c.UnsafeKey()
	var material [48]byte
	_, err := io.ReadFull(hkdf.New(sha256.New, key[:], binding, []byte("AeroTUN-v3-XChaCha20-Poly1305-data")), material[:])
	if err != nil {
		return nil, [16]byte{}, err
	}
	a, err := chacha20poly1305.NewX(material[:32])
	var prefix [16]byte
	copy(prefix[:], material[32:])
	clear(material[:])
	clear(key[:])
	return a, prefix, err
}
func makeSecureSession(local, remote uint32, mtu int, comp bool, tx, rx *noise.CipherState, binding []byte) (*secureSession, error) {
	s := &secureSession{local: local, remote: remote, peerMTU: mtu, compression: comp, born: time.Now()}
	var err error
	s.tx, s.txPrefix, err = deriveDataState(tx, binding)
	if err != nil {
		return nil, err
	}
	s.rx, s.rxPrefix, err = deriveDataState(rx, binding)
	return s, err
}

// A responder publishes only a pending receive session until authenticated data
// proves the peer also possesses the PSK and completed the fresh handshake.
func (p *secureProtocol) handleHandshake(b []byte, mtu int, source ...netip.AddrPort) (response []byte, peerMTU int, err error) {
	if len(b) < 32 || len(b) > handshakeMaxSize || binary.BigEndian.Uint16(b) != MagicHeader || b[2] != 3 || !p.authenticateHandshake(b) {
		return nil, 0, errPacket
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	now := time.Now()
	switch b[3] {
	case secureInit:
		if p.cfg.Mode != "server" {
			return nil, 0, errPacket
		}
		if len(b) < 48 {
			return nil, 0, errPacket
		}
		var address netip.AddrPort
		if len(source) > 0 {
			address = source[0]
		}
		current := p.addressCookie(b[:16], address, now.Unix()/30)
		previous := p.addressCookie(b[:16], address, now.Unix()/30-1)
		supplied := b[len(b)-32 : len(b)-16]
		if !hmac.Equal(current[:], supplied) && !hmac.Equal(previous[:], supplied) {
			out := append([]byte(nil), b[:16]...)
			out[3] = secureCookie
			out = append(out, current[:]...)
			return p.finishHandshake(out), 0, nil
		}

		fingerprint := sha256.Sum256(b)
		if now.Sub(p.lastRequest) < time.Second {
			return nil, 0, errHandshakeRate
		}
		p.lastRequest = now
		if fingerprint == p.cachedInit && now.Sub(p.cachedAt) < 5*time.Second {
			return p.cachedResponse, 0, nil
		}
		hs, e := p.newHandshake(false, b[:16])
		if e != nil {
			return nil, 0, e
		}
		meta, _, _, e := hs.ReadMessage(nil, b[16:len(b)-32])
		if e != nil || !hmac.Equal(hs.PeerStatic(), p.peer) {
			return nil, 0, errPacket
		}
		remote, remoteMTU, comp, e := parseSecureMetadata(meta)
		if e != nil || remote != binary.BigEndian.Uint32(b[4:8]) {
			return nil, 0, errPacket
		}
		local, e := randomIndex()
		if e != nil {
			return nil, 0, e
		}
		old := p.sessions.Load()
		for local == remote || (old.current != nil && local == old.current.local) || (old.previous != nil && local == old.previous.local) {
			local, e = randomIndex()
			if e != nil {
				return nil, 0, e
			}
		}
		var h [16]byte
		copy(h[:], b[:16])
		h[3] = secureResponse
		out, c1, c2, e := hs.WriteMessage(append([]byte(nil), h[:]...), secureMetadata(local, mtu, p.cfg.Compression))
		if e != nil {
			return nil, 0, e
		}
		if c1 == nil || c2 == nil {
			return nil, 0, errPacket
		}
		session, e := makeSecureSession(local, remote, remoteMTU, comp && p.cfg.Compression, c2, c1, hs.ChannelBinding())
		if e != nil {
			return nil, 0, e
		}
		p.sessions.Store(&secureSessions{current: old.current, previous: old.previous, previousUntil: old.previousUntil, pending: session})
		out = p.finishHandshake(out)
		p.cachedInit = fingerprint
		p.cachedResponse = out
		p.cachedAt = now
		return out, 0, nil
	case secureCookie:
		if p.cfg.Mode != "client" || len(b) != 48 || p.attempt == nil || now.Sub(p.attempt.started) > 5*time.Second || !hmac.Equal(b[4:16], p.attempt.header[4:16]) {
			return nil, 0, errPacket
		}
		packet := append([]byte(nil), p.attempt.packet...)
		copy(packet[len(packet)-32:len(packet)-16], b[16:32])
		packet = p.finishHandshake(packet[:len(packet)-16])
		p.attempt.packet = packet
		p.attempt.lastSent = now
		return packet, 0, nil
	case secureResponse:
		if p.cfg.Mode != "client" || p.attempt == nil || now.Sub(p.attempt.started) > 5*time.Second || !hmac.Equal(b[4:16], p.attempt.header[4:16]) {
			return nil, 0, errPacket
		}
		hs := p.attempt.state
		meta, c1, c2, e := hs.ReadMessage(nil, b[16:len(b)-16])
		if e != nil {
			p.attempt = nil
			return nil, 0, e
		}
		remote, remoteMTU, comp, e := parseSecureMetadata(meta)
		if e != nil {
			return nil, 0, e
		}
		if c1 == nil || c2 == nil {
			return nil, 0, errPacket
		}
		local := binary.BigEndian.Uint32(b[4:8])
		if local == remote {
			return nil, 0, errPacket
		}
		session, e := makeSecureSession(local, remote, remoteMTU, comp && p.cfg.Compression, c1, c2, hs.ChannelBinding())
		if e != nil {
			return nil, 0, e
		}
		old := p.sessions.Load()
		p.sessions.Store(&secureSessions{current: session, previous: old.current, previousUntil: now.Add(securePreviousGrace)})
		p.attempt = nil
		p.handshakes.Add(1)
		return nil, remoteMTU, nil
	}
	return nil, 0, errPacket
}
func (p *secureProtocol) sealInto(dst []byte, nonce *[24]byte, flags byte, payload []byte) ([]byte, error) {
	s := p.sessions.Load().current
	if s == nil || time.Since(s.born) > secureRejectAfter {
		return nil, errSession
	}
	kind := secureRaw
	switch flags {
	case 0:
	case FlagCompressedLZ4:
		if !s.compression {
			return nil, errPacket
		}
		kind = secureCompressed
	case FlagKeepalive:
		kind = secureKeepalive
	default:
		return nil, errPacket
	}
	if cap(dst) < 16+len(payload)+16 {
		return nil, errPacket
	}
	n := s.counter.Add(1)
	if n == 0 || n >= 1<<60 {
		return nil, errSession
	}
	dst = dst[:16]
	secureHeader(dst, kind, s.remote, n)
	copy(nonce[:16], s.txPrefix[:])
	binary.BigEndian.PutUint64(nonce[16:], n)
	return s.tx.Seal(dst, nonce[:], payload, dst), nil
}
func (p *secureProtocol) compressionEnabled() bool {
	s := p.sessions.Load().current
	return s != nil && s.compression
}
func (p *secureProtocol) openInto(dst []byte, nonce *[24]byte, b []byte) (payload []byte, flags byte, peerMTU int, err error) {
	if len(b) < 32 || cap(dst) < len(b)-32 || binary.BigEndian.Uint16(b) != MagicHeader || b[2] != 3 {
		return nil, 0, 0, errPacket
	}
	switch b[3] {
	case secureRaw:
	case secureCompressed:
		flags = FlagCompressedLZ4
	case secureKeepalive:
		flags = FlagKeepalive
	default:
		return nil, 0, 0, errPacket
	}
	slots := p.sessions.Load()
	index := binary.BigEndian.Uint32(b[4:8])
	var s *secureSession
	for _, candidate := range []*secureSession{slots.current, slots.pending, slots.previous} {
		if candidate != nil && candidate.local == index {
			s = candidate
			break
		}
	}
	if s == nil || time.Since(s.born) > secureRejectAfter || (s == slots.previous && time.Now().After(slots.previousUntil)) || (flags == FlagCompressedLZ4 && !s.compression) {
		return nil, 0, 0, errSession
	}
	n := binary.BigEndian.Uint64(b[8:16])
	if n >= 1<<60 {
		return nil, 0, 0, errPacket
	}
	s.rxMu.Lock()
	valid := s.replay.valid(n)
	s.rxMu.Unlock()
	if !valid {
		return nil, 0, 0, errReplay
	}
	copy(nonce[:16], s.rxPrefix[:])
	binary.BigEndian.PutUint64(nonce[16:], n)
	payload, err = s.rx.Open(dst[:0], nonce[:], b[16:], b[:16])
	if err != nil {
		return nil, 0, 0, err
	}
	if flags == FlagKeepalive && (len(payload) != 2 || binary.BigEndian.Uint16(payload) < 576 || binary.BigEndian.Uint16(payload) > 9000) {
		return nil, 0, 0, errPacket
	}
	s.rxMu.Lock()
	valid = s.replay.accept(n)
	s.rxMu.Unlock()
	if !valid {
		return nil, 0, 0, errReplay
	}
	if s == slots.pending {
		p.mu.Lock()
		nowSlots := p.sessions.Load()
		if nowSlots.pending == s {
			p.sessions.Store(&secureSessions{current: s, previous: nowSlots.current, previousUntil: time.Now().Add(securePreviousGrace)})
			p.handshakes.Add(1)
		}
		p.mu.Unlock()
		peerMTU = s.peerMTU
	}
	return
}

// Retain at most the current session, one bounded-grace predecessor, and one
// short-lived pending handshake. Expired keys are no longer reachable by the
// protocol; Go does not guarantee physical erasure of discarded heap objects.
func (p *secureProtocol) expire() {
	p.mu.Lock()
	defer p.mu.Unlock()
	old := p.sessions.Load()
	next := *old
	now := time.Now()
	changed := false
	if next.current != nil && now.Sub(next.current.born) > secureRejectAfter {
		next.current = nil
		changed = true
	}
	if next.previous != nil && now.After(next.previousUntil) {
		next.previous = nil
		changed = true
	}
	if next.pending != nil && now.Sub(next.pending.born) > 5*time.Second {
		next.pending = nil
		changed = true
	}
	if changed {
		p.sessions.Store(&next)
	}
	if p.attempt != nil && now.Sub(p.attempt.started) > 5*time.Second {
		p.attempt = nil
	}
	if len(p.cachedResponse) > 0 && now.Sub(p.cachedAt) > 5*time.Second {
		p.cachedResponse = nil
		p.cachedInit = [32]byte{}
	}
}

// A short-lived stateless cookie binds expensive DH work to a reachable UDP
// source. A captured initiation cannot be moved to another source without a new
// PSK MAC. Current/previous 30-second buckets tolerate packet delay.
func (p *secureProtocol) addressCookie(header []byte, address netip.AddrPort, bucket int64) (out [16]byte) {
	h := hmac.New(sha256.New, p.cookieSecret[:])
	h.Write([]byte("AeroTUN-v3-source-cookie"))
	h.Write(header[4:16])
	var encoded [64]byte
	b := address.AppendTo(encoded[:0])
	h.Write(b)
	var clock [8]byte
	binary.BigEndian.PutUint64(clock[:], uint64(bucket))
	h.Write(clock[:])
	copy(out[:], h.Sum(nil)[:16])
	return
}