package proxy

import (
	"fmt"
	"log"
	"net"
	"sync"
	"time"
)

// udpSession tracks the state of a single client's UDP "connection"
type udpSession struct {
	clientAddr  *net.UDPAddr
	backendConn *net.UDPConn
	lastActive  time.Time
}

// RunUDPProxy starts a UDP reverse proxy, establishing a NAT-like session mapping for clients
func RunUDPProxy(port int, targetBackend string) error {
	addr := fmt.Sprintf(":%d", port)
	proxyAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return err
	}

	backendAddr, err := net.ResolveUDPAddr("udp", targetBackend)
	if err != nil {
		return err
	}

	proxyConn, err := net.ListenUDP("udp", proxyAddr)
	if err != nil {
		return err
	}
	defer proxyConn.Close()

	log.Printf("Starting UDP proxy on %s, forwarding to %s\n", addr, targetBackend)

	var mu sync.RWMutex
	sessions := make(map[string]*udpSession)

	// Background routine to clean up inactive sessions
	go func() {
		for {
			time.Sleep(1 * time.Minute)
			mu.Lock()
			for k, s := range sessions {
				// If a client hasn't sent or received packets in 3 minutes, close it
				if time.Since(s.lastActive) > 3*time.Minute {
					s.backendConn.Close()
					delete(sessions, k)
				}
			}
			mu.Unlock()
		}
	}()

	buffer := make([]byte, 65535)
	for {
		n, clientAddr, err := proxyConn.ReadFromUDP(buffer)
		if err != nil {
			log.Printf("UDP Read error: %v", err)
			continue
		}

		clientKey := clientAddr.String()
		mu.RLock()
		session, exists := sessions[clientKey]
		mu.RUnlock()

		if !exists {
			// New client session detected, dial the backend
			bConn, err := net.DialUDP("udp", nil, backendAddr)
			if err != nil {
				log.Printf("Failed to connect to UDP backend: %v", err)
				continue
			}

			session = &udpSession{
				clientAddr:  clientAddr,
				backendConn: bConn,
				lastActive:  time.Now(),
			}

			mu.Lock()
			sessions[clientKey] = session
			mu.Unlock()

			// Start a goroutine to read from the backend and send back to this specific client
			go func(s *udpSession) {
				bBuf := make([]byte, 65535)
				for {
					bn, err := s.backendConn.Read(bBuf)
					if err != nil {
						break // Connection closed or error
					}
					s.lastActive = time.Now()
					proxyConn.WriteToUDP(bBuf[:bn], s.clientAddr)
				}
			}(session)
		}

		// Forward the client's packet to the backend
		session.lastActive = time.Now()
		session.backendConn.Write(buffer[:n])
	}
}