// Package kcp is a Go port of the mhy-kcp Rust crate used by the original
// GaC-hkrpg-ps gameserver (kcp/src/kcp.rs). It implements the mihoyo KCP wire
// variant: a standard KCP segment extended with a 4-byte token after conv,
// giving a 28-byte header (KCPOverhead).
//
// Wire format (all little-endian):
//
//	conv(4) token(4) cmd(1) frg(1) wnd(2) ts(4) sn(4) una(4) len(4) data...
package kcp

import (
	"encoding/binary"
	"errors"
	"fmt"
)

// Protocol constants (mirrored from the Rust crate).
const (
	RtoNdl = 20  // no delay min rto
	RtoMin = 100 // normal min rto
	RtoDef = 200
	RtoMax = 60000

	CmdPush = 81 // cmd: push data
	CmdAck  = 82 // cmd: ack
	CmdWask = 83 // cmd: window probe (ask)
	CmdWins = 84 // cmd: window size (tell)

	AskSend = 1 // need to send CmdWask
	AskTell = 2 // need to send CmdWins

	WndSnd = 32
	WndRcv = 258 // must >= max fragment size

	MtuDef = 1400

	// Overhead is the KCP header size (conv+token+cmd+frg+wnd+ts+sn+una+len).
	Overhead = 28

	Interval = 100
	DeadLink = 20

	ThreshInit = 2
	ThreshMin  = 2

	ProbeInit  = 7000   // 7 secs to probe window size
	ProbeLimit = 120000 // up to 120 secs to probe window

	FastackLimit = 5 // max times to trigger fastack
)

// Sentinel errors mirroring kcp/src/error.rs.
var (
	ErrRecvQueueEmpty     = errors.New("recv queue is empty")
	ErrExpectingFragment  = errors.New("expecting fragment")
	ErrUserBufTooSmall    = errors.New("user buffer is too small")
	ErrUserBufTooBig      = errors.New("user buffer is too big")
	ErrNeedUpdate         = errors.New("update must be called at least once")
	ErrInvalidMtu         = errors.New("invalid mtu")
	ErrUnsupportedCmd     = errors.New("unsupported cmd")
	ErrConvInconsistent   = errors.New("conv inconsistent")
	ErrTokenMismatch      = errors.New("token mismatch")
	ErrInvalidSegmentSize = errors.New("invalid segment size")
	ErrInvalidSegmentData = errors.New("invalid segment data size")
)

// GetConv reads the conv field from a raw KCP packet.
func GetConv(buf []byte) uint32 {
	return binary.LittleEndian.Uint32(buf)
}

// GetSn reads the 32-bit field at the standard KCP sn offset (12).
// Ported as-is from the Rust crate (unused by the server).
func GetSn(buf []byte) uint32 {
	return binary.LittleEndian.Uint32(buf[12:])
}

func timeDiff(later, earlier uint32) int32 {
	return int32(later - earlier)
}

func bound(lower, v, upper uint32) uint32 {
	if v < lower {
		v = lower
	}
	if v > upper {
		v = upper
	}
	return v
}

type segment struct {
	conv  uint32
	token uint32
	cmd   uint8
	frg   uint8
	wnd   uint16
	ts    uint32
	sn    uint32
	una   uint32

	resendts uint32
	rto      uint32
	fastack  uint32
	xmit     uint32
	data     []byte
}

func (s *segment) encode(buf []byte) int {
	binary.LittleEndian.PutUint32(buf[0:], s.conv)
	binary.LittleEndian.PutUint32(buf[4:], s.token)
	buf[8] = s.cmd
	buf[9] = s.frg
	binary.LittleEndian.PutUint16(buf[10:], s.wnd)
	binary.LittleEndian.PutUint32(buf[12:], s.ts)
	binary.LittleEndian.PutUint32(buf[16:], s.sn)
	binary.LittleEndian.PutUint32(buf[20:], s.una)
	binary.LittleEndian.PutUint32(buf[24:], uint32(len(s.data)))
	copy(buf[Overhead:], s.data)
	return Overhead + len(s.data)
}

// Output receives fully assembled KCP packets ready for UDP transmission.
type Output func(data []byte) error

// KCP is the protocol control block.
type KCP struct {
	conv  uint32
	token uint32

	mtu   int
	mss   int
	state int32

	sndUna uint32
	sndNxt uint32
	rcvNxt uint32

	ssthresh uint16

	rxRttval uint32
	rxSrtt   uint32
	rxRto    uint32
	rxMinrto uint32

	sndWnd uint16
	rcvWnd uint16
	rmtWnd uint16
	cwnd   uint16
	probe  uint32

	current  uint32
	interval uint32
	tsFlush  uint32
	xmit     uint32

	nodelay   bool
	updated   bool
	tsProbe   uint32
	probeWait uint32

	deadLink uint32
	incr     int

	sndQueue []segment
	rcvQueue []segment
	sndBuf   []segment
	rcvBuf   []segment

	acklist    [][2]uint32
	buf        []byte
	fastresend uint32
	fastlimit  uint32
	nocwnd     bool
	stream     bool

	inputConv bool

	output Output
}

// New creates a KCP control object.
// conv must be equal on both endpoints; token authenticates the conversation;
// stream enables stream mode; output is called to emit packets.
func New(conv, token uint32, stream bool, output Output) *KCP {
	return &KCP{
		conv:      conv,
		token:     token,
		sndWnd:    WndSnd,
		rcvWnd:    WndRcv,
		rmtWnd:    WndRcv,
		mtu:       MtuDef,
		mss:       MtuDef - Overhead,
		stream:    stream,
		buf:       make([]byte, 0, (MtuDef+Overhead)*3),
		rxRto:     RtoDef,
		rxMinrto:  RtoMin,
		interval:  Interval,
		tsFlush:   Interval,
		ssthresh:  ThreshInit,
		fastlimit: FastackLimit,
		deadLink:  DeadLink,
		output:    output,
	}
}

// Conv returns the current conversation id.
func (k *KCP) Conv() uint32 { return k.conv }

// State returns the connection state (-1 means dead link).
func (k *KCP) State() int32 { return k.state }

// IsDeadLink reports whether resends exceeded the dead link limit.
func (k *KCP) IsDeadLink() bool { return k.state != 0 }

func (k *KCP) moveBuf() {
	for len(k.rcvBuf) > 0 {
		seg := &k.rcvBuf[0]
		if seg.sn == k.rcvNxt && len(k.rcvQueue) < int(k.rcvWnd) {
			k.rcvNxt++
		} else {
			break
		}
		k.rcvQueue = append(k.rcvQueue, k.rcvBuf[0])
		k.rcvBuf = k.rcvBuf[1:]
	}
}

// Recv receives merged (fragment-assembled) data into buf.
// Returns ErrRecvQueueEmpty when no complete message is pending.
func (k *KCP) Recv(buf []byte) (int, error) {
	if len(k.rcvQueue) == 0 {
		return 0, ErrRecvQueueEmpty
	}

	peekSize, err := k.peekSize()
	if err != nil {
		return 0, err
	}
	if peekSize > len(buf) {
		return 0, ErrUserBufTooSmall
	}

	recover := len(k.rcvQueue) >= int(k.rcvWnd)

	written := 0
	for len(k.rcvQueue) > 0 {
		seg := k.rcvQueue[0]
		k.rcvQueue = k.rcvQueue[1:]
		copy(buf[written:], seg.data)
		written += len(seg.data)
		if seg.frg == 0 {
			break
		}
	}

	k.moveBuf()

	// fast recover
	if len(k.rcvQueue) < int(k.rcvWnd) && recover {
		k.probe |= AskTell
	}

	return written, nil
}

func (k *KCP) peekSize() (int, error) {
	if len(k.rcvQueue) == 0 {
		return 0, ErrRecvQueueEmpty
	}
	seg := &k.rcvQueue[0]
	if seg.frg == 0 {
		return len(seg.data), nil
	}
	if len(k.rcvQueue) < int(seg.frg)+1 {
		return 0, ErrExpectingFragment
	}
	length := 0
	for i := range k.rcvQueue {
		length += len(k.rcvQueue[i].data)
		if k.rcvQueue[i].frg == 0 {
			break
		}
	}
	return length, nil
}

// Send queues buf for transmission, splitting it into mss-sized fragments.
func (k *KCP) Send(buf []byte) (int, error) {
	sent := 0

	if k.stream && len(k.sndQueue) > 0 {
		old := &k.sndQueue[len(k.sndQueue)-1]
		if l := len(old.data); l < k.mss {
			capacity := k.mss - l
			extend := len(buf)
			if extend > capacity {
				extend = capacity
			}
			old.data = append(old.data, buf[:extend]...)
			buf = buf[extend:]
			old.frg = 0
			sent += extend
		}
		if len(buf) == 0 {
			return sent, nil
		}
	}

	count := 1
	if len(buf) > k.mss {
		count = (len(buf) + k.mss - 1) / k.mss
	}
	if count >= WndRcv {
		return 0, ErrUserBufTooBig
	}
	if count < 1 {
		count = 1
	}

	for i := 0; i < count; i++ {
		size := len(buf)
		if size > k.mss {
			size = k.mss
		}

		seg := segment{data: append([]byte(nil), buf[:size]...)}
		buf = buf[size:]

		if k.stream {
			seg.frg = 0
		} else {
			seg.frg = uint8(count - i - 1)
		}

		k.sndQueue = append(k.sndQueue, seg)
		sent += size
	}

	return sent, nil
}

func (k *KCP) updateAck(rtt uint32) {
	if k.rxSrtt == 0 {
		k.rxSrtt = rtt
		k.rxRttval = rtt / 2
	} else {
		delta := rtt
		if rtt < k.rxSrtt {
			delta = k.rxSrtt - rtt
		} else {
			delta = rtt - k.rxSrtt
		}
		k.rxRttval = (3*k.rxRttval + delta) / 4
		k.rxSrtt = uint32((7*uint64(k.rxSrtt) + uint64(rtt)) / 8)
		if k.rxSrtt < 1 {
			k.rxSrtt = 1
		}
	}
	rto := k.rxSrtt
	if m := 4 * k.rxRttval; m > k.interval {
		rto += m
	} else {
		rto += k.interval
	}
	k.rxRto = bound(k.rxMinrto, rto, RtoMax)
}

func (k *KCP) shrinkBuf() {
	if len(k.sndBuf) > 0 {
		k.sndUna = k.sndBuf[0].sn
	} else {
		k.sndUna = k.sndNxt
	}
}

func (k *KCP) parseAck(sn uint32) {
	if timeDiff(sn, k.sndUna) < 0 || timeDiff(sn, k.sndNxt) >= 0 {
		return
	}
	for i := range k.sndBuf {
		if sn == k.sndBuf[i].sn {
			k.sndBuf = append(k.sndBuf[:i], k.sndBuf[i+1:]...)
			break
		} else if sn < k.sndBuf[i].sn {
			break
		}
	}
}

func (k *KCP) parseUna(una uint32) {
	for len(k.sndBuf) > 0 && timeDiff(una, k.sndBuf[0].sn) > 0 {
		k.sndBuf = k.sndBuf[1:]
	}
}

func (k *KCP) parseFastack(sn, ts uint32) {
	if timeDiff(sn, k.sndUna) < 0 || timeDiff(sn, k.sndNxt) >= 0 {
		return
	}
	for i := range k.sndBuf {
		seg := &k.sndBuf[i]
		if timeDiff(sn, seg.sn) < 0 {
			break
		} else if sn != seg.sn && timeDiff(ts, seg.ts) >= 0 {
			seg.fastack++
		}
	}
}

func (k *KCP) ackPush(sn, ts uint32) {
	k.acklist = append(k.acklist, [2]uint32{sn, ts})
}

func (k *KCP) parseData(newseg segment) {
	sn := newseg.sn

	if timeDiff(sn, k.rcvNxt+uint32(k.rcvWnd)) >= 0 || timeDiff(sn, k.rcvNxt) < 0 {
		return
	}

	repeat := false
	newIndex := len(k.rcvBuf)
	for i := len(k.rcvBuf) - 1; i >= 0; i-- {
		if k.rcvBuf[i].sn == sn {
			repeat = true
			break
		}
		if timeDiff(sn, k.rcvBuf[i].sn) > 0 {
			break
		}
		newIndex--
	}

	if !repeat {
		k.rcvBuf = append(k.rcvBuf, segment{})
		copy(k.rcvBuf[newIndex+1:], k.rcvBuf[newIndex:])
		k.rcvBuf[newIndex] = newseg
	}

	k.moveBuf()
}

// Input feeds a raw UDP payload (possibly containing several segments) into KCP.
func (k *KCP) Input(buf []byte) error {
	if len(buf) < Overhead {
		return fmt.Errorf("%w: bufsize=%d, at least %d", ErrInvalidSegmentSize, len(buf), Overhead)
	}

	flag := false
	maxAck := uint32(0)
	oldUna := k.sndUna
	latestTs := uint32(0)

	pos := 0
	for len(buf)-pos >= Overhead {
		segBuf := buf[pos:]
		conv := binary.LittleEndian.Uint32(segBuf[0:])
		if conv != k.conv {
			if k.inputConv {
				k.conv = conv
				k.inputConv = false
			} else {
				return fmt.Errorf("%w: expected %d got %d", ErrConvInconsistent, k.conv, conv)
			}
		}

		token := binary.LittleEndian.Uint32(segBuf[4:])
		cmd := segBuf[8]
		frg := segBuf[9]
		wnd := binary.LittleEndian.Uint16(segBuf[10:])
		ts := binary.LittleEndian.Uint32(segBuf[12:])
		sn := binary.LittleEndian.Uint32(segBuf[16:])
		una := binary.LittleEndian.Uint32(segBuf[20:])
		length := int(binary.LittleEndian.Uint32(segBuf[24:]))

		if len(buf)-pos-Overhead < length {
			return fmt.Errorf("%w: payload length=%d remaining=%d", ErrInvalidSegmentData, length, len(buf)-pos-Overhead)
		}

		switch cmd {
		case CmdPush, CmdAck, CmdWask, CmdWins:
		default:
			return fmt.Errorf("%w: cmd=%d", ErrUnsupportedCmd, cmd)
		}

		if token != k.token {
			return fmt.Errorf("%w: got %d expected %d", ErrTokenMismatch, token, k.token)
		}

		k.rmtWnd = wnd
		k.parseUna(una)
		k.shrinkBuf()

		switch cmd {
		case CmdAck:
			rtt := timeDiff(k.current, ts)
			if rtt >= 0 {
				k.updateAck(uint32(rtt))
			}
			k.parseAck(sn)
			k.shrinkBuf()

			if !flag {
				flag = true
				maxAck = sn
				latestTs = ts
			} else if timeDiff(sn, maxAck) > 0 && timeDiff(ts, latestTs) > 0 {
				maxAck = sn
				latestTs = ts
			}
		case CmdPush:
			if timeDiff(sn, k.rcvNxt+uint32(k.rcvWnd)) < 0 {
				k.ackPush(sn, ts)
				if timeDiff(sn, k.rcvNxt) >= 0 {
					data := make([]byte, length)
					copy(data, segBuf[Overhead:Overhead+length])

					seg := segment{
						conv:  conv,
						token: token,
						cmd:   cmd,
						frg:   frg,
						wnd:   wnd,
						ts:    ts,
						sn:    sn,
						una:   una,
						data:  data,
					}
					k.parseData(seg)
				}
			}
		case CmdWask:
			k.probe |= AskTell
		case CmdWins:
			// nothing
		}

		pos += Overhead + length
	}

	if flag {
		k.parseFastack(maxAck, latestTs)
	}

	if timeDiff(k.sndUna, oldUna) > 0 && k.cwnd < k.rmtWnd {
		mss := k.mss
		if k.cwnd < k.ssthresh {
			k.cwnd++
			k.incr += mss
		} else {
			if k.incr < mss {
				k.incr = mss
			}
			k.incr += (mss*mss)/k.incr + (mss / 16)
			if (int(k.cwnd)+1)*mss <= k.incr {
				k.cwnd = uint16((k.incr + mss - 1) / mss)
			}
		}
		if k.cwnd > k.rmtWnd {
			k.cwnd = k.rmtWnd
			k.incr = int(k.rmtWnd) * mss
		}
	}

	return nil
}

func (k *KCP) wndUnused() uint16 {
	if len(k.rcvQueue) < int(k.rcvWnd) {
		return k.rcvWnd - uint16(len(k.rcvQueue))
	}
	return 0
}

// appendSeg appends the wire encoding of seg to buf and returns the new buf.
func appendSeg(buf []byte, seg *segment) []byte {
	buf = append(buf, make([]byte, Overhead+len(seg.data))...)
	seg.encode(buf[len(buf)-Overhead-len(seg.data):])
	return buf
}

func (k *KCP) probeWndSize() {
	if k.rmtWnd == 0 {
		if k.probeWait == 0 {
			k.probeWait = ProbeInit
			k.tsProbe = k.current + k.probeWait
		} else if timeDiff(k.current, k.tsProbe) >= 0 {
			if k.probeWait < ProbeInit {
				k.probeWait = ProbeInit
			}
			k.probeWait += k.probeWait / 2
			if k.probeWait > ProbeLimit {
				k.probeWait = ProbeLimit
			}
			k.tsProbe = k.current + k.probeWait
			k.probe |= AskSend
		}
	} else {
		k.tsProbe = 0
		k.probeWait = 0
	}
}

// FlushAck writes pending ACKs immediately.
func (k *KCP) FlushAck() error {
	if !k.updated {
		return ErrNeedUpdate
	}

	seg := segment{
		conv:  k.conv,
		token: k.token,
		cmd:   CmdAck,
		wnd:   k.wndUnused(),
		una:   k.rcvNxt,
	}

	for _, ack := range k.acklist {
		if len(k.buf)+Overhead > k.mtu {
			if err := k.output(k.buf); err != nil {
				return err
			}
			k.buf = k.buf[:0]
		}
		seg.sn = ack[0]
		seg.ts = ack[1]
		k.buf = append(k.buf, make([]byte, Overhead+len(seg.data))...)
		seg.encode(k.buf[len(k.buf)-Overhead-len(seg.data):])
	}
	k.acklist = k.acklist[:0]

	return nil
}

// Flush moves queued data into the send buffer, handles retransmission and
// congestion state, and emits packets through the output callback.
func (k *KCP) Flush() error {
	if !k.updated {
		return ErrNeedUpdate
	}

	seg := segment{
		conv:  k.conv,
		token: k.token,
		cmd:   CmdAck,
		wnd:   k.wndUnused(),
		una:   k.rcvNxt,
	}
	hdrWnd := seg.wnd
	for _, ack := range k.acklist {
		if len(k.buf)+Overhead > k.mtu {
			if err := k.output(k.buf); err != nil {
				return err
			}
			k.buf = k.buf[:0]
		}
		seg.sn = ack[0]
		seg.ts = ack[1]
		k.buf = appendSeg(k.buf, &seg)
	}
	k.acklist = k.acklist[:0]

	k.probeWndSize()

	// flush window probing commands
	if k.probe&AskSend != 0 {
		seg.cmd = CmdWask
		if len(k.buf)+Overhead > k.mtu {
			if err := k.output(k.buf); err != nil {
				return err
			}
			k.buf = k.buf[:0]
		}
		k.buf = appendSeg(k.buf, &seg)
	}
	if k.probe&AskTell != 0 {
		seg.cmd = CmdWins
		if len(k.buf)+Overhead > k.mtu {
			if err := k.output(k.buf); err != nil {
				return err
			}
			k.buf = k.buf[:0]
		}
		k.buf = appendSeg(k.buf, &seg)
	}
	k.probe = 0

	// calculate window size
	cwnd := k.sndWnd
	if k.rmtWnd < cwnd {
		cwnd = k.rmtWnd
	}
	if !k.nocwnd && k.cwnd < cwnd {
		cwnd = k.cwnd
	}

	// move data from snd_queue to snd_buf
	for timeDiff(k.sndNxt, k.sndUna+uint32(cwnd)) < 0 {
		if len(k.sndQueue) == 0 {
			break
		}
		newseg := k.sndQueue[0]
		k.sndQueue = k.sndQueue[1:]

		newseg.conv = k.conv
		newseg.token = k.token
		newseg.cmd = CmdPush
		newseg.wnd = seg.wnd
		newseg.ts = k.current
		newseg.sn = k.sndNxt
		k.sndNxt++
		newseg.una = k.rcvNxt
		newseg.resendts = k.current
		newseg.rto = k.rxRto
		newseg.fastack = 0
		newseg.xmit = 0
		k.sndBuf = append(k.sndBuf, newseg)
	}

	// calculate resent
	resent := uint32(0)
	if k.fastresend > 0 {
		resent = k.fastresend
	} else {
		resent = ^uint32(0)
	}

	rtomin := uint32(0)
	if !k.nodelay {
		rtomin = k.rxRto >> 3
	}

	lost := false
	change := 0

	for i := range k.sndBuf {
		seg := &k.sndBuf[i]
		needSend := false

		if seg.xmit == 0 {
			needSend = true
			seg.xmit++
			seg.rto = k.rxRto
			seg.resendts = k.current + seg.rto + rtomin
		} else if timeDiff(k.current, seg.resendts) >= 0 {
			needSend = true
			seg.xmit++
			k.xmit++
			if !k.nodelay {
				seg.rto += max(seg.rto, k.rxRto)
			} else {
				seg.rto += seg.rto / 2
			}
			seg.resendts = k.current + seg.rto
			lost = true
		} else if seg.fastack >= resent && (seg.xmit <= k.fastlimit || k.fastlimit == 0) {
			needSend = true
			seg.xmit++
			seg.fastack = 0
			seg.resendts = k.current + seg.rto
			change++
		}

		if needSend {
			seg.ts = k.current
			seg.wnd = hdrWnd
			seg.una = k.rcvNxt

			need := Overhead + len(seg.data)
			if len(k.buf)+need > k.mtu {
				if err := k.output(k.buf); err != nil {
					return err
				}
				k.buf = k.buf[:0]
			}

			k.buf = appendSeg(k.buf, seg)

			if seg.xmit >= k.deadLink {
				k.state = -1
			}
		}
	}

	// flush all data in buffer
	if len(k.buf) > 0 {
		if err := k.output(k.buf); err != nil {
			return err
		}
		k.buf = k.buf[:0]
	}

	// update ssthresh
	if change > 0 {
		inflight := k.sndNxt - k.sndUna
		k.ssthresh = uint16(inflight) / 2
		if k.ssthresh < ThreshMin {
			k.ssthresh = ThreshMin
		}
		k.cwnd = k.ssthresh + uint16(resent)
		k.incr = int(k.cwnd) * k.mss
	}

	if lost {
		k.ssthresh = cwnd / 2
		if k.ssthresh < ThreshMin {
			k.ssthresh = ThreshMin
		}
		k.cwnd = 1
		k.incr = k.mss
	}

	if k.cwnd < 1 {
		k.cwnd = 1
		k.incr = k.mss
	}

	return nil
}

// Update drives the KCP state machine; call it regularly (every ~100ms) and
// after Input/Send, passing the current timestamp in the same unit used for
// session timing.
func (k *KCP) Update(current uint32) error {
	k.current = current

	if !k.updated {
		k.updated = true
		k.tsFlush = k.current
	}

	slap := timeDiff(k.current, k.tsFlush)
	if slap < -10000 || slap >= 10000 {
		k.tsFlush = k.current
		slap = 0
	}

	if slap >= 0 {
		k.tsFlush += k.interval
		if timeDiff(k.current, k.tsFlush) >= 0 {
			k.tsFlush = k.current + k.interval
		}
		if err := k.Flush(); err != nil {
			return err
		}
	}

	return nil
}

// SetNodelay configures latency behaviour (nodelay, interval, resend, nc).
func (k *KCP) SetNodelay(nodelay bool, interval int32, resend int32, nc bool) {
	if nodelay {
		k.nodelay = true
		k.rxMinrto = RtoNdl
	} else {
		k.nodelay = false
		k.rxMinrto = RtoMin
	}

	switch {
	case interval < 10:
		k.interval = 10
	case interval > 5000:
		k.interval = 5000
	default:
		k.interval = uint32(interval)
	}

	if resend >= 0 {
		k.fastresend = uint32(resend)
	}

	k.nocwnd = nc
}

// Now reports the current internal clock (used by tests/tools driving Update).
func (k *KCP) Now() uint32 { return k.current }