// Package stream 提供 TCP/UDP Stream 代理功能。
//
// 该文件实现第四层（传输层）代理，支持 MySQL、PostgreSQL、DNS 等服务的代理转发。
// 与 HTTP 代理不同，Stream 代理不解析应用层协议，而是进行透明的双向数据转发。
//
// 主要功能：
//   - TCP 代理：支持 TCP 连接的代理转发
//   - UDP 代理：支持 UDP 数据报的代理转发
//   - 负载均衡：支持轮询和最少连接算法
//   - 健康检查：定期检查后端服务可用性
//   - 会话管理：UDP 会话自动过期清理
//
// 使用示例：
//
//	server := stream.NewServer()
//	err := server.AddUpstream("mysql", []stream.TargetSpec{
//	    {Addr: "db1:3306", Weight: 1},
//	    {Addr: "db2:3306", Weight: 2},
//	}, "round_robin", stream.HealthCheckSpec{Enabled: true})
//	if err != nil {
//	    log.Fatal(err)
//	}
//
//	server.ListenTCP(":3306")
//	server.Start()
//	defer server.Stop()
//
// 作者：xfy
package stream

import (
	"hash/fnv"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

// 负载均衡方法常量。
const (
	balanceMethodIPHash = "ip_hash"
)

// Balancer Stream 代理（L4 层）负载均衡器接口。
//
// Stream Balancer 特性（区别于 HTTP Balancer）：
//   - 仅 Select(): 按算法策略选择健康目标
//   - 无 SelectExcluding(): Stream 代理无 failover 重试机制
//
// 语义差异说明：
//   - Stream 代理工作在传输层（L4），连接建立后直接转发数据，无应用层重试
//   - HTTP 代理工作在应用层（L7），支持 next_upstream 配置的失败重试
//   - 因此 Stream Balancer 接口签名更简单，仅需要 Select 方法
//   - HTTP Balancer 需要 SelectExcluding 用于排除失败节点
//   - 两种 Balancer 接口签名不同，不可合并
type Balancer interface {
	Select(targets []*Target) *Target
}

// roundRobin 简单轮询负载均衡器。
//
// 使用原子计数器实现线程安全的轮询选择，每次选择后计数器递增，
// 确保请求均匀分布到所有健康目标。
type roundRobin struct {
	counter     atomic.Uint64
	healthyPool sync.Pool
}

func newRoundRobin() Balancer {
	rr := &roundRobin{}
	rr.healthyPool = sync.Pool{
		New: func() any {
			// 用指针包装 slice，避免 Put 时的装箱分配
			s := make([]*Target, 0, 64)
			return &s
		},
	}
	return rr
}

// Select 选择下一个目标（轮询算法）。
//
// 从健康目标列表中按轮询顺序选择。
func (r *roundRobin) Select(targets []*Target) *Target {
	// 从池中获取 healthy slice 并复用
	healthyPtr := r.healthyPool.Get().(*[]*Target) //nolint:errcheck // pool always returns valid *[]*Target
	healthy := *healthyPtr
	healthy = healthy[:0] // 清空但保留容量
	for _, t := range targets {
		if t.healthy.Load() {
			healthy = append(healthy, t)
		}
	}
	*healthyPtr = healthy
	if len(healthy) == 0 {
		r.healthyPool.Put(healthyPtr)
		return nil
	}
	idx := r.counter.Add(1) - 1
	result := healthy[idx%uint64(len(healthy))]
	r.healthyPool.Put(healthyPtr)
	return result
}

// leastConn 最少连接负载均衡器。
//
// 跟踪每个目标服务器的活跃连接数，总是选择连接数最少的目标。
// 适合连接持续时间差异较大的场景。
type leastConn struct{}

// newLeastConn 创建最少连接均衡器。
//
// 返回实现了最少连接算法的 Balancer 接口实例。
//
// 返回值：
//   - Balancer: 最少连接负载均衡器实例
func newLeastConn() Balancer {
	return &leastConn{}
}

// Select 选择连接最少的目标。
//
// 遍历所有健康目标，选择当前连接数最少的目标服务器。
// 只选择标记为健康的目标，如果无健康目标则返回 nil。
//
// 参数：
//   - targets: 目标服务器列表
//
// 返回值：
//   - *Target: 选中的目标服务器，无可用目标时返回 nil
func (l *leastConn) Select(targets []*Target) *Target {
	var selected *Target
	var minConns int64 = -1
	for _, t := range targets {
		if !t.healthy.Load() {
			continue
		}
		conns := atomic.LoadInt64(&t.conns)
		if selected == nil || conns < minConns {
			selected = t
			minConns = conns
		}
	}
	return selected
}

// weightedRoundRobin 加权轮询负载均衡器。
//
// 根据目标服务器的权重分配请求，权重高的目标获得更多请求。
// 使用原子计数器确保线程安全，支持不同权重的目标混合使用。
type weightedRoundRobin struct {
	counter     atomic.Uint64
	healthyPool sync.Pool
}

// newWeightedRoundRobin 创建加权轮询均衡器。
//
// 返回实现了加权轮询算法的 Balancer 接口实例。
//
// 返回值：
//   - Balancer: 加权轮询负载均衡器实例
func newWeightedRoundRobin() Balancer {
	w := &weightedRoundRobin{}
	w.healthyPool = sync.Pool{
		New: func() any {
			s := make([]*Target, 0, 64)
			return &s
		},
	}
	return w
}

// Select 选择下一个目标（加权轮询算法）。
//
// 从健康目标列表中按权重比例选择。
func (w *weightedRoundRobin) Select(targets []*Target) *Target {
	healthyPtr := w.healthyPool.Get().(*[]*Target) //nolint:errcheck // pool always returns valid *[]*Target
	healthy := *healthyPtr
	healthy = healthy[:0]
	for _, t := range targets {
		if t.healthy.Load() {
			healthy = append(healthy, t)
		}
	}
	*healthyPtr = healthy
	if len(healthy) == 0 {
		w.healthyPool.Put(healthyPtr)
		return nil
	}

	// 计算总权重
	totalWeight := 0
	for _, t := range healthy {
		if t.weight <= 0 {
			totalWeight++ // 最小权重为 1
		} else {
			totalWeight += t.weight
		}
	}

	// 使用原子计数器确定位置
	idx := w.counter.Add(1) - 1
	pos := int(idx % uint64(totalWeight))

	// 找到对应位置的目标
	currentWeight := 0
	for _, t := range healthy {
		weight := t.weight
		if weight <= 0 {
			weight = 1
		}
		currentWeight += weight
		if pos < currentWeight {
			w.healthyPool.Put(healthyPtr)
			return t
		}
	}

	result := healthy[len(healthy)-1]
	w.healthyPool.Put(healthyPtr)
	return result
}

// ipHash IP 哈希负载均衡器。
//
// 基于客户端 IP 地址计算哈希值，将同一 IP 的请求分配到固定目标。
// 适合需要会话保持的场景，确保相同客户端总是被路由到同一服务器。
type ipHash struct {
	healthyPool sync.Pool
}

// newIPHash 创建 IP 哈希均衡器。
//
// 返回实现了 IP 哈希算法的 Balancer 接口实例。
//
// 返回值：
//   - Balancer: IP 哈希负载均衡器实例
func newIPHash() Balancer {
	ih := &ipHash{}
	ih.healthyPool = sync.Pool{
		New: func() any {
			s := make([]*Target, 0, 64)
			return &s
		},
	}
	return ih
}

// Select 选择下一个目标（IP 哈希算法）。
//
// 委托给 SelectByIP，使用空客户端 IP。
func (i *ipHash) Select(targets []*Target) *Target {
	return i.SelectByIP(targets, "")
}

// SelectByIP 根据客户端 IP 选择目标。
//
// 使用 FNV-64a 哈希算法对客户端 IP 进行哈希，从健康目标列表中选择对应的目标。
// 如果 clientIP 为空，则回退到轮询选择。
func (i *ipHash) SelectByIP(targets []*Target, clientIP string) *Target {
	healthyPtr := i.healthyPool.Get().(*[]*Target) //nolint:errcheck // pool always returns valid *[]*Target
	healthy := *healthyPtr
	healthy = healthy[:0]
	for _, t := range targets {
		if t.healthy.Load() {
			healthy = append(healthy, t)
		}
	}
	*healthyPtr = healthy
	if len(healthy) == 0 {
		i.healthyPool.Put(healthyPtr)
		return nil
	}

	// 使用 FNV-64a 哈希
	h := fnv.New64a()
	h.Write([]byte(clientIP))
	hash := h.Sum64()

	idx := hash % uint64(len(healthy))
	result := healthy[idx]
	i.healthyPool.Put(healthyPtr)
	return result
}

// Server TCP/UDP Stream 代理服务器。
type Server struct {
	// listeners TCP 监听器映射，按 upstream 名称索引
	listeners map[string]net.Listener
	// udpServers UDP 服务器映射
	udpServers map[string]*udpServer
	// upstreams 上游配置映射
	upstreams map[string]*Upstream
	// connCount 当前连接数
	connCount int64
	// mu 读写锁，保护并发访问
	mu sync.RWMutex
	// running 运行状态标志
	running atomic.Bool
}

// Upstream Stream 上游配置。
type Upstream struct {
	balancer  Balancer
	healthChk *HealthChecker
	name      string
	targets   []*Target
	mu        sync.RWMutex
}

// Target Stream 代理（L4 层）的目标服务器。
//
// Stream Target 特性（区别于 HTTP Target）：
//   - 简单地址：仅支持 host:port 格式，无 URL 解析
//   - 无 DNS 缓存：直接连接目标地址，无需动态 DNS 解析
//   - 无 failover：Stream 代理无重试机制，仅记录健康状态
//
// 语义差异说明：
//   - Stream 代理工作在传输层（L4），直接转发 TCP/UDP 数据
//   - HTTP 代理工作在应用层（L7），需要 URL 解析和 DNS 动态解析
//   - 因此 Stream Target 保持简单结构，HTTP Target 需要 DNS 缓存等复杂功能
//   - 两种 Target 必须保持独立定义，不可合并
type Target struct {
	// addr 目标地址（host:port）
	addr string
	// weight 权重
	weight int
	// healthy 健康状态
	healthy atomic.Bool
	// conns 当前连接数
	conns int64
}

// HealthChecker Stream 健康检查器。
type HealthChecker struct {
	// upstream 所属上游
	upstream *Upstream
	// stopCh 停止信号通道
	stopCh chan struct{}
	// interval 检查间隔
	interval time.Duration
	// timeout 检查超时
	timeout time.Duration
}

// Config Stream 配置。
type Config struct {
	// Listen 监听地址
	Listen string
	// Protocol 协议类型（tcp 或 udp）
	Protocol string
	// Upstream 上游配置
	Upstream UpstreamSpec
}

// UpstreamSpec 上游配置规格。
type UpstreamSpec struct {
	Name        string
	LoadBalance string
	Targets     []TargetSpec
	HealthCheck HealthCheckSpec
}

// TargetSpec 目标配置规格。
type TargetSpec struct {
	// Addr 目标地址（host:port）
	Addr string
	// Weight 权重
	Weight int
}

// HealthCheckSpec 健康检查配置规格。
type HealthCheckSpec struct {
	// Interval 检查间隔
	Interval time.Duration
	// Timeout 检查超时
	Timeout time.Duration
	// Enabled 是否启用
	Enabled bool
}

// NewServer 创建 Stream 代理服务器实例。
//
// 初始化 TCP/UDP 第四层代理服务器，包含监听器、上游配置和
// 会话管理的空映射。创建后需通过 AddUpstream、ListenTCP/UDP
// 等方法配置后再启动。
//
// 返回值：
//   - *Server: 初始化的 Stream 代理服务器实例
func NewServer() *Server {
	return &Server{
		listeners:  make(map[string]net.Listener),
		udpServers: make(map[string]*udpServer),
		upstreams:  make(map[string]*Upstream),
	}
}

// AddUpstream 添加上游配置。
func (s *Server) AddUpstream(name string, targets []TargetSpec, lbType string, hcSpec HealthCheckSpec) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 创建目标列表
	tgts := make([]*Target, len(targets))
	for i, t := range targets {
		tgts[i] = &Target{
			addr:   t.Addr,
			weight: t.Weight,
		}
		tgts[i].healthy.Store(true) // 初始假设健康
	}

	// 创建负载均衡器
	var balancer Balancer
	switch lbType {
	case "round_robin", "":
		balancer = newRoundRobin()
	case "weighted_round_robin":
		balancer = newWeightedRoundRobin()
	case "least_conn":
		balancer = newLeastConn()
	case balanceMethodIPHash:
		balancer = newIPHash()
	default:
		balancer = newRoundRobin()
	}

	upstream := &Upstream{
		name:     name,
		targets:  tgts,
		balancer: balancer,
	}

	// 启动健康检查
	if hcSpec.Enabled {
		upstream.healthChk = &HealthChecker{
			upstream: upstream,
			interval: hcSpec.Interval,
			timeout:  hcSpec.Timeout,
			stopCh:   make(chan struct{}),
		}
		go upstream.healthChk.Start()
	}

	s.upstreams[name] = upstream
	return nil
}

// ListenTCP 开始监听 TCP 端口。
func (s *Server) ListenTCP(addr string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	s.listeners[addr] = listener
	return nil
}

// ListenUDP 开始监听 UDP 端口。
func (s *Server) ListenUDP(addr string, upstreamName string, timeout time.Duration) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 查找上游配置
	upstream, exists := s.upstreams[upstreamName]
	if !exists {
		return io.ErrClosedPipe
	}

	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return err
	}

	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		return err
	}

	// 创建 UDP 服务器
	udpSrv := newUDPServer(conn, upstream, timeout)
	s.udpServers[addr] = udpSrv

	return nil
}

// Start 启动 Stream 服务器。
func (s *Server) Start() error {
	s.running.Store(true)

	s.mu.RLock()
	defer s.mu.RUnlock()

	// 启动 TCP 监听器
	for addr, listener := range s.listeners {
		go s.acceptLoop(addr, listener)
	}

	// 启动 UDP 服务器
	for _, udpSrv := range s.udpServers {
		go udpSrv.serve()
		go udpSrv.startCleanupTicker()
	}

	return nil
}

// acceptLoop 接受连接循环。
//
// 在单独的 goroutine 中运行，持续接受 TCP 连接。
// 当服务器停止时，该函数返回。
//
// 参数：
//   - addr: 监听地址
//   - listener: TCP 监听器实例
func (s *Server) acceptLoop(addr string, listener net.Listener) {
	for s.running.Load() {
		conn, err := listener.Accept()
		if err != nil {
			if !s.running.Load() {
				return // 正常关闭
			}
			continue
		}

		atomic.AddInt64(&s.connCount, 1)
		go s.handleConnection(conn, addr)
	}
}

// handleConnection 处理单个连接。
//
// 处理客户端连接的完整生命周期：
// 1. 选择上游目标
// 2. 建立到目标服务器的连接
// 3. 在客户端和目标之间双向转发数据
// 4. 处理连接关闭和错误
//
// 参数：
//   - clientConn: 客户端连接
//   - addr: 监听地址
func (s *Server) handleConnection(clientConn net.Conn, _ string) {
	defer func() {
		_ = clientConn.Close()
		atomic.AddInt64(&s.connCount, -1)
	}()

	s.mu.RLock()
	// 根据监听地址找到对应 upstream（简化：用第一个）
	var upstream *Upstream
	for _, up := range s.upstreams {
		upstream = up
		break
	}
	s.mu.RUnlock()

	if upstream == nil {
		return // 无上游配置
	}

	// 选择目标
	target := upstream.Select()
	if target == nil {
		return // 无可用目标
	}

	atomic.AddInt64(&target.conns, 1)
	defer func() { atomic.AddInt64(&target.conns, -1) }()

	// 连接目标
	targetConn, err := net.DialTimeout("tcp", target.addr, 10*time.Second)
	if err != nil {
		target.healthy.Store(false)
		return
	}
	defer func() { _ = targetConn.Close() }()

	// 双向数据转发：任一方向完成/出错时立即关闭双端连接，迫使另一方向退出
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		_, _ = io.Copy(targetConn, clientConn)
		_ = clientConn.Close()
		_ = targetConn.Close()
	}()
	go func() {
		defer wg.Done()
		_, _ = io.Copy(clientConn, targetConn)
		_ = targetConn.Close()
		_ = clientConn.Close()
	}()
	wg.Wait()
}

// Select 选择健康的上游目标。
func (u *Upstream) Select() *Target {
	u.mu.RLock()
	defer u.mu.RUnlock()

	// 获取健康目标列表
	healthyTargets := make([]*Target, 0)
	for _, t := range u.targets {
		if t.healthy.Load() {
			healthyTargets = append(healthyTargets, t)
		}
	}

	if len(healthyTargets) == 0 {
		return nil
	}

	// 使用负载均衡器选择
	return u.balancer.Select(healthyTargets)
}

// Start 启动健康检查。
func (h *HealthChecker) Start() {
	ticker := time.NewTicker(h.interval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			h.check()
		case <-h.stopCh:
			return
		}
	}
}

// check 执行健康检查。
//
// 尝试 TCP 连接到所有目标，根据连接结果更新健康状态。
// 连接成功标记为健康，连接失败标记为不健康。
func (h *HealthChecker) check() {
	for _, target := range h.upstream.targets {
		conn, err := net.DialTimeout("tcp", target.addr, h.timeout)
		if err != nil {
			target.healthy.Store(false)
		} else {
			_ = conn.Close()
			target.healthy.Store(true)
		}
	}
}

// Stats Stream 服务器统计。
type Stats struct {
	// Connections 当前活跃连接数量
	Connections int64

	// Listeners 当前监听器数量（TCP + UDP）
	Listeners int

	// Upstreams 上游配置数量
	Upstreams int
}

// udpSession UDP 会话，管理客户端到后端的映射。
//
// 每个 UDP 会话代表一个客户端与一个后端目标的映射关系。
// 会话包含客户端地址、后端连接、最后活跃时间等信息。
// 会话在空闲超时后自动清理。
type udpSession struct {
	lastActive time.Time
	targetConn net.Conn
	target     *Target
	clientAddr *net.UDPAddr
	srv        *udpServer
	mu         sync.RWMutex
	closeOnce  sync.Once
}

// udpServer UDP 服务器，管理多个客户端会话。
//
// 负责监听 UDP 端口，管理客户端会话的生命周期。
// 支持会话自动过期清理和优雅停止。
type udpServer struct {
	conn     *net.UDPConn
	sessions map[string]*udpSession
	upstream *Upstream
	stopCh   chan struct{}
	wg       sync.WaitGroup
	timeout  time.Duration
	mu       sync.RWMutex
	running  atomic.Bool
}

// newUDPServer 创建新的 UDP 服务器。
//
// 根据 UDP 连接、上游配置和超时时间创建 UDP 服务器实例。
// 如果超时时间小于等于 0，使用默认 60 秒。
//
// 参数：
//   - conn: UDP 连接
//   - upstream: 上游配置
//   - timeout: 会话空闲超时时间
//
// 返回值：
//   - *udpServer: 创建的 UDP 服务器实例
func newUDPServer(conn *net.UDPConn, upstream *Upstream, timeout time.Duration) *udpServer {
	if timeout <= 0 {
		timeout = 60 * time.Second // 默认 60 秒超时
	}
	return &udpServer{
		conn:     conn,
		sessions: make(map[string]*udpSession),
		upstream: upstream,
		timeout:  timeout,
		stopCh:   make(chan struct{}),
	}
}

// sessionKey 从 UDP 地址生成会话键。
//
// 使用 UDP 地址的字符串表示作为会话映射的唯一键。
//
// 参数：
//   - addr: UDP 地址
//
// 返回值：
//   - string: 会话键字符串
func sessionKey(addr *net.UDPAddr) string {
	return addr.String()
}

// getSession 获取现有会话（线程安全）。
//
// 根据客户端地址查找现有会话，如果找到则更新最后活跃时间。
// 如果会话不存在返回 nil。
//
// 参数：
//   - clientAddr: 客户端 UDP 地址
//
// 返回值：
//   - *udpSession: 找到的会话，不存在时返回 nil
func (s *udpServer) getSession(clientAddr *net.UDPAddr) *udpSession {
	s.mu.RLock()
	defer s.mu.RUnlock()

	session, exists := s.sessions[sessionKey(clientAddr)]
	if !exists {
		return nil
	}

	// 更新最后活动时间
	session.mu.Lock()
	session.lastActive = time.Now()
	session.mu.Unlock()

	return session
}

// getOrCreateSession 获取或创建会话（线程安全）。
//
// 首先尝试获取现有会话，如果不存在则创建新会话：
// 1. 选择后端目标
// 2. 建立到目标的 UDP 连接
// 3. 创建会话并启动响应监听 goroutine
//
// 使用双重检查锁定模式避免重复创建会话。
//
// 参数：
//   - clientAddr: 客户端 UDP 地址
//
// 返回值：
//   - *udpSession: 获取或创建的会话
//   - error: 创建失败时返回错误
func (s *udpServer) getOrCreateSession(clientAddr *net.UDPAddr) (*udpSession, error) {
	// 先尝试获取现有会话
	session := s.getSession(clientAddr)
	if session != nil {
		return session, nil
	}

	// 需要创建新会话，获取写锁
	s.mu.Lock()
	defer s.mu.Unlock()

	// 双重检查：可能另一个 goroutine 已经创建了会话
	if existingSession, exists := s.sessions[sessionKey(clientAddr)]; exists {
		existingSession.mu.Lock()
		existingSession.lastActive = time.Now()
		existingSession.mu.Unlock()
		return existingSession, nil
	}

	// 选择后端目标
	target := s.upstream.Select()
	if target == nil {
		return nil, io.ErrClosedPipe
	}

	// 连接到后端（使用 UDP 连接）
	targetAddr, err := net.ResolveUDPAddr("udp", target.addr)
	if err != nil {
		return nil, err
	}

	targetConn, err := net.DialUDP("udp", nil, targetAddr)
	if err != nil {
		return nil, err
	}

	atomic.AddInt64(&target.conns, 1)

	// 创建新会话
	session = &udpSession{
		clientAddr: clientAddr,
		targetConn: targetConn,
		target:     target,
		lastActive: time.Now(),
		srv:        s,
	}

	s.sessions[sessionKey(clientAddr)] = session

	// 启动后端响应监听
	s.wg.Add(1)
	go session.handleBackendResponse()

	return session, nil
}

// removeSession 移除会话（线程安全）。
//
// 关闭会话并删除会话映射。
//
// 参数：
//   - clientAddr: 客户端 UDP 地址
func (s *udpServer) removeSession(clientAddr *net.UDPAddr) {
	s.mu.Lock()
	defer s.mu.Unlock()

	key := sessionKey(clientAddr)
	if session, exists := s.sessions[key]; exists {
		session.close()
		delete(s.sessions, key)
	}
}

// close 关闭会话。
//
// 使用 sync.Once 确保会话只关闭一次。
// 关闭后端目标连接。
func (sess *udpSession) close() {
	sess.closeOnce.Do(func() {
		if sess.targetConn != nil {
			_ = sess.targetConn.Close()
		}
		if sess.target != nil {
			atomic.AddInt64(&sess.target.conns, -1)
		}
	})
}

// handleBackendResponse 处理后端响应并转发回客户端。
//
// 在单独的 goroutine 中运行，持续监听后端响应：
// 1. 读取后端数据
// 2. 转发到客户端
// 3. 处理超时和错误
//
// 当会话超时、连接错误或服务器停止时返回。
func (sess *udpSession) handleBackendResponse() {
	defer sess.srv.wg.Done()

	buf := make([]byte, 65535)
	for {
		// 设置读取超时
		_ = sess.targetConn.SetReadDeadline(time.Now().Add(sess.srv.timeout))

		n, err := sess.targetConn.Read(buf)
		if err != nil {
			// 超时或其他错误，检查是否需要关闭
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				// 检查是否超过空闲超时
				sess.mu.RLock()
				lastActive := sess.lastActive
				sess.mu.RUnlock()

				if time.Since(lastActive) >= sess.srv.timeout {
					sess.srv.removeSession(sess.clientAddr)
					return
				}
				continue
			}
			// 其他错误，关闭会话
			sess.srv.removeSession(sess.clientAddr)
			return
		}

		// 更新活动时间
		sess.mu.Lock()
		sess.lastActive = time.Now()
		sess.mu.Unlock()

		// 发送回客户端
		_, err = sess.srv.conn.WriteToUDP(buf[:n], sess.clientAddr)
		if err != nil {
			// 写入客户端失败，关闭会话
			sess.srv.removeSession(sess.clientAddr)
			return
		}
	}
}

// serve 启动 UDP 服务循环。
//
// 在单独的 goroutine 中运行，持续监听 UDP 数据报：
// 1. 接收客户端数据报
// 2. 获取或创建会话
// 3. 转发数据到后端
//
// 当服务器停止时返回。
func (s *udpServer) serve() {
	s.running.Store(true)

	buf := make([]byte, 65535)
	for s.running.Load() {
		// 设置读取超时，以便定期检查 stopCh
		_ = s.conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))

		n, clientAddr, err := s.conn.ReadFromUDP(buf)
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				// 检查是否需要停止
				select {
				case <-s.stopCh:
					return
				default:
					continue
				}
			}
			continue
		}

		// 获取或创建会话
		session, err := s.getOrCreateSession(clientAddr)
		if err != nil {
			continue
		}

		// 转发数据到后端
		_, err = session.targetConn.Write(buf[:n])
		if err != nil {
			// 写入失败，移除会话
			s.removeSession(clientAddr)
		}
	}
}

// startCleanupTicker 启动定期清理过期会话的 ticker。
//
// 每 10 秒执行一次清理，移除空闲超时的会话。
// 当服务器停止时返回。
func (s *udpServer) startCleanupTicker() {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			s.cleanupExpiredSessions()
		case <-s.stopCh:
			return
		}
	}
}

// cleanupExpiredSessions 清理过期会话。
//
// 遍历所有会话，移除空闲时间超过 timeout 的会话。
// 必须在持有写锁的情况下调用。
func (s *udpServer) cleanupExpiredSessions() {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := time.Now()
	for key, session := range s.sessions {
		session.mu.RLock()
		lastActive := session.lastActive
		session.mu.RUnlock()

		if now.Sub(lastActive) >= s.timeout {
			session.close()
			delete(s.sessions, key)
		}
	}
}