package redirect

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"git.okcu.io/root/drover-go/internal/divert"
	"git.okcu.io/root/drover-go/internal/socks5"
)

// UDPInjector is the minimal subset of *divert.Handle the UDPProxy
// needs to reinject return-path packets. Defined as an interface so
// tests can stub it out without spinning up a real WinDivert handle.
type UDPInjector interface {
	Send(buf []byte, addr UDPInjectAddr) (int, error)
}

// UDPInjectAddr describes the WinDivert addr fields that matter for
// reinjection (we don't need the full 64-byte union here — only flags
// determine direction + checksum status). Production code uses the
// adapter (see DivertHandleInjector) to convert between this and the
// real *idivert.Address.
type UDPInjectAddr struct {
	// Outbound=false → packet will be delivered as inbound (kernel
	// rcv path), which is exactly what we want when fabricating a
	// "remote → local" reply for Discord.
	Outbound bool
}

// UDPConfig configures the UDPProxy.
type UDPConfig struct {
	SOCKS5  socks5.Config
	LocalIP net.IP // local LAN IP we use as the dst on fabricated reply packets

	// Injector is used to reinject return-path packets back to Discord
	// via the WinDivert NETWORK handle. Required.
	Injector UDPInjector

	// LogPrefix is prepended to all log lines emitted by the proxy.
	// Empty defaults to "udp-proxy: ".
	LogPrefix string
}

// udpFlow tracks one (Discord_src → real_dst) UDP flow for the
// purpose of routing relay responses back to Discord.
type udpFlow struct {
	// realDst* identifies the upstream UDP target (the same key the
	// SOCKS5 relay puts in DST.ADDR/DST.PORT on the inbound envelope).
	realDstIP   [4]byte
	realDstPort uint16

	// discordSrc* identifies the Discord side of the flow — used as
	// the dst on fabricated reply packets so the kernel matches the
	// connect()-bound socket.
	discordSrcIP   [4]byte
	discordSrcPort uint16

	lastUsed time.Time
}

// UDPProxy is the SOCKS5 UDP relay manager. The engine's diverterLoop
// calls Forward on outbound UDP packets from target processes; the
// proxy lazily opens a single UDP ASSOCIATE control TCP + relay UDP
// socket on first use, and shares them across all UDP flows. Inbound
// responses are read from the relay socket, decap'd, and reinjected
// as fabricated IPv4+UDP packets via the WinDivert NETWORK handle.
type UDPProxy struct {
	cfg UDPConfig

	// Lazy-opened on first Forward call.
	ctrlMu    sync.Mutex
	ctrlConn  net.Conn       // SOCKS5 control TCP — must stay open for relay validity
	relayAddr *net.UDPAddr   // upstream relay UDP endpoint
	relayConn net.PacketConn // local UDP socket bound to talk to relay

	flowMu sync.RWMutex
	// Keyed by realDstIP:realDstPort — the relay responds with these
	// in the SOCKS5 envelope, so this is our reverse lookup.
	flowsByDst map[flowDstKey]*udpFlow

	// Atomic stats counters for diagnostics
	fwdPackets      uint64
	fwdBytes        uint64
	recvPackets     uint64
	injectedPackets uint64

	wg     sync.WaitGroup
	ctx    context.Context
	cancel context.CancelFunc

	// Idle TTL for udpFlow entries (default 5 minutes per RFC 4787).
	IdleTTL time.Duration
}

type flowDstKey struct {
	ip   [4]byte
	port uint16
}

// NewUDP constructs a UDPProxy. It does not yet open the SOCKS5 UDP
// ASSOCIATE — that happens lazily on the first Forward call.
func NewUDP(cfg UDPConfig) (*UDPProxy, error) {
	if cfg.Injector == nil {
		return nil, errors.New("UDPConfig.Injector is required")
	}
	if cfg.LocalIP == nil || cfg.LocalIP.To4() == nil {
		return nil, errors.New("UDPConfig.LocalIP must be IPv4")
	}
	if cfg.LogPrefix == "" {
		cfg.LogPrefix = "udp-proxy: "
	}
	ctx, cancel := context.WithCancel(context.Background())
	u := &UDPProxy{
		cfg:        cfg,
		flowsByDst: map[flowDstKey]*udpFlow{},
		ctx:        ctx,
		cancel:     cancel,
		IdleTTL:    5 * time.Minute,
	}
	u.wg.Add(1)
	go u.sweepLoop()
	return u, nil
}

// Forward is called from the engine's diverterLoop on each outbound
// UDP packet from a target process. It:
//
//  1. Lazy-opens the SOCKS5 UDP association on first call.
//  2. Records the flow keyed by (dstIP,dstPort) so the relay-response
//     reader can route the reply back to the right Discord port.
//  3. Encapsulates the payload in a SOCKS5 UDP datagram (RFC 1928 §7)
//     and forwards it to the relay endpoint.
//
// Returns nil on success or any error encountered (caller may log
// but should generally drop the packet on failure — UDP loss is
// expected at the wire).
func (u *UDPProxy) Forward(srcIP net.IP, srcPort uint16, dstIP net.IP, dstPort uint16, payload []byte) error {
	srcV4 := srcIP.To4()
	dstV4 := dstIP.To4()
	if srcV4 == nil || dstV4 == nil {
		return errors.New("UDPProxy.Forward: src/dst must be IPv4")
	}

	if err := u.ensureAssociated(); err != nil {
		return fmt.Errorf("ensure assoc: %w", err)
	}

	// Record/refresh flow for the return path
	var dKey flowDstKey
	copy(dKey.ip[:], dstV4)
	dKey.port = dstPort

	u.flowMu.Lock()
	fl, ok := u.flowsByDst[dKey]
	if !ok {
		fl = &udpFlow{}
		u.flowsByDst[dKey] = fl
	}
	copy(fl.realDstIP[:], dstV4)
	fl.realDstPort = dstPort
	copy(fl.discordSrcIP[:], srcV4)
	fl.discordSrcPort = srcPort
	fl.lastUsed = time.Now()
	u.flowMu.Unlock()

	// Encap and send to relay
	envelope, err := socks5.EncapUDPv4(dstIP, dstPort, payload)
	if err != nil {
		return fmt.Errorf("encap: %w", err)
	}
	n, err := u.relayConn.WriteTo(envelope, u.relayAddr)
	if err != nil {
		return fmt.Errorf("write to relay: %w", err)
	}
	atomic.AddUint64(&u.fwdPackets, 1)
	atomic.AddUint64(&u.fwdBytes, uint64(n))
	return nil
}

// Stats returns counters for diagnostics: forwarded outbound packets,
// inbound packets received from relay, inbound packets successfully
// reinjected to Discord. All atomic; safe to read concurrently.
func (u *UDPProxy) Stats() (fwdPkts, fwdBytes, recvPkts, injectedPkts uint64) {
	return atomic.LoadUint64(&u.fwdPackets),
		atomic.LoadUint64(&u.fwdBytes),
		atomic.LoadUint64(&u.recvPackets),
		atomic.LoadUint64(&u.injectedPackets)
}

// ensureAssociated opens the SOCKS5 UDP association on first use and
// reuses it forever (until Close). The relay endpoint stays valid as
// long as the control TCP is open, per RFC 1928 §6.
func (u *UDPProxy) ensureAssociated() error {
	u.ctrlMu.Lock()
	defer u.ctrlMu.Unlock()
	if u.ctrlConn != nil && u.relayAddr != nil && u.relayConn != nil {
		return nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	relay, ctrl, err := socks5.AssociateUDP(ctx, u.cfg.SOCKS5)
	if err != nil {
		return err
	}

	// Bind a local UDP socket to talk to the relay. Bind on 0.0.0.0:0
	// so the kernel picks an ephemeral port; we'll use this socket as
	// both the writer (Forward) AND the reader (relayReadLoop).
	pc, err := net.ListenPacket("udp4", ":0")
	if err != nil {
		ctrl.Close()
		return fmt.Errorf("listen relay socket: %w", err)
	}

	u.ctrlConn = ctrl
	u.relayAddr = relay
	u.relayConn = pc

	log.Printf("%sSOCKS5 UDP ASSOCIATE relay=%s local=%s", u.cfg.LogPrefix, relay, pc.LocalAddr())

	// Spawn the reader goroutine.
	u.wg.Add(1)
	go u.relayReadLoop()

	// Spawn a control-conn watcher: if the proxy closes the control
	// TCP for any reason, our relay endpoint is invalidated. Mark
	// state for re-association on next Forward.
	u.wg.Add(1)
	go u.ctrlWatcher()
	return nil
}

func (u *UDPProxy) ctrlWatcher() {
	defer u.wg.Done()
	// Read forever from ctrlConn; per RFC 1928 §6 the proxy doesn't
	// send anything on this conn after the UDP ASSOCIATE reply, so
	// any read-completion (with or without bytes) means the conn is
	// gone. This is a fire-and-forget watcher — it doesn't actively
	// re-associate; ensureAssociated() will do that on next Forward.
	one := make([]byte, 1)
	for {
		// Use a generous read deadline so we wake up periodically to
		// honor ctx cancellation.
		_ = u.ctrlConn.SetReadDeadline(time.Now().Add(30 * time.Second))
		_, err := u.ctrlConn.Read(one)
		if err == nil {
			continue // unexpected data; keep monitoring
		}
		if ne, ok := err.(net.Error); ok && ne.Timeout() {
			select {
			case <-u.ctx.Done():
				return
			default:
			}
			continue
		}
		// Real error — control conn is dead. Tear down so next Forward
		// re-associates.
		log.Printf("%scontrol TCP closed: %v — relay invalidated", u.cfg.LogPrefix, err)
		u.ctrlMu.Lock()
		if u.ctrlConn != nil {
			u.ctrlConn.Close()
			u.ctrlConn = nil
		}
		if u.relayConn != nil {
			u.relayConn.Close()
			u.relayConn = nil
		}
		u.relayAddr = nil
		u.ctrlMu.Unlock()
		return
	}
}

// relayReadLoop reads inbound datagrams from the relay socket.
// Datagrams from the relay are SOCKS5 UDP envelopes (RFC 1928 §7);
// we decap, look up the corresponding Discord flow by the envelope's
// DST.ADDR/DST.PORT (which contains the ORIGIN of the response), and
// reinject a fabricated IPv4+UDP packet as inbound via WinDivert.
func (u *UDPProxy) relayReadLoop() {
	defer u.wg.Done()

	buf := make([]byte, 65535)
	for {
		select {
		case <-u.ctx.Done():
			return
		default:
		}
		// Snapshot relay conn under lock; if torn down by ctrlWatcher
		// we need to bail out.
		u.ctrlMu.Lock()
		pc := u.relayConn
		relay := u.relayAddr
		u.ctrlMu.Unlock()
		if pc == nil {
			return
		}

		_ = pc.SetReadDeadline(time.Now().Add(2 * time.Second))
		n, fromAddr, err := pc.ReadFrom(buf)
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Timeout() {
				continue
			}
			// Likely closed — exit.
			return
		}
		atomic.AddUint64(&u.recvPackets, 1)

		// Sanity-check source: relay datagrams come from the relay's
		// known address. Ignore anything else (in particular some
		// SOCKS5 implementations bind 0.0.0.0; we accept any port match
		// loosely, but require IP match when available).
		fromUDP, ok := fromAddr.(*net.UDPAddr)
		if !ok {
			continue
		}
		if relay != nil && relay.IP != nil && !relay.IP.Equal(net.IPv4zero) {
			if !fromUDP.IP.Equal(relay.IP) || fromUDP.Port != relay.Port {
				// Not from our relay — drop.
				continue
			}
		}

		srcIP, srcPort, payload, derr := socks5.DecapUDPv4(buf[:n])
		if derr != nil {
			log.Printf("%sdecap error: %v", u.cfg.LogPrefix, derr)
			continue
		}

		// Look up the Discord flow by (origin IP, origin port)
		v4 := srcIP.To4()
		if v4 == nil {
			continue
		}
		var key flowDstKey
		copy(key.ip[:], v4)
		key.port = srcPort

		u.flowMu.RLock()
		fl, ok := u.flowsByDst[key]
		u.flowMu.RUnlock()
		if !ok {
			// No active flow for this origin; drop.
			continue
		}

		// Mark the flow as recently used (touched by inbound).
		u.flowMu.Lock()
		fl.lastUsed = time.Now()
		u.flowMu.Unlock()

		// Fabricate IPv4+UDP packet:
		//   src = real_origin (the proxy's relay tells us this in the envelope)
		//   dst = local LAN IP we bound on
		//   srcPort = real origin port
		//   dstPort = Discord ephemeral port (so kernel matches the connect()-bound socket)
		discordIP := net.IPv4(fl.discordSrcIP[0], fl.discordSrcIP[1], fl.discordSrcIP[2], fl.discordSrcIP[3])
		// Some Discord sockets bind to local LAN IP, others bind 0.0.0.0
		// (which the SOCKET layer reports as 0.0.0.0). When discord's
		// reported srcIP is 0.0.0.0 the kernel's connect-bound socket
		// will still match dst=our LocalIP. But to be safe for the
		// non-zero case (sockets bound to specific local IP), use the
		// recorded discord side IP if it is non-zero; otherwise fall
		// back to LocalIP.
		dstIP := discordIP
		if discordIP.Equal(net.IPv4zero) {
			dstIP = u.cfg.LocalIP
		}

		pkt, berr := divert.BuildIPv4UDPInbound(srcIP, dstIP, srcPort, fl.discordSrcPort, payload)
		if berr != nil {
			log.Printf("%sbuild packet error: %v", u.cfg.LogPrefix, berr)
			continue
		}

		// Reinject as inbound. WinDivert flag bits we set: IPChecksum
		// (we computed it), UDPChecksum (we computed it). Outbound bit
		// stays clear — kernel delivers via inbound path.
		if _, serr := u.cfg.Injector.Send(pkt, UDPInjectAddr{Outbound: false}); serr != nil {
			log.Printf("%sinject error: %v", u.cfg.LogPrefix, serr)
		} else {
			atomic.AddUint64(&u.injectedPackets, 1)
		}
	}
}

// sweepLoop garbage-collects stale udpFlow entries. UDP "flows" are
// stateless — there's no FIN-equivalent — so we rely on idle timeout.
// 5 minutes matches RFC 4787 NAT requirements (REQ-5).
func (u *UDPProxy) sweepLoop() {
	defer u.wg.Done()
	tk := time.NewTicker(time.Minute)
	defer tk.Stop()
	for {
		select {
		case <-u.ctx.Done():
			return
		case <-tk.C:
			cutoff := time.Now().Add(-u.IdleTTL)
			u.flowMu.Lock()
			for k, f := range u.flowsByDst {
				if f.lastUsed.Before(cutoff) {
					delete(u.flowsByDst, k)
				}
			}
			u.flowMu.Unlock()
		}
	}
}

// Close tears down the UDPProxy: cancels reader goroutines, closes
// the relay UDP socket and the SOCKS5 control TCP. Safe to call
// multiple times.
func (u *UDPProxy) Close() error {
	u.cancel()
	u.ctrlMu.Lock()
	if u.relayConn != nil {
		_ = u.relayConn.Close()
		u.relayConn = nil
	}
	if u.ctrlConn != nil {
		_ = u.ctrlConn.Close()
		u.ctrlConn = nil
	}
	u.relayAddr = nil
	u.ctrlMu.Unlock()
	u.wg.Wait()
	return nil
}

// FlowCount returns the current number of tracked UDP flows. Test
// helper.
func (u *UDPProxy) FlowCount() int {
	u.flowMu.RLock()
	defer u.flowMu.RUnlock()
	return len(u.flowsByDst)
}