package transport

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

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

// UDPConn 是对 UDP 连接的轻量封装。
// server 侧：共享同一个 net.UDPConn，Send 时通过 peerAddr 指定对端地址。
// peer 侧：独立的 net.UDPConn，已通过 Dial 连接到 server，Send 直接写即可。
type UDPConn struct {
	conn     *net.UDPConn
	peerAddr *net.UDPAddr    // server 侧为对端地址；peer 侧为 nil（连接模式下直接 Write）
	raw      syscall.RawConn // 底层 syscall 句柄，用于 Linux socket timestamping

	logger                 latencylog.Logger
	txTimestampDebugLogger TXTimestampDebugLogger
	nodeRole               string     // 日志中记录的节点角色，例如 "server" 或 "peer"
	nodeID                 string     // 日志中记录的节点 ID
	writeMu                sync.Mutex // 保护 Send 的互斥锁
	closeOnce              sync.Once
	closeErr               error
}

// 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 连接封装。
// peerAddr 为 nil 时表示 peer 侧已连接模式（conn 已 Dial 到 server）。
// peerAddr 非 nil 时表示 server 侧，Send 时需要指定目标地址。
func NewUDPConn(conn *net.UDPConn, peerAddr *net.UDPAddr, opts ...UDPOption) (*UDPConn, error) {
	udpConn := &UDPConn{
		conn:     conn,
		peerAddr: peerAddr,
		logger:   latencylog.NoopLogger{},
	}

	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 数据报并发送。
// 多个 goroutine 可以并发调用，内部会串行化写入。
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 数据报并发送到指定地址。
// 主要用于 server 侧向特定 peer 发送消息。
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 连接读取一条完整协议消息。
// 返回解码后的消息和来源地址（peer 侧来源地址始终为 server 地址）。
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 处理。
// 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 连接，保证重复调用安全。
// 注意：server 侧多个 UDPConn 共享同一个 net.UDPConn 时，
// 只应由 UDPHub 负责关闭底层连接，不应通过此方法关闭。
func (c *UDPConn) Close() error {
	c.closeOnce.Do(func() {
		c.closeErr = c.conn.Close()
	})

	return c.closeErr
}