// Package l4 implements a small but production-shaped Layer-4 TCP proxy:
// an accept loop, bidirectional copy with backpressure and half-close,
// optional PROXY-protocol v1 emission, connection tracking, and graceful
// drain with a deadline.
//
// It is deliberately stdlib-only. Go's runtime netpoller is itself an
// epoll/kqueue event loop, so the idiomatic "two goroutines per
// connection" shape here compiles down to exactly the event-driven I/O
// you would write by hand against epoll in C — without the ceremony.
package l4

import (
	"context"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

// Metrics are the counters a real L4 proxy exports. BytesUp is
// client->origin; BytesDown is origin->client. Created/Active/Closed let
// you watch connection lifecycle (and, with a sampler, connection
// churn — see gw-04).
type Metrics struct {
	Accepted  atomic.Int64
	Active    atomic.Int64
	Closed    atomic.Int64
	DialFails atomic.Int64
	BytesUp   atomic.Int64
	BytesDown atomic.Int64
}

// Proxy forwards every accepted client connection to a single origin.
type Proxy struct {
	// ListenAddr is the bind address; default ":8080". Use "127.0.0.1:0"
	// in tests to get an ephemeral port (read it back via Addr()).
	ListenAddr string
	// OriginAddr is the single upstream "host:port" we forward to.
	OriginAddr string
	// SendProxyV1, when true, prepends a PROXY-protocol v1 header to the
	// origin connection so the origin learns the real client address.
	SendProxyV1 bool
	// DialTimeout bounds how long we wait to connect to the origin.
	DialTimeout time.Duration
	// DrainTimeout bounds graceful drain on ctx cancellation. Zero means
	// wait forever for in-flight connections to finish.
	DrainTimeout time.Duration

	Metrics Metrics

	mu     sync.Mutex
	ln     net.Listener
	active map[net.Conn]struct{} // for force-close on drain deadline
	wg     sync.WaitGroup
}

// Listen binds the listener. Call before Serve; lets tests read Addr().
func (p *Proxy) Listen() error {
	addr := p.ListenAddr
	if addr == "" {
		addr = ":8080"
	}
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	p.mu.Lock()
	p.ln = ln
	p.active = map[net.Conn]struct{}{}
	p.mu.Unlock()
	return nil
}

// Addr returns the bound address (valid after Listen).
func (p *Proxy) Addr() net.Addr {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.ln == nil {
		return nil
	}
	return p.ln.Addr()
}

// Run is Listen + Serve.
func (p *Proxy) Run(ctx context.Context) error {
	if err := p.Listen(); err != nil {
		return err
	}
	return p.Serve(ctx)
}

// Serve runs the accept loop until ctx is cancelled, then drains.
//
// Ordering matters: cancelling ctx closes the listener (so no NEW
// connections), then we wait for in-flight connections up to
// DrainTimeout. In Kubernetes you would ALSO fail the readiness probe
// before cancelling so the load balancer removes you from rotation
// first (see gw-09).
func (p *Proxy) Serve(ctx context.Context) error {
	p.mu.Lock()
	ln := p.ln
	p.mu.Unlock()
	if ln == nil {
		panic("l4: Serve called before Listen")
	}

	// Closing the listener unblocks Accept on cancellation.
	go func() { <-ctx.Done(); ln.Close() }()

	for {
		conn, err := ln.Accept()
		if err != nil {
			select {
			case <-ctx.Done():
				p.drain()
				return nil
			default:
				// Transient accept error (e.g. EMFILE). A real proxy
				// would back off briefly; we just continue.
				continue
			}
		}
		p.Metrics.Accepted.Add(1)
		p.wg.Add(1)
		go func() {
			defer p.wg.Done()
			p.handle(conn)
		}()
	}
}

func (p *Proxy) dialTimeout() time.Duration {
	if p.DrainTimeout < 0 || p.DialTimeout == 0 {
		return 2 * time.Second
	}
	return p.DialTimeout
}

func (p *Proxy) handle(client net.Conn) {
	p.track(client)
	defer p.untrack(client)
	p.Metrics.Active.Add(1)
	defer func() { p.Metrics.Active.Add(-1); p.Metrics.Closed.Add(1) }()
	defer client.Close()

	origin, err := net.DialTimeout("tcp", p.OriginAddr, p.dialTimeout())
	if err != nil {
		p.Metrics.DialFails.Add(1)
		return
	}
	p.track(origin)
	defer p.untrack(origin)
	defer origin.Close()

	// TCP_NODELAY on both sides: disable Nagle so small request/response
	// messages aren't delayed waiting to coalesce (the 40ms-stall trap).
	setNoDelay(client)
	setNoDelay(origin)

	if p.SendProxyV1 {
		if err := WriteProxyV1(origin, client); err != nil {
			return
		}
	}

	// Two coupled copies. The copy blocks on a slow Write, which stops the
	// corresponding Read: that IS backpressure, with a bounded buffer.
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		pipe(origin, client, &p.Metrics.BytesUp) // client -> origin
	}()
	go func() {
		defer wg.Done()
		pipe(client, origin, &p.Metrics.BytesDown) // origin -> client
	}()
	wg.Wait()
}

// pipe copies src -> dst, updating counter as it goes (so live metrics
// move during long-lived connections, not only at close), then
// half-closes dst's write side so the peer sees a FIN on this direction
// while the reverse direction stays open.
func pipe(dst, src net.Conn, counter *atomic.Int64) {
	buf := make([]byte, 32*1024) // bounded buffer == bounded memory == backpressure
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[:nr])
			if nw > 0 {
				counter.Add(int64(nw))
			}
			if ew != nil || nw < nr {
				break
			}
		}
		if er != nil { // includes io.EOF (peer sent FIN)
			break
		}
	}
	if cw, ok := dst.(interface{ CloseWrite() error }); ok {
		_ = cw.CloseWrite()
	}
}

func setNoDelay(c net.Conn) {
	if tc, ok := c.(*net.TCPConn); ok {
		_ = tc.SetNoDelay(true)
	}
}

func (p *Proxy) track(c net.Conn) {
	p.mu.Lock()
	if p.active != nil {
		p.active[c] = struct{}{}
	}
	p.mu.Unlock()
}

func (p *Proxy) untrack(c net.Conn) {
	p.mu.Lock()
	delete(p.active, c)
	p.mu.Unlock()
}

// drain waits for in-flight connections, bounded by DrainTimeout. On the
// deadline it force-closes the stragglers (which unblocks their copy
// goroutines).
func (p *Proxy) drain() {
	done := make(chan struct{})
	go func() { p.wg.Wait(); close(done) }()

	if p.DrainTimeout <= 0 {
		<-done
		return
	}
	select {
	case <-done:
	case <-time.After(p.DrainTimeout):
		p.forceCloseAll()
		<-done
	}
}

func (p *Proxy) forceCloseAll() {
	p.mu.Lock()
	conns := make([]net.Conn, 0, len(p.active))
	for c := range p.active {
		conns = append(conns, c)
	}
	p.mu.Unlock()
	for _, c := range conns {
		_ = c.Close()
	}
}