package native_test

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"testing"
	"time"

	omni "git.yuntuo.hk/k2safe/go-omni"
	"git.yuntuo.hk/k2safe/go-omni/transport/native"
	netrouter "git.yuntuo.hk/k2safe/go-omni/transport/nethttp"
)

func TestServerKeepAliveAndApplicationBridge(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/hello", func(ctx *omni.Context) error {
		return ctx.OK(map[string]string{"message": "hello"})
	})
	server, listener, done := startServer(t, application)
	defer stopServer(t, server, listener, done)

	conn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	reader := bufio.NewReader(conn)

	for index, closeAfter := range []bool{false, true} {
		connectionHeader := "keep-alive"
		if closeAfter {
			connectionHeader = "close"
		}
		fmt.Fprintf(conn, "GET /hello HTTP/1.1\r\nHost: test\r\nConnection: %s\r\n\r\n", connectionHeader)
		request := &http.Request{Method: http.MethodGet}
		response, err := http.ReadResponse(reader, request)
		if err != nil {
			t.Fatalf("response %d: %v", index, err)
		}
		body, err := io.ReadAll(response.Body)
		response.Body.Close()
		if err != nil || response.StatusCode != http.StatusOK || !strings.Contains(string(body), "\"message\":\"hello\"") {
			t.Fatalf("response %d: status=%d body=%s error=%v", index, response.StatusCode, body, err)
		}
	}
}

func TestServerRejectsUnsupportedBodyAndLimits(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.POST("/body", func(ctx *omni.Context) error { return ctx.OK("ok") })

	server := native.New(application)
	server.MaxBodyBytes = 4
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	waitForServer(t, listener.Addr().String())
	defer stopServer(t, server, listener, done)

	for _, requestText := range []string{
		"POST /body HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\n\r\n4\r\ntest\r\n0\r\n\r\n",
		"POST /body HTTP/1.1\r\nHost: test\r\nContent-Length: 5\r\n\r\ntest!",
	} {
		conn, err := net.Dial("tcp", listener.Addr().String())
		if err != nil {
			t.Fatal(err)
		}
		_, _ = io.WriteString(conn, requestText)
		response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: http.MethodPost})
		if err != nil {
			conn.Close()
			t.Fatal(err)
		}
		response.Body.Close()
		if response.StatusCode != http.StatusNotImplemented && response.StatusCode != http.StatusRequestEntityTooLarge {
			t.Fatalf("unexpected protocol status %d", response.StatusCode)
		}
		conn.Close()
	}
}

func TestServerShutdownCancelsActiveHandler(t *testing.T) {
	started := make(chan struct{})
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/wait", func(ctx *omni.Context) error {
		close(started)
		<-ctx.Context().Done()
		return ctx.Context().Err()
	})
	server, listener, done := startServer(t, application)

	conn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	_, _ = io.WriteString(conn, "GET /wait HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("handler did not start")
	}
	shutdownContext, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := server.Shutdown(shutdownContext); err != nil {
		t.Fatal(err)
	}
	response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusInternalServerError {
		t.Fatalf("expected canceled handler to become 500, got %d", response.StatusCode)
	}
	conn.Close()
	select {
	case err := <-done:
		if err != nil && !errors.Is(err, net.ErrClosed) {
			t.Fatal(err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("server did not stop")
	}
}

func TestServerConcurrentRequestLimit(t *testing.T) {
	started := make(chan struct{})
	release := make(chan struct{})
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/work", func(ctx *omni.Context) error {
		close(started)
		<-release
		return ctx.OK("done")
	})
	server := native.New(application)
	server.MaxConcurrentRequests = 1
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	waitForServer(t, listener.Addr().String())
	defer func() {
		stopServer(t, server, listener, done)
	}()

	first, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer first.Close()
	_, _ = io.WriteString(first, "GET /work HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("first handler did not start")
	}

	second, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer second.Close()
	_, _ = io.WriteString(second, "GET /work HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	response, err := http.ReadResponse(bufio.NewReader(second), &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("expected 503 at concurrency limit, got %d", response.StatusCode)
	}
	close(release)
	response, err = http.ReadResponse(bufio.NewReader(first), &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusOK {
		t.Fatalf("expected first request to complete, got %d", response.StatusCode)
	}
}

func TestServerReadHeaderTimeout(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/health", func(ctx *omni.Context) error { return ctx.OK("ok") })
	server := native.New(application)
	server.ReadHeaderTimeout = 25 * time.Millisecond
	server.IdleTimeout = 0
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	waitForServer(t, listener.Addr().String())
	defer stopServer(t, server, listener, done)

	conn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	conn.SetReadDeadline(time.Now().Add(time.Second))
	_, _ = io.WriteString(conn, "GET /health HTTP/1.1\r\nHost: test")
	buffer := make([]byte, 1)
	if _, err := conn.Read(buffer); err == nil {
		t.Fatal("expected incomplete headers to time out and close")
	}
}

func TestAppRunWithNativeRunner(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/health", func(ctx *omni.Context) error { return ctx.OK("ok") })
	address := freeAddress(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- application.RunWith(ctx, native.New(nil), address) }()
	waitForServer(t, address)
	response, err := http.Get("http://" + address + "/health")
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusOK {
		t.Fatalf("unexpected status %d", response.StatusCode)
	}
	cancel()
	select {
	case err := <-done:
		if err != nil {
			t.Fatal(err)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("RunWith did not stop")
	}
}

func startServer(t *testing.T, handler http.Handler) (*native.Server, net.Listener, <-chan error) {
	t.Helper()
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	server := native.New(handler)
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	waitForServer(t, listener.Addr().String())
	return server, listener, done
}

func stopServer(t *testing.T, server *native.Server, listener net.Listener, done <-chan error) {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	_ = server.Shutdown(ctx)
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("server did not stop")
	}
	_ = listener.Close()
}

func freeAddress(t *testing.T) string {
	t.Helper()
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	address := listener.Addr().String()
	listener.Close()
	return address
}

func waitForServer(t *testing.T, address string) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for {
		conn, err := net.DialTimeout("tcp", address, 50*time.Millisecond)
		if err == nil {
			conn.Close()
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("server did not start: %v", err)
		}
		time.Sleep(10 * time.Millisecond)
	}
}