package v5

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"time"

	"github.com/go-gost/core/hosts"
	"github.com/go-gost/core/limiter"
	"github.com/go-gost/core/logger"
	"github.com/go-gost/core/observer/stats"
	"github.com/go-gost/core/resolver"
	"github.com/go-gost/gosocks5"
	ctxvalue "github.com/go-gost/x/ctx"
	ictx "github.com/go-gost/x/internal/ctx"
	xnet "github.com/go-gost/x/internal/net"
	"github.com/go-gost/x/internal/net/udp"
	"github.com/go-gost/x/internal/util/socks"
	traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper"
	metrics "github.com/go-gost/x/metrics/wrapper"
	xstats "github.com/go-gost/x/observer/stats"
	stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
	xrecorder "github.com/go-gost/x/recorder"
)

func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, network string, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
	log = log.WithFields(map[string]any{
		"network": network,
		"cmd":     network,
	})

	if !h.md.enableUDP {
		reply := gosocks5.NewReply(gosocks5.NotAllowed, nil)
		log.Trace(reply)
		log.Error("socks5: UDP relay is disabled")
		return reply.Write(conn)
	}

	// Bind the client-facing relay socket to the TCP control connection's local
	// IP — the address advertised to the client as BND.ADDR — so replies sent
	// from this socket carry the matching source IP. A wildcard bind lets the
	// kernel pick the interface's primary IP, which mismatches BND.ADDR and
	// makes compliant clients drop the replies (RFC 1928 §6). This is what
	// breaks UDP associate under multiple IP aliases (metadata.interface).
	host, _, _ := net.SplitHostPort(conn.LocalAddr().String())

	lc := xnet.ListenConfig{
		Netns: h.options.Netns,
	}

	cc, err := lc.ListenPacket(ctx, network, net.JoinHostPort(host, "0"))
	if err != nil {
		log.Error(err)
		reply := gosocks5.NewReply(gosocks5.Failure, nil)
		log.Trace(reply)
		reply.Write(conn)
		return err
	}
	defer cc.Close()

	log = log.WithFields(map[string]any{
		"src":  cc.LocalAddr().String(),
		"bind": cc.LocalAddr().String(),
	})
	ro.SrcAddr = cc.LocalAddr().String()

	// Verify the upstream chain is reachable before replying Success,
	// per RFC 1928 §7.
	var buf bytes.Buffer
	c, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), network, "") // UDP association
	ro.Route = buf.String()
	if err != nil {
		log.Error(err)
		reply := gosocks5.NewReply(gosocks5.Failure, nil)
		log.Trace(reply)
		reply.Write(conn)
		return err
	}
	defer c.Close()

	pc, ok := c.(net.PacketConn)
	if !ok {
		err := errors.New("socks5: wrong connection type")
		log.Error(err)
		reply := gosocks5.NewReply(gosocks5.Failure, nil)
		log.Trace(reply)
		reply.Write(conn)
		return err
	}

	saddr := gosocks5.Addr{}
	saddr.ParseFrom(cc.LocalAddr().String())

	saddr.Host = host
	if v := net.ParseIP(h.md.publicAddr); v != nil {
		saddr.Host = h.md.publicAddr
	}
	saddr.Type = 0
	reply := gosocks5.NewReply(gosocks5.Succeeded, &saddr)
	log.Trace(reply)
	if err := reply.Write(conn); err != nil {
		log.Error(err)
		return err
	}

	log.Debugf("bind on %s OK", cc.LocalAddr())
	pc = metrics.WrapPacketConn(ro.Service, pc)

	{
		pStats := xstats.Stats{}
		cc = stats_wrapper.WrapPacketConn(cc, &pStats)

		defer func() {
			ro.InputBytes += pStats.Get(stats.KindInputBytes)
			ro.OutputBytes += pStats.Get(stats.KindOutputBytes)
		}()

		clientID := ctxvalue.ClientIDFromContext(ctx)
		cc = traffic_wrapper.WrapPacketConn(
			cc,
			h.limiter,
			string(clientID),
			limiter.ServiceOption(h.options.Service),
			limiter.ScopeOption(limiter.ScopeClient),
			limiter.NetworkOption(network),
			limiter.ClientOption(string(clientID)),
			limiter.SrcOption(conn.RemoteAddr().String()),
		)
		if h.options.Observer != nil {
			pstats := h.stats.Stats(string(clientID))
			pstats.Add(stats.KindTotalConns, 1)
			pstats.Add(stats.KindCurrentConns, 1)
			defer pstats.Add(stats.KindCurrentConns, -1)
			cc = stats_wrapper.WrapPacketConn(cc, pstats)
		}
	}

	tcpAddr, ok := conn.RemoteAddr().(*net.TCPAddr)
	if !ok {
		err := fmt.Errorf("socks5 udp: unexpected remote address type %T", conn.RemoteAddr())
		log.Error(err)
		return err
	}
	pc1 := socks.UDPConn(
		&filteredPacketConn{
			PacketConn: cc,
			filterIP:   tcpAddr.IP,
		},
		h.md.udpBufferSize)

	if h.md.udpResolveDomain {
		pc1 = &domainResolvePacketConn{
			PacketConn: pc1,
			resolver:   h.options.Router.Options().Resolver,
			hostMapper: h.options.Router.Options().HostMapper,
			log:        log,
		}
	}

	r := udp.NewRelay(pc1, pc).
		WithService(h.options.Service).
		WithBypass(h.options.Bypass).
		WithBufferSize(h.md.udpBufferSize).
		WithLogger(log)

	go r.Run(ctx)

	t := time.Now()
	log.Debugf("%s <-> %s", conn.RemoteAddr(), cc.LocalAddr())
	io.Copy(io.Discard, conn)
	log.WithFields(map[string]any{"duration": time.Since(t)}).
		Debugf("%s >-< %s", conn.RemoteAddr(), cc.LocalAddr())

	return nil
}

// filteredPacketConn implements SOCKS5 RFC 1928 UDP relay security by filtering
// incoming packets to only accept those from the client IP that established the TCP control connection.
type filteredPacketConn struct {
	net.PacketConn
	filterIP net.IP // The expected client IP from the SOCKS5 TCP connection
}

func (f *filteredPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
	for {
		n, addr, err = f.PacketConn.ReadFrom(p)
		if err != nil {
			return
		}

		udpAddr, ok := addr.(*net.UDPAddr)
		if !ok {
			continue
		}
		if udpAddr.IP.Equal(f.filterIP) {
			return
		}
	}
}

// domainResolvePacketConn wraps a net.PacketConn and resolves domain
// addresses to IPs in WriteTo calls. This guarantees that SOCKS5 UDP
// response datagrams never carry ATYP=Domain (0x03), which some clients
// (e.g., tun2proxy, Surge) cannot parse.
type domainResolvePacketConn struct {
	net.PacketConn
	resolver   resolver.Resolver
	hostMapper hosts.HostMapper
	log        logger.Logger
}

func (c *domainResolvePacketConn) WriteTo(b []byte, addr net.Addr) (int, error) {
	host, portStr, err := net.SplitHostPort(addr.String())
	if err != nil {
		return c.PacketConn.WriteTo(b, addr)
	}
	if net.ParseIP(host) != nil {
		// Already an IP address — pass through unchanged.
		return c.PacketConn.WriteTo(b, addr)
	}

	// Resolve the domain using the router's infrastructure
	// (host mapper → resolver → system fallback) so the SOCKS5 encoder
	// will use AddrIPv4 or AddrIPv6 instead of AddrDomain.
	var ips []net.IP
	if c.hostMapper != nil {
		ips, _ = c.hostMapper.Lookup(context.Background(), "ip", host)
	}
	if len(ips) == 0 && c.resolver != nil {
		ips, err = c.resolver.Resolve(context.Background(), "ip", host)
		if err != nil && c.log != nil {
			c.log.Warnf("socks5 udp: resolve %s: %v", host, err)
		}
	}
	if len(ips) == 0 {
		ips, err = net.LookupIP(host)
	}
	if err != nil || len(ips) == 0 {
		// Cannot resolve the domain — drop the datagram rather than
		// forwarding an ATYP=Domain response the client cannot parse.
		if c.log != nil {
			c.log.Warnf("socks5 udp: dropping datagram to %s (DNS resolution failed)", host)
		}
		if err == nil {
			err = fmt.Errorf("socks5 udp: DNS resolution failed for %s", host)
		}
		return 0, err
	}

	// Pick the first IPv4 address from the result set, falling back to
	// the first entry when no IPv4 is available.
	ip := ips[0]
	for _, candidate := range ips {
		if candidate.To4() != nil {
			ip = candidate
			break
		}
	}
	port, _ := strconv.Atoi(portStr)
	return c.PacketConn.WriteTo(b, &net.UDPAddr{IP: ip, Port: port})
}