package connect

import (
	"fmt"
	"net"
	"time"

	"github.com/pion/stun"
	"go.uber.org/zap"
)

// STUNClient STUN 客户端 - 用于 NAT 探测和候选地址采集
type STUNClient struct {
	servers []string
	logger  *zap.Logger
	timeout time.Duration
}

// NewSTUNClient 创建 STUN 客户端
func NewSTUNClient(servers []string, logger *zap.Logger) *STUNClient {
	return &STUNClient{
		servers: servers,
		logger:  logger,
		timeout: 5 * time.Second,
	}
}

// DiscoverAddress 发现外部地址（通过单个 STUN 服务器）
func (c *STUNClient) DiscoverAddress(server string) (*net.UDPAddr, error) {
	host, port, err := net.SplitHostPort(server)
	if err != nil {
		return nil, fmt.Errorf("STUN 服务器地址格式错误：%w", err)
	}

	udpAddr, err := net.ResolveUDPAddr("udp4", net.JoinHostPort(host, port))
	if err != nil {
		return nil, fmt.Errorf("解析 UDP 地址失败：%w", err)
	}

	conn, err := net.DialUDP("udp4", nil, udpAddr)
	if err != nil {
		return nil, fmt.Errorf("连接 STUN 服务器失败：%w", err)
	}
	defer conn.Close()

	conn.SetDeadline(time.Now().Add(c.timeout))

	// 构建 STUN Binding Request
	msg, err := stun.Build(stun.BindingRequest, stun.TransactionID)
	if err != nil {
		return nil, fmt.Errorf("构建 STUN 请求失败：%w", err)
	}

	// 发送请求
	if _, err := conn.Write(msg.Raw); err != nil {
		return nil, fmt.Errorf("发送 STUN 请求失败：%w", err)
	}

	// 读取响应
	buf := make([]byte, 1024)
	n, err := conn.Read(buf)
	if err != nil {
		return nil, fmt.Errorf("读取 STUN 响应失败：%w", err)
	}

	// 解析响应
	res := &stun.Message{Raw: buf[:n]}
	if err := res.Decode(); err != nil {
		return nil, fmt.Errorf("解码 STUN 响应失败：%w", err)
	}

	// 提取 XOR-MAPPED-ADDRESS
	var xorAddr stun.XORMappedAddress
	if err := xorAddr.GetFrom(res); err != nil {
		return nil, fmt.Errorf("提取外部地址失败：%w", err)
	}

	c.logger.Debug("STUN 查询成功",
		zap.String("server", server),
		zap.String("external_addr", xorAddr.String()))

	return &net.UDPAddr{
		IP:   xorAddr.IP,
		Port: xorAddr.Port,
	}, nil
}

// CollectCandidates 收集候选地址（通过多个 STUN 服务器）
func (c *STUNClient) CollectCandidates() []string {
	var candidates []string

	for _, server := range c.servers {
		addr, err := c.DiscoverAddress(server)
		if err != nil {
			c.logger.Debug("STUN 服务器查询失败",
				zap.String("server", server),
				zap.Error(err))
			continue
		}

		candidates = append(candidates, addr.String())
		c.logger.Debug("收集到候选地址",
			zap.String("server", server),
			zap.String("candidate", addr.String()))
	}

	return candidates
}

// GetExternalIP 获取外部 IP（兼容旧 API）
func (c *STUNClient) GetExternalIP() (string, error) {
	for _, server := range c.servers {
		addr, err := c.DiscoverAddress(server)
		if err == nil {
			return addr.IP.String(), nil
		}
	}
	return "", fmt.Errorf("所有 STUN 服务器均查询失败")
}