package gateway

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"net"
	"sync"
	"time"
)

const (
	apolloENetChannels        = 48
	apolloENetMaximumPacket   = 4096
	apolloENetMaximumPayload  = 2048
	apolloENetMaximumCommands = 32
	apolloENetMaximumPending  = 128
	apolloENetMaximumReorder  = 64
	// These are the pinned ENet fork defaults. The connection remains subject
	// to the stricter ten-second no-receive peer deadline below.
	apolloENetTimeoutLimit    = 32
	apolloENetTimeoutMinimum  = 5 * time.Second
	apolloENetTimeoutMaximum  = 30 * time.Second
	apolloENetPeerTimeout     = 10 * time.Second
	apolloENetPeerIDMask      = 0x0fff
	apolloENetSentTimeFlag    = 0x8000
	apolloENetCompressedFlag  = 0x4000
	apolloENetSessionMask     = 0x3000
	apolloENetSessionShift    = 12
	apolloENetCommandMask     = 0x0f
	apolloENetAcknowledged    = 0x80
	apolloENetUnsequenced     = 0x40
	apolloENetConnect         = 2
	apolloENetVerifyConnect   = 3
	apolloENetDisconnect      = 4
	apolloENetPing            = 5
	apolloENetSendReliable    = 6
	apolloENetSendUnsequenced = 9
	apolloENetBandwidthLimit  = 10
	apolloENetThrottleConfig  = 11
)

var errApolloENet = errors.New("apollo ENet malformed")

type apolloENetState uint8

const (
	apolloENetConnecting apolloENetState = iota
	apolloENetConnected
	apolloENetDisconnecting
	apolloENetClosed
)

type apolloENetChannel struct {
	nextOutgoing uint16
	lastIncoming uint16
	hasIncoming  bool
	incoming     map[uint16][]byte
}

type apolloENetPendingKey struct {
	channel  uint8
	sequence uint16
}

type apolloENetPending struct {
	packet    []byte
	firstSent time.Time
	sentTime  time.Time
	timeout   time.Duration
	attempts  uint8
}

// apolloENetPeer is deliberately scoped to the Apollo adapter. It implements
// one negotiated ENet peer over one connected UDP socket and exposes no
// reusable transport abstraction.
type apolloENetPeer struct {
	conn *net.UDPConn
	now  func() time.Time

	mu              sync.Mutex
	state           apolloENetState
	peerID          uint16
	inboundSession  uint8
	outboundSession uint8
	connectID       uint32
	channels        [apolloENetChannels]apolloENetChannel
	pending         map[apolloENetPendingKey]*apolloENetPending
	unsequenced     uint16
	rtt             time.Duration
	variance        time.Duration
	reliableSent    uint64
	retransmits     uint64
	lastReceive     time.Time
	lastSend        time.Time
	lastPing        time.Time
	disconnectAck   chan struct{}
	disconnectSeq   uint16
	onPayload       func(uint8, bool, []byte)
	onDisconnect    func(error)
	done            chan struct{}
	closeOnce       sync.Once
}

func newApolloENetPeer(conn *net.UDPConn, now func() time.Time) (*apolloENetPeer, error) {
	if conn == nil {
		return nil, ErrProviderMalformed
	}
	if now == nil {
		now = time.Now
	}
	return &apolloENetPeer{
		conn: conn, now: now, state: apolloENetConnecting, pending: make(map[apolloENetPendingKey]*apolloENetPending),
		rtt: 500 * time.Millisecond, variance: time.Millisecond, done: make(chan struct{}),
	}, nil
}

func (p *apolloENetPeer) Connect(ctx context.Context, connectData uint32) error {
	if p == nil || connectData == 0 {
		return ErrProviderMalformed
	}
	var id [4]byte
	if _, err := rand.Read(id[:]); err != nil {
		return err
	}
	p.mu.Lock()
	if p.state != apolloENetConnecting {
		p.mu.Unlock()
		return ErrProviderMalformed
	}
	p.connectID = binary.BigEndian.Uint32(id[:])
	now := p.now()
	packet := apolloENetConnectPacket(now, p.connectID, connectData)
	p.pending[apolloENetPendingKey{channel: 0xff, sequence: 1}] = &apolloENetPending{packet: append([]byte(nil), packet...), firstSent: now, sentTime: now, timeout: apolloENetRetransmitTimeout(p.rtt, p.variance, 1), attempts: 1}
	err := p.writeLocked(packet)
	p.mu.Unlock()
	if err != nil {
		p.close(err)
		return err
	}
	for {
		if err := p.readOnce(ctx); err != nil {
			p.close(err)
			return err
		}
		p.mu.Lock()
		connected := p.state == apolloENetConnected
		p.mu.Unlock()
		if connected {
			go p.run()
			return nil
		}
	}
}

func apolloENetConnectPacket(now time.Time, connectID, data uint32) []byte {
	packet := make([]byte, 52)
	apolloENetHeader(packet, apolloENetPeerIDMask, 0, now)
	packet[4] = apolloENetConnect | apolloENetAcknowledged
	packet[5] = 0xff
	binary.BigEndian.PutUint16(packet[6:8], 1)
	binary.BigEndian.PutUint16(packet[8:10], 0)
	packet[10], packet[11] = 0xff, 0xff
	binary.BigEndian.PutUint32(packet[12:16], 1400)
	binary.BigEndian.PutUint32(packet[16:20], 32768)
	binary.BigEndian.PutUint32(packet[20:24], apolloENetChannels)
	binary.BigEndian.PutUint32(packet[24:28], 0)
	binary.BigEndian.PutUint32(packet[28:32], 0)
	binary.BigEndian.PutUint32(packet[32:36], 5000)
	binary.BigEndian.PutUint32(packet[36:40], 2)
	binary.BigEndian.PutUint32(packet[40:44], 2)
	binary.BigEndian.PutUint32(packet[44:48], connectID)
	binary.BigEndian.PutUint32(packet[48:52], data)
	return packet
}

func apolloENetHeader(packet []byte, peerID uint16, session uint8, now time.Time) {
	value := peerID&apolloENetPeerIDMask | (uint16(session&3) << apolloENetSessionShift) | apolloENetSentTimeFlag
	binary.BigEndian.PutUint16(packet[:2], value)
	binary.BigEndian.PutUint16(packet[2:4], uint16(now.UnixMilli()))
}

func (p *apolloENetPeer) run() {
	ticker := time.NewTicker(25 * time.Millisecond)
	defer ticker.Stop()
	for {
		select {
		case <-p.done:
			return
		case <-ticker.C:
			if err := p.maintain(); err != nil {
				p.close(err)
				return
			}
		default:
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
			err := p.readOnce(ctx)
			cancel()
			if err != nil && !errors.Is(err, context.DeadlineExceeded) && !isApolloENetTimeout(err) {
				p.close(err)
				return
			}
		}
	}
}

func (p *apolloENetPeer) readOnce(ctx context.Context) error {
	if p == nil || p.conn == nil {
		return ErrProviderDisconnected
	}
	deadline := time.Now().Add(100 * time.Millisecond)
	if contextDeadline, ok := ctx.Deadline(); ok && contextDeadline.Before(deadline) {
		deadline = contextDeadline
	}
	if err := p.conn.SetReadDeadline(deadline); err != nil {
		return err
	}
	buffer := make([]byte, apolloENetMaximumPacket+1)
	count, err := p.conn.Read(buffer)
	if err != nil {
		if networkErr, ok := err.(net.Error); ok && networkErr.Timeout() {
			return context.DeadlineExceeded
		}
		return err
	}
	if count < 4 || count > apolloENetMaximumPacket {
		return errApolloENet
	}
	return p.handleDatagram(buffer[:count])
}

func isApolloENetTimeout(err error) bool {
	return errors.Is(err, context.DeadlineExceeded)
}

func (p *apolloENetPeer) handleDatagram(packet []byte) error {
	if len(packet) < 4 || len(packet) > apolloENetMaximumPacket {
		return errApolloENet
	}
	header := binary.BigEndian.Uint16(packet[:2])
	if header&apolloENetCompressedFlag != 0 {
		return errApolloENet
	}
	peerID := header & apolloENetPeerIDMask
	session := uint8((header & apolloENetSessionMask) >> apolloENetSessionShift)
	offset := 2
	sentTime := uint16(0)
	if header&apolloENetSentTimeFlag != 0 {
		if len(packet) < 4 {
			return errApolloENet
		}
		sentTime = binary.BigEndian.Uint16(packet[2:4])
		offset = 4
	}
	p.mu.Lock()
	state := p.state
	if state == apolloENetClosed || (state == apolloENetConnected && (peerID != p.peerID || session != p.inboundSession)) {
		p.mu.Unlock()
		return errApolloENet
	}
	p.lastReceive = p.now()
	p.mu.Unlock()
	commands := 0
	for offset < len(packet) {
		commands++
		if commands > apolloENetMaximumCommands || len(packet)-offset < 4 {
			return errApolloENet
		}
		command := packet[offset] & apolloENetCommandMask
		flags := packet[offset]
		channel := packet[offset+1]
		sequence := binary.BigEndian.Uint16(packet[offset+2 : offset+4])
		consumed, err := p.handleCommand(command, flags, channel, sequence, sentTime, packet[offset:])
		if err != nil || consumed < 4 || consumed > len(packet)-offset {
			return errApolloENet
		}
		offset += consumed
	}
	return nil
}

func (p *apolloENetPeer) handleCommand(command, flags, channel uint8, sequence, sentTime uint16, data []byte) (int, error) {
	switch command {
	case 1:
		if len(data) < 8 {
			return 0, errApolloENet
		}
		return 8, p.acknowledge(channel, binary.BigEndian.Uint16(data[4:6]), binary.BigEndian.Uint16(data[6:8]))
	case apolloENetVerifyConnect:
		if len(data) < 44 {
			return 0, errApolloENet
		}
		return 44, p.verifyConnect(sequence, sentTime, data[:44])
	case apolloENetDisconnect:
		if len(data) < 8 {
			return 0, errApolloENet
		}
		if flags&apolloENetAcknowledged != 0 {
			if err := p.sendAcknowledge(channel, sequence, sentTime); err != nil {
				return 0, err
			}
		}
		return 8, ErrProviderDisconnected
	case apolloENetPing:
		if flags&apolloENetAcknowledged != 0 {
			if err := p.sendAcknowledge(channel, sequence, sentTime); err != nil {
				return 0, err
			}
		}
		return 4, nil
	case apolloENetSendReliable:
		if len(data) < 6 {
			return 0, errApolloENet
		}
		length := int(binary.BigEndian.Uint16(data[4:6]))
		if length > apolloENetMaximumPayload || len(data) < 6+length {
			return 0, errApolloENet
		}
		if flags&apolloENetAcknowledged == 0 || channel >= apolloENetChannels {
			return 0, errApolloENet
		}
		if err := p.sendAcknowledge(channel, sequence, sentTime); err != nil {
			return 0, err
		}
		deliver, err := p.acceptReliable(channel, sequence, data[6:6+length])
		if err != nil {
			return 0, err
		}
		for _, message := range deliver {
			p.deliver(channel, true, message)
		}
		return 6 + length, nil
	case apolloENetSendUnsequenced:
		if len(data) < 8 {
			return 0, errApolloENet
		}
		length := int(binary.BigEndian.Uint16(data[6:8]))
		if flags&apolloENetUnsequenced == 0 || channel >= apolloENetChannels || length > apolloENetMaximumPayload || len(data) < 8+length {
			return 0, errApolloENet
		}
		p.deliver(channel, false, data[8:8+length])
		return 8 + length, nil
	case apolloENetBandwidthLimit:
		if len(data) < 12 {
			return 0, errApolloENet
		}
		return 12, nil
	case apolloENetThrottleConfig:
		if len(data) < 16 {
			return 0, errApolloENet
		}
		return 16, nil
	case apolloENetConnect, 7, 8, 12:
		return 0, errApolloENet
	default:
		return 0, errApolloENet
	}
}

func (p *apolloENetPeer) verifyConnect(sequence, sentTime uint16, data []byte) error {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.state != apolloENetConnecting || binary.BigEndian.Uint32(data[40:44]) != p.connectID || binary.BigEndian.Uint32(data[16:20]) != apolloENetChannels {
		return errApolloENet
	}
	p.peerID = binary.BigEndian.Uint16(data[4:6])
	p.inboundSession = data[6]
	p.outboundSession = data[7]
	p.state = apolloENetConnected
	delete(p.pending, apolloENetPendingKey{channel: 0xff, sequence: 1})
	return p.sendAcknowledgeLocked(0xff, sequence, sentTime)
}

func (p *apolloENetPeer) acknowledge(channel uint8, sequence, sentTime uint16) error {
	p.mu.Lock()
	if p.state == apolloENetDisconnecting && channel == 0xff && sequence == p.disconnectSeq && p.disconnectAck != nil {
		close(p.disconnectAck)
		p.disconnectAck = nil
		p.mu.Unlock()
		return nil
	}
	key := apolloENetPendingKey{channel: channel, sequence: sequence}
	pending, ok := p.pending[key]
	if !ok {
		p.mu.Unlock()
		return nil
	}
	delete(p.pending, key)
	measured := p.now().Sub(pending.sentTime)
	if measured < 0 || measured > 30*time.Second {
		p.mu.Unlock()
		return errApolloENet
	}
	delta := durationAbs(p.rtt - measured)
	p.variance += (delta - p.variance) / 4
	p.rtt += (measured - p.rtt) / 8
	p.mu.Unlock()
	_ = sentTime
	return nil
}

func durationAbs(value time.Duration) time.Duration {
	if value < 0 {
		return -value
	}
	return value
}

func (p *apolloENetPeer) acceptReliable(channel uint8, sequence uint16, payload []byte) ([][]byte, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	state := &p.channels[channel]
	if !state.hasIncoming {
		state.hasIncoming = true
		state.lastIncoming = sequence
		return [][]byte{append([]byte(nil), payload...)}, nil
	}
	if sequence == state.lastIncoming+1 {
		state.lastIncoming = sequence
		deliver := [][]byte{append([]byte(nil), payload...)}
		for {
			next := state.lastIncoming + 1
			queued, ok := state.incoming[next]
			if !ok {
				return deliver, nil
			}
			delete(state.incoming, next)
			state.lastIncoming = next
			deliver = append(deliver, queued)
		}
	}
	if apolloENetSequenceGreater(sequence, state.lastIncoming) {
		if uint16(sequence-state.lastIncoming) > 1024 || len(state.incoming) >= apolloENetMaximumReorder {
			return nil, errApolloENet
		}
		if state.incoming == nil {
			state.incoming = make(map[uint16][]byte)
		}
		if _, duplicate := state.incoming[sequence]; !duplicate {
			state.incoming[sequence] = append([]byte(nil), payload...)
		}
	}
	return nil, nil
}

func apolloENetSequenceGreater(first, second uint16) bool {
	return (first > second && first-second <= 32768) || (first < second && second-first > 32768)
}

func (p *apolloENetPeer) deliver(channel uint8, reliable bool, payload []byte) {
	p.mu.Lock()
	callback := p.onPayload
	p.mu.Unlock()
	if callback != nil {
		callback(channel, reliable, append([]byte(nil), payload...))
	}
}

func (p *apolloENetPeer) SendReliable(channel uint8, payload []byte) error {
	if channel >= apolloENetChannels || len(payload) == 0 || len(payload) > apolloENetMaximumPayload {
		return ErrProviderMalformed
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.state != apolloENetConnected || len(p.pending) >= apolloENetMaximumPending {
		return ErrProviderDisconnected
	}
	state := &p.channels[channel]
	state.nextOutgoing++
	if state.nextOutgoing == 0 {
		state.nextOutgoing++
	}
	now := p.now()
	packet := make([]byte, 10+len(payload))
	apolloENetHeader(packet, p.peerID, p.outboundSession, now)
	packet[4] = apolloENetSendReliable | apolloENetAcknowledged
	packet[5] = channel
	binary.BigEndian.PutUint16(packet[6:8], state.nextOutgoing)
	binary.BigEndian.PutUint16(packet[8:10], uint16(len(payload)))
	copy(packet[10:], payload)
	key := apolloENetPendingKey{channel: channel, sequence: state.nextOutgoing}
	p.pending[key] = &apolloENetPending{packet: append([]byte(nil), packet...), firstSent: now, sentTime: now, timeout: apolloENetRetransmitTimeout(p.rtt, p.variance, 1), attempts: 1}
	if err := p.writeLocked(packet); err != nil {
		delete(p.pending, key)
		return err
	}
	p.reliableSent++
	return nil
}

func (p *apolloENetPeer) SendUnsequenced(channel uint8, payload []byte) error {
	if channel >= apolloENetChannels || len(payload) == 0 || len(payload) > apolloENetMaximumPayload {
		return ErrProviderMalformed
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.state != apolloENetConnected {
		return ErrProviderDisconnected
	}
	p.unsequenced++
	packet := make([]byte, 12+len(payload))
	apolloENetHeader(packet, p.peerID, p.outboundSession, p.now())
	packet[4] = apolloENetSendUnsequenced | apolloENetUnsequenced
	packet[5] = channel
	binary.BigEndian.PutUint16(packet[6:8], 0)
	binary.BigEndian.PutUint16(packet[8:10], p.unsequenced)
	binary.BigEndian.PutUint16(packet[10:12], uint16(len(payload)))
	copy(packet[12:], payload)
	return p.writeLocked(packet)
}

func (p *apolloENetPeer) sendAcknowledge(channel uint8, sequence, sentTime uint16) error {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.sendAcknowledgeLocked(channel, sequence, sentTime)
}

func (p *apolloENetPeer) sendAcknowledgeLocked(channel uint8, sequence, sentTime uint16) error {
	if p.state == apolloENetClosed {
		return ErrProviderDisconnected
	}
	packet := make([]byte, 12)
	peerID := p.peerID
	session := p.outboundSession
	if p.state == apolloENetConnecting {
		peerID, session = apolloENetPeerIDMask, 0
	}
	apolloENetHeader(packet, peerID, session, p.now())
	packet[4] = 1
	packet[5] = channel
	binary.BigEndian.PutUint16(packet[6:8], 0)
	binary.BigEndian.PutUint16(packet[8:10], sequence)
	binary.BigEndian.PutUint16(packet[10:12], sentTime)
	return p.writeLocked(packet)
}

func (p *apolloENetPeer) maintain() error {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.state != apolloENetConnected && p.state != apolloENetDisconnecting {
		return nil
	}
	now := p.now()
	if !p.lastReceive.IsZero() && now.Sub(p.lastReceive) > apolloENetPeerTimeout {
		return ErrProviderTimeout
	}
	for key, pending := range p.pending {
		if now.Sub(pending.sentTime) < pending.timeout {
			continue
		}
		if now.Sub(pending.firstSent) >= apolloENetTimeoutMaximum || (apolloENetExceededTimeoutLimit(pending.attempts) && now.Sub(pending.firstSent) >= apolloENetTimeoutMinimum) {
			return ErrProviderTimeout
		}
		pending.timeout = apolloENetRetransmitTimeout(p.rtt, p.variance, pending.attempts)
		pending.attempts++
		p.retransmits++
		pending.sentTime = now
		binary.BigEndian.PutUint16(pending.packet[2:4], uint16(now.UnixMilli()))
		if err := p.writeLocked(pending.packet); err != nil {
			return err
		}
		p.pending[key] = pending
	}
	if now.Sub(p.lastPing) >= 500*time.Millisecond {
		p.lastPing = now
		return p.sendPingLocked()
	}
	return nil
}

func apolloENetRetransmitTimeout(rtt, variance time.Duration, attempts uint8) time.Duration {
	if rtt < time.Millisecond {
		rtt = time.Millisecond
	}
	if variance < time.Millisecond {
		variance = time.Millisecond
	}
	base := rtt + minDuration(rtt, 4*variance)
	if base > apolloENetTimeoutMaximum/5 {
		base = apolloENetTimeoutMaximum / 5
	}
	if attempts == 0 {
		attempts = 1
	}
	if attempts > apolloENetTimeoutLimit {
		attempts = apolloENetTimeoutLimit
	}
	return base * time.Duration(attempts)
}

func apolloENetExceededTimeoutLimit(attempts uint8) bool {
	return attempts >= 6 // 1 << (attempts - 1) reaches the fork's limit of 32.
}

func minDuration(first, second time.Duration) time.Duration {
	if first < second {
		return first
	}
	return second
}

func (p *apolloENetPeer) sendPingLocked() error {
	if len(p.pending) >= apolloENetMaximumPending {
		return ErrProviderTimeout
	}
	state := &p.channels[0]
	state.nextOutgoing++
	if state.nextOutgoing == 0 {
		state.nextOutgoing++
	}
	now := p.now()
	packet := make([]byte, 8)
	apolloENetHeader(packet, p.peerID, p.outboundSession, now)
	packet[4] = apolloENetPing | apolloENetAcknowledged
	packet[5] = 0
	binary.BigEndian.PutUint16(packet[6:8], state.nextOutgoing)
	p.pending[apolloENetPendingKey{channel: 0, sequence: state.nextOutgoing}] = &apolloENetPending{packet: append([]byte(nil), packet...), firstSent: now, sentTime: now, timeout: apolloENetRetransmitTimeout(p.rtt, p.variance, 1), attempts: 1}
	if err := p.writeLocked(packet); err != nil {
		delete(p.pending, apolloENetPendingKey{channel: 0, sequence: state.nextOutgoing})
		return err
	}
	p.reliableSent++
	return nil
}

func (p *apolloENetPeer) telemetry() ProviderTelemetry {
	if p == nil {
		return ProviderTelemetry{}
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	return ProviderTelemetry{
		ControlRTT: p.rtt, ControlJitter: p.variance, ReliableSent: p.reliableSent,
		ReliableRetransmits: p.retransmits, PendingReliable: uint64(len(p.pending)),
	}
}

func (p *apolloENetPeer) writeLocked(packet []byte) error {
	if len(packet) < 4 || len(packet) > apolloENetMaximumPacket {
		return ErrProviderMalformed
	}
	count, err := p.conn.Write(packet)
	if err != nil || count != len(packet) {
		if err != nil {
			return err
		}
		return ErrProviderDisconnected
	}
	p.lastSend = p.now()
	return nil
}

func (p *apolloENetPeer) Disconnect(ctx context.Context) error {
	if p == nil {
		return nil
	}
	p.mu.Lock()
	if p.state == apolloENetClosed {
		p.mu.Unlock()
		return nil
	}
	p.state = apolloENetDisconnecting
	ack := make(chan struct{})
	p.disconnectAck, p.disconnectSeq = ack, 1
	packet := make([]byte, 12)
	apolloENetHeader(packet, p.peerID, p.outboundSession, p.now())
	packet[4] = apolloENetDisconnect | apolloENetAcknowledged
	packet[5] = 0xff
	binary.BigEndian.PutUint16(packet[6:8], 1)
	if err := p.writeLocked(packet); err != nil {
		p.mu.Unlock()
		p.close(err)
		return err
	}
	p.mu.Unlock()
	deadline := time.NewTimer(2 * time.Second)
	defer deadline.Stop()
	select {
	case <-ack:
		p.close(nil)
		return nil
	case <-p.done:
		return nil
	case <-ctx.Done():
		p.close(ctx.Err())
		return ctx.Err()
	case <-deadline.C:
		p.close(ErrProviderTimeout)
		return ErrProviderTimeout
	}
}

func (p *apolloENetPeer) close(err error) {
	if p == nil {
		return
	}
	p.closeOnce.Do(func() {
		p.mu.Lock()
		p.state = apolloENetClosed
		callback := p.onDisconnect
		p.mu.Unlock()
		close(p.done)
		_ = p.conn.Close()
		if callback != nil && err != nil {
			callback(err)
		}
	})
}