package client

import (
	"fmt"
	"log"
	"net"
	pb "net-tunnel/pkg/proto"
	"time"
)

func (c *Client) handleUDPData(msg *pb.Message_ProxyData) {
	log.Printf("Received UDP proxy data for connection: %v", msg.ProxyData.ConnId)
	hostPort := net.JoinHostPort(msg.ProxyData.ProxyConfig.LocalIp, fmt.Sprintf("%d", msg.ProxyData.ProxyConfig.LocalPort))

	// 获取或创建 UDP 连接
	udpConnKey := msg.ProxyData.ConnId
	var conn *net.UDPConn

	existingConn, ok := c.udpConnections.Load(udpConnKey)
	if ok {
		conn = existingConn.(*net.UDPConn)
	} else {
		udpAddr, err := net.ResolveUDPAddr("udp", hostPort)
		if err != nil {
			log.Printf("Failed to resolve UDP address: %v", err)
			return
		}
		conn, err = net.DialUDP("udp", nil, udpAddr)
		if err != nil {
			log.Printf("Failed to dial UDP address: %v", err)
			return
		}
		c.udpConnections.Store(udpConnKey, conn)

		go c.readFromUDPConn(msg, conn)
	}

	_, err := conn.Write(msg.ProxyData.Data)
	if err != nil {
		log.Printf("Failed to write data to UDP connection: %v", err)
		conn.Close()
		c.udpConnections.Delete(udpConnKey)
		return
	}
}

func (c *Client) readFromUDPConn(msg *pb.Message_ProxyData, conn *net.UDPConn) {
	buffer := make([]byte, 4096)
	for {
		err := conn.SetReadDeadline(time.Now().Add(time.Minute * 10))
		if err != nil {
			log.Printf("Failed to set read deadline: %v", err)
			return
		}
		n, _, err := conn.ReadFromUDP(buffer)
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue
			}
			log.Printf("Failed to read from UDP connection: %v", err)
			return
		}
		if n > 0 {
			msg := &pb.Message{
				Content: &pb.Message_ProxyData{
					ProxyData: &pb.ProxyData{
						ConnId:      msg.ProxyData.ConnId,
						ProxyConfig: msg.ProxyData.ProxyConfig,
						Data:        buffer[:n],
					},
				},
			}
			if err := c.stream.Send(msg); err != nil {
				log.Printf("Failed to send proxy data: %v", err)
				break
			}
		}
	}

	conn.Close()
	c.udpConnections.Delete(msg.ProxyData.ConnId)
	log.Printf("UDP connection closed: %v", msg.ProxyData.ConnId)
}