package forward

import (
	"git.makemake.in/kzkzzzz/mycommon/mylog"
	"log"
	"math/rand"
	"net"
	"time"
)

var _ IForward = (*UDP)(nil)

type UDP struct {
	forwardInfo      *Info
	conn             *net.UDPConn
	packetBufferSize int
	messageTimeout   time.Duration
}

func NewUDP(forwardInfo *Info) *UDP {
	return &UDP{
		forwardInfo:      forwardInfo,
		packetBufferSize: 2048,
		messageTimeout:   time.Second * 2,
	}
}

func (u *UDP) Forward() error {
	udpLocalAddr, err := net.ResolveUDPAddr("udp", u.forwardInfo.LocalAddr)
	if err != nil {
		return err
	}

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

	log.Printf("[UDP] [%s] %s -> %s", u.forwardInfo.Name, u.forwardInfo.LocalAddr, u.forwardInfo.TargetAddr)

	go func() {
		err := u.handleConn(conn)
		if err != nil {
			mylog.Error(err)
		}
	}()

	return nil
}

func (u *UDP) handleConn(conn *net.UDPConn) error {
	defer func() {
		mylog.Warnf("udp forward stop %s -> %+v", u.forwardInfo.LocalAddr, u.forwardInfo.TargetAddr)
	}()

	buffer := make([]byte, u.packetBufferSize) // 创建一个缓冲区

	for {
		// 接收数据
		n, clientAddr, err := conn.ReadFromUDP(buffer)
		if err != nil {
			mylog.Error(err)
			return err
		}

		targetAddr := u.forwardInfo.TargetAddr[rand.Intn(len(u.forwardInfo.TargetAddr))]

		mylog.Debugf("udp forward %s -> %s", clientAddr.String(), targetAddr)

		// 解析目标地址
		//remoteAddr, err := net.ResolveUDPAddr("udp", targetAddr)
		//if err != nil {
		//	mylog.Error(err)
		//	return err
		//}

		dialer := &net.Dialer{
			Timeout: u.messageTimeout,
		}
		remoteRawConn, err := dialer.Dial("udp", targetAddr)
		if err != nil {
			mylog.Error(err)
			continue
		}

		remoteRawConn.SetWriteDeadline(time.Now().Add(u.messageTimeout))
		remoteRawConn.SetReadDeadline(time.Now().Add(u.messageTimeout))

		remoteConn := remoteRawConn.(*net.UDPConn)

		//remoteConn, err := net.DialUDP("udp", nil, remoteAddr)
		//dial, err := reuseport.Dial("udp", "", targetAddr)

		// 将数据转发到目标地址
		_, err = remoteConn.Write(buffer[:n])
		if err != nil {
			mylog.Error(err)
			continue
		}

		res := make([]byte, u.packetBufferSize)
		resLen, _, err := remoteConn.ReadFromUDP(res)
		if err != nil {
			mylog.Error(err)
			continue
		}

		// 将目标地址的响应返回给客户端
		_, err = conn.WriteToUDP(res[:resLen], clientAddr)
		if err != nil {
			mylog.Error(err)
			continue
		}

	}

}

func (u *UDP) Stop() {
	if u.conn != nil {
		u.conn.Close()
	}
}