package vless

import (
	"encoding/binary"
	"io"
	"net"
	"sync"
	"time"
)

const udpIdleTimeout = 5 * time.Minute

// relayTCP dials the target and relays data bidirectionally.
// The VLESS response header is written before relaying begins.
func relayTCP(clientConn net.Conn, target string) error {
	remote, err := net.DialTimeout("tcp", target, 10*time.Second)
	if err != nil {
		return err
	}
	defer remote.Close()

	if err := WriteResponse(clientConn); err != nil {
		return err
	}

	var wg sync.WaitGroup
	wg.Add(2)

	// client → remote
	go func() {
		defer wg.Done()
		io.Copy(remote, clientConn)
		remote.(*net.TCPConn).CloseWrite()
	}()

	// remote → client
	go func() {
		defer wg.Done()
		io.Copy(clientConn, remote)
	}()

	wg.Wait()
	return nil
}

// relayUDP handles VLESS UDP-over-TCP relay.
// UDP datagrams are framed as [2B length BE][payload] on the TCP stream.
func relayUDP(clientConn net.Conn, target string) error {
	remoteAddr, err := net.ResolveUDPAddr("udp", target)
	if err != nil {
		return err
	}

	udpConn, err := net.DialUDP("udp", nil, remoteAddr)
	if err != nil {
		return err
	}
	defer udpConn.Close()

	if err := WriteResponse(clientConn); err != nil {
		return err
	}

	var wg sync.WaitGroup
	wg.Add(2)

	// client TCP → remote UDP: read [2B len][payload] frames, send as UDP datagrams.
	go func() {
		defer wg.Done()
		defer udpConn.Close()
		var lenBuf [2]byte
		for {
			if _, err := io.ReadFull(clientConn, lenBuf[:]); err != nil {
				return
			}
			payloadLen := binary.BigEndian.Uint16(lenBuf[:])
			if payloadLen == 0 {
				continue
			}
			buf := make([]byte, payloadLen)
			if _, err := io.ReadFull(clientConn, buf); err != nil {
				return
			}
			udpConn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if _, err := udpConn.Write(buf); err != nil {
				return
			}
		}
	}()

	// remote UDP → client TCP: read UDP datagrams, write as [2B len][payload] frames.
	go func() {
		defer wg.Done()
		buf := make([]byte, 64*1024)
		var lenBuf [2]byte
		for {
			udpConn.SetReadDeadline(time.Now().Add(udpIdleTimeout))
			n, err := udpConn.Read(buf)
			if err != nil {
				return
			}
			binary.BigEndian.PutUint16(lenBuf[:], uint16(n))
			if _, err := clientConn.Write(lenBuf[:]); err != nil {
				return
			}
			if _, err := clientConn.Write(buf[:n]); err != nil {
				return
			}
		}
	}()

	wg.Wait()
	return nil
}