package l4

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"net"
	"strings"
	"testing"
	"time"
)

// startEchoUpper is a test origin that upper-cases each line it receives.
func startEchoUpper(t *testing.T) string {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go func(c net.Conn) {
				defer c.Close()
				sc := bufio.NewScanner(c)
				for sc.Scan() {
					fmt.Fprintf(c, "%s\n", strings.ToUpper(sc.Text()))
				}
			}(c)
		}
	}()
	return ln.Addr().String()
}

// startProxy boots a Proxy on an ephemeral port and returns it + a cancel.
func startProxy(t *testing.T, p *Proxy) (context.CancelFunc, func()) {
	t.Helper()
	if p.ListenAddr == "" {
		p.ListenAddr = "127.0.0.1:0"
	}
	if err := p.Listen(); err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { _ = p.Serve(ctx); close(done) }()
	return cancel, func() { <-done }
}

func TestProxyForwards(t *testing.T) {
	origin := startEchoUpper(t)
	p := &Proxy{OriginAddr: origin}
	cancel, wait := startProxy(t, p)
	defer wait()
	defer cancel()

	c, err := net.Dial("tcp", p.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	fmt.Fprintf(c, "hello\n")
	line, err := bufio.NewReader(c).ReadString('\n')
	if err != nil {
		t.Fatal(err)
	}
	if got := strings.TrimSpace(line); got != "HELLO" {
		t.Fatalf("got %q, want HELLO", got)
	}
	if p.Metrics.Accepted.Load() < 1 {
		t.Fatalf("Accepted not incremented")
	}
	if p.Metrics.BytesUp.Load() == 0 || p.Metrics.BytesDown.Load() == 0 {
		t.Fatalf("byte counters not moving: up=%d down=%d",
			p.Metrics.BytesUp.Load(), p.Metrics.BytesDown.Load())
	}
}

func TestParseProxyV1(t *testing.T) {
	in := "PROXY TCP4 203.0.113.7 10.0.0.5 56324 8080\r\nrest"
	br := bufio.NewReader(strings.NewReader(in))
	info, err := ParseProxyV1(br)
	if err != nil {
		t.Fatal(err)
	}
	if info.SrcIP.String() != "203.0.113.7" || info.SrcPort != 56324 {
		t.Fatalf("bad src: %v:%d", info.SrcIP, info.SrcPort)
	}
	if info.DstPort != 8080 {
		t.Fatalf("bad dst port: %d", info.DstPort)
	}
	// The reader must still hold the payload after the header.
	rest, _ := io.ReadAll(br)
	if string(rest) != "rest" {
		t.Fatalf("payload not preserved: %q", rest)
	}
}

func TestParseProxyV1Bad(t *testing.T) {
	for _, in := range []string{
		"GET / HTTP/1.1\r\n",
		"PROXY TCP4 1.2.3.4\r\n",
		"PROXY TCP9 1.2.3.4 5.6.7.8 1 2\r\n",
		"PROXY TCP4 not-an-ip 5.6.7.8 1 2\r\n",
		"PROXY TCP4 1.2.3.4 5.6.7.8 99999 2\r\n",
	} {
		br := bufio.NewReader(strings.NewReader(in))
		if _, err := ParseProxyV1(br); err == nil {
			t.Fatalf("expected error for %q", in)
		}
	}
}

func TestProxyProtocolEndToEnd(t *testing.T) {
	// Origin parses the PROXY header and reports the src port back.
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go func(c net.Conn) {
				defer c.Close()
				br := bufio.NewReader(c)
				info, err := ParseProxyV1(br)
				if err != nil {
					fmt.Fprintf(c, "ERR\n")
					return
				}
				fmt.Fprintf(c, "%s:%d\n", info.SrcIP, info.SrcPort)
			}(c)
		}
	}()

	p := &Proxy{OriginAddr: ln.Addr().String(), SendProxyV1: true}
	cancel, wait := startProxy(t, p)
	defer wait()
	defer cancel()

	c, err := net.Dial("tcp", p.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	clientPort := c.LocalAddr().(*net.TCPAddr).Port

	line, err := bufio.NewReader(c).ReadString('\n')
	if err != nil {
		t.Fatal(err)
	}
	want := fmt.Sprintf(":%d", clientPort)
	if !strings.Contains(line, want) {
		t.Fatalf("origin saw %q, expected client port %d", strings.TrimSpace(line), clientPort)
	}
}

func TestHalfClose(t *testing.T) {
	// Origin drains its read side to EOF, THEN replies "DONE". This only
	// works if the proxy propagates the client's FIN (client->origin)
	// without tearing down the origin->client direction.
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go func(c net.Conn) {
				defer c.Close()
				_, _ = io.Copy(io.Discard, c) // read until FIN from client
				fmt.Fprintf(c, "DONE\n")      // reverse direction must still be open
			}(c)
		}
	}()

	p := &Proxy{OriginAddr: ln.Addr().String()}
	cancel, wait := startProxy(t, p)
	defer wait()
	defer cancel()

	c, err := net.Dial("tcp", p.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	fmt.Fprintf(c, "payload")
	c.(*net.TCPConn).CloseWrite() // FIN on client->origin only

	c.SetReadDeadline(time.Now().Add(3 * time.Second))
	line, err := bufio.NewReader(c).ReadString('\n')
	if err != nil {
		t.Fatalf("reverse direction closed prematurely: %v", err)
	}
	if strings.TrimSpace(line) != "DONE" {
		t.Fatalf("got %q, want DONE", line)
	}
}

func TestGracefulDrain(t *testing.T) {
	// Origin holds the connection open (never replies), so the client
	// connection is "in flight" when we cancel.
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			// hold it: read forever
			go func(c net.Conn) { io.Copy(io.Discard, c) }(c)
		}
	}()

	p := &Proxy{OriginAddr: ln.Addr().String(), DrainTimeout: 200 * time.Millisecond}
	cancel, wait := startProxy(t, p)

	c, err := net.Dial("tcp", p.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	fmt.Fprintf(c, "hold open")
	time.Sleep(50 * time.Millisecond)

	// Cancel: Serve must drain (force-close the straggler after the
	// deadline) and return promptly.
	start := time.Now()
	cancel()
	wait() // returns when Serve returns
	if elapsed := time.Since(start); elapsed > 2*time.Second {
		t.Fatalf("drain took too long: %v", elapsed)
	}

	// New connections must be refused after shutdown.
	if c2, err := net.DialTimeout("tcp", p.Addr().String(), 200*time.Millisecond); err == nil {
		c2.Close()
		t.Fatalf("expected new dial to fail after shutdown")
	}
}