package main

import (
	"context"
	"fmt"
	"net"
	"runtime"
	"sync"
	"time"
)

var (
	udpWg     sync.WaitGroup
	clientMap sync.Map // key: 客户端地址字符串，value: *clientConn，保存客户端对应的目标服务器连接
)

// clientConn 保存目标服务器的 UDP 连接及其状态
type clientConn struct {
	conn       *net.UDPConn
	lastActive time.Time
	mu         sync.Mutex // 保护 conn 读写和 lastActive 字段的并发安全
}

// 写数据到目标 UDP 服务器，线程安全
func (c *clientConn) write(data []byte) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	// 设置写超时时间为2秒
	c.conn.SetWriteDeadline(time.Now().Add(2 * time.Second))
	_, err := c.conn.Write(data)
	return err
}

// 读数据从目标 UDP 服务器，线程安全
func (c *clientConn) read(buffer []byte) (int, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	// 设置读超时时间为2秒
	c.conn.SetReadDeadline(time.Now().Add(2 * time.Second))
	return c.conn.Read(buffer)
}

// 更新连接最后活跃时间，线程安全
func (c *clientConn) updateLastActive() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.lastActive = time.Now()
}

// 处理单个 UDP 客户端的数据转发
func handleUdpTraffic(listener *net.UDPConn, c *clientConn, clientAddr *net.UDPAddr, data []byte) {
	// 更新客户端最后活跃时间
	c.updateLastActive()

	// 将收到的数据写入到目标 UDP 服务器
	if err := c.write(data); err != nil {
		fmt.Printf("[WARN] UDP 转发失败 %s -> %s: %v\n", clientAddr, c.conn.RemoteAddr(), err)
		return
	}

	// 读取目标 UDP 服务器响应
	buffer := make([]byte, 4096)
	n, err := c.read(buffer)
	if err != nil {
		// 如果是超时错误，忽略日志，减少日志噪音
		if netErr, ok := err.(net.Error); !ok || !netErr.Timeout() {
			fmt.Printf("[WARN] UDP 读取目标响应失败: %v\n", err)
		}
		return
	}

	// 将目标服务器响应写回给客户端
	listener.SetWriteDeadline(time.Now().Add(2 * time.Second))
	_, err = listener.WriteToUDP(buffer[:n], clientAddr)
	if err != nil {
		fmt.Printf("[WARN] UDP 发送响应失败: %v\n", err)
	}
}

// 定期清理长时间未活跃的客户端连接
func cleanupInactiveClients(ctx context.Context, interval, maxIdle time.Duration) {
	ticker := time.NewTicker(interval) // 定时器，按 interval 频率执行清理
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			now := time.Now()
			// 遍历所有客户端连接，判断是否超过最大空闲时间
			clientMap.Range(func(key, value interface{}) bool {
				c := value.(*clientConn)

				c.mu.Lock()
				inactive := now.Sub(c.lastActive) > maxIdle
				c.mu.Unlock()

				if inactive {
					// 从 map 删除连接，防止新请求复用已关闭连接
					clientMap.Delete(key)

					fmt.Printf("[INFO] 关闭长时间未活跃客户端连接 %s\n", key)
					c.conn.Close() // 关闭 UDP 连接
				}
				return true
			})
		case <-ctx.Done():
			// 上下文取消，退出清理协程
			return
		}
	}
}

// 启动 UDP 代理服务
func StartUdpProxy(ctx context.Context, listenAddr, targetAddr string) {
	defer udpWg.Done()

	// 解析本地监听 UDP 地址
	localAddr, err := net.ResolveUDPAddr("udp", listenAddr)
	if err != nil {
		fmt.Printf("[ERROR] 解析本地地址失败: %v\n", err)
		return
	}

	// 监听本地 UDP 端口
	listener, err := net.ListenUDP("udp", localAddr)
	if err != nil {
		fmt.Printf("[ERROR] 监听 UDP 失败: %v\n", err)
		return
	}
	defer listener.Close()

	// 解析目标 UDP 服务器地址
	targetUDPAddr, err := net.ResolveUDPAddr("udp", targetAddr)
	if err != nil {
		fmt.Printf("[ERROR] 解析目标地址失败: %v\n", err)
		return
	}

	fmt.Printf("[INFO] UDP 代理启动 %s -> %s\n", listenAddr, targetAddr)

	// 启动后台协程定期清理空闲连接
	go cleanupInactiveClients(ctx, 1*time.Minute, 5*time.Minute)

	// 根据 CPU 核心数启动多个工作协程，提高并发处理能力
	var numWorkers = runtime.NumCPU()
	for i := 0; i < numWorkers; i++ {
		udpWg.Add(1)
		go func(workerID int) {
			defer udpWg.Done()
			buffer := make([]byte, 4096) // 每个协程独立 buffer

			for {
				select {
				case <-ctx.Done():
					// 收到退出信号，结束协程
					return
				default:
					// 设置读取超时，避免阻塞过久
					listener.SetReadDeadline(time.Now().Add(1 * time.Second))
					n, clientAddr, err := listener.ReadFromUDP(buffer)
					if err != nil {
						// 超时错误正常，继续循环
						if ne, ok := err.(net.Error); ok && ne.Timeout() {
							continue
						}

						// 如果上下文已取消，退出协程
						if ctx.Err() != nil {
							return
						}

						// 监听器关闭，退出协程
						if err.Error() == "use of closed network connection" {
							return
						}

						fmt.Printf("[WARN] UDP 读取数据失败: %v\n", err)
						continue
					}

					clientKey := clientAddr.String()
					val, exists := clientMap.Load(clientKey)

					var c *clientConn
					if !exists {
						// 新客户端，创建连接到目标服务器的 UDP 连接
						conn, err := net.DialUDP("udp", nil, targetUDPAddr)
						if err != nil {
							fmt.Printf("[WARN] 连接目标服务器失败: %v\n", err)
							continue
						}
						c = &clientConn{
							conn:       conn,
							lastActive: time.Now(),
						}
						clientMap.Store(clientKey, c)
					} else {
						// 已存在连接，复用
						c = val.(*clientConn)
					}

					// 同步调用，避免因并发过多导致压力大
					handleUdpTraffic(listener, c, clientAddr, buffer[:n])
				}
			}
		}(i)
	}

	// 等待退出信号
	<-ctx.Done()
	fmt.Println("[INFO] 收到退出信号，正在关闭 UDP 代理...")

	listener.Close() // 关闭监听

	// 关闭所有客户端到目标服务器的连接
	clientMap.Range(func(key, value interface{}) bool {
		c := value.(*clientConn)
		c.conn.Close()
		return true
	})

	fmt.Println("[INFO] UDP 代理已关闭")
}
