package transport

import (
	"fmt"
	"net"
	"sync"
	"syscall"

	"omnisocketgo/cmd/internal/latencylog"
	"omnisocketgo/cmd/internal/protocol"
)

// UDPConn 是对 UDP 连接的轻量封装。
// server 侧共享一个 net.UDPConn，通过 SendTo 指定目标地址；
// peer 侧使用已连接的 net.UDPConn，直接 Send 即可。
type UDPConn struct {
	conn     *net.UDPConn
	peerAddr *net.UDPAddr
	raw      syscall.RawConn

	logger                 latencylog.Logger
	txTimestampDebugLogger TXTimestampDebugLogger
	txPacketSeq            uint32
	pendingTX              map[uint32]udpTXPendingRecord
	nodeRole               string
	nodeID                 string
	writeMu                sync.Mutex
	closeOnce              sync.Once
	closeErr               error
}

type udpTXPendingRecord struct {
	msg                protocol.Message
	sendCallIndex      int
	bytesWritten       int
	expectedTXID       uint32
	observedTimestamps map[string]int64
}

// UDPOption 用于为 UDPConn 注入可选行为。
type UDPOption func(*UDPConn)

// WithUDPLogger 为 UDP 连接注入业务消息日志上下文。
func WithUDPLogger(logger latencylog.Logger, nodeRole, nodeID string) UDPOption {
	return func(conn *UDPConn) {
		conn.logger = logger
		conn.nodeRole = nodeRole
		conn.nodeID = nodeID
	}
}

// WithUDPTXTimestampDebugLogger 为 UDP 连接注入 TX errqueue 调试日志器。
func WithUDPTXTimestampDebugLogger(logger TXTimestampDebugLogger) UDPOption {
	return func(conn *UDPConn) {
		conn.txTimestampDebugLogger = logger
	}
}

// NewUDPConn 创建 UDP transport 连接封装。
func NewUDPConn(conn *net.UDPConn, peerAddr *net.UDPAddr, opts ...UDPOption) (*UDPConn, error) {
	udpConn := &UDPConn{
		conn:      conn,
		peerAddr:  peerAddr,
		logger:    latencylog.NoopLogger{},
		pendingTX: make(map[uint32]udpTXPendingRecord),
	}

	for _, opt := range opts {
		opt(udpConn)
	}

	if udpConn.logger == nil {
		udpConn.logger = latencylog.NoopLogger{}
	}

	if err := udpConn.initUDPLinuxTimestamping(); err != nil {
		return nil, err
	}

	return udpConn, nil
}

// Send 将一条协议消息编码为 UDP 数据报并发送。
func (c *UDPConn) Send(msg protocol.Message) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()

	latencylog.LogMessageEvent(c.logger, c.nodeRole, c.nodeID, latencylog.EventSendHandoffBegin, msg)
	if err := c.sendMessageLinux(msg); err != nil {
		return fmt.Errorf("transport: udp send message: %w", err)
	}
	latencylog.LogMessageEvent(c.logger, c.nodeRole, c.nodeID, latencylog.EventSendHandoffEnd, msg)

	return nil
}

// SendTo 将一条协议消息编码为 UDP 数据报并发送到指定地址。
func (c *UDPConn) SendTo(msg protocol.Message, addr *net.UDPAddr) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()

	latencylog.LogMessageEvent(c.logger, c.nodeRole, c.nodeID, latencylog.EventSendHandoffBegin, msg)
	if err := c.sendMessageToLinux(msg, addr); err != nil {
		return fmt.Errorf("transport: udp send message to %s: %w", addr, err)
	}
	latencylog.LogMessageEvent(c.logger, c.nodeRole, c.nodeID, latencylog.EventSendHandoffEnd, msg)

	return nil
}

// Receive 从 UDP 连接读取一条完整协议消息。
func (c *UDPConn) Receive() (protocol.Message, *net.UDPAddr, error) {
	msg, addr, err := c.receiveMessageLinux()
	if err != nil {
		return protocol.Message{}, nil, fmt.Errorf("transport: udp receive message: %w", err)
	}

	return msg, addr, nil
}

// ReceiveLoop 持续从 UDP 连接读取消息并交给 handler 处理。
func (c *UDPConn) ReceiveLoop(handler func(protocol.Message, *net.UDPAddr) error) error {
	for {
		msg, addr, err := c.Receive()
		if err != nil {
			return fmt.Errorf("transport: udp receive loop read: %w", err)
		}

		if err := handler(msg, addr); err != nil {
			return fmt.Errorf("transport: udp receive loop handler: %w", err)
		}
	}
}

// Close 关闭底层 UDP 连接，保证重复调用安全。
func (c *UDPConn) Close() error {
	c.closeOnce.Do(func() {
		c.closeErr = c.conn.Close()
	})

	return c.closeErr
}