package native_test

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"sync/atomic"
	"testing"
	"time"

	omni "git.yuntuo.hk/k2safe/go-omni"
	apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
	"git.yuntuo.hk/k2safe/go-omni/transport/native"
)

func TestApplicationRejectsUnreadBodyPromptly(t *testing.T) {
	for _, framing := range []string{"Content-Length: 1048577\r\n", "Transfer-Encoding: chunked\r\n"} {
		t.Run(framing[:8], func(t *testing.T) {
			var calls atomic.Int32
			server := native.New(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				calls.Add(1)
				w.Header().Set("Content-Type", "application/json")
				w.WriteHeader(http.StatusRequestEntityTooLarge)
				_, _ = io.WriteString(w, `{"code":"BODY_TOO_LARGE"}`)
			}))
			// Much longer than the client's deadline: a draining implementation fails.
			server.ReadTimeout = 10 * time.Second
			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) }()
			defer stopServer(t, server, listener, done)
			conn, err := net.DialTimeout("tcp", listener.Addr().String(), time.Second)
			if err != nil {
				t.Fatal(err)
			}
			defer conn.Close()
			_ = conn.SetDeadline(time.Now().Add(time.Second))
			if _, err = fmt.Fprintf(conn, "POST /upload HTTP/1.1\r\nHost: local\r\n%s\r\n", framing); err != nil {
				t.Fatal(err)
			}
			response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: "POST"})
			if err != nil {
				t.Fatal(err)
			}
			body, err := io.ReadAll(response.Body)
			_ = response.Body.Close()
			if err != nil || response.StatusCode != 413 || !response.Close || string(body) != `{"code":"BODY_TOO_LARGE"}` {
				t.Fatalf("%+v %s %v", response, body, err)
			}
			requirePeerClose(t, conn)
			if calls.Load() != 1 {
				t.Fatalf("handler calls=%d", calls.Load())
			}
		})
	}
}

func TestBodyRejectionWithBusinessHTTP200DoesNotDrain(t *testing.T) {
	app := omni.New(omni.WithErrors(apierrors.Catalog{DefaultStatus: 200, Definitions: map[string]apierrors.Definition{
		"BODY_TOO_LARGE": {Code: "LIMIT_REJECTED", Message: "body limit exceeded"},
	}}))
	app.HandleWithBodyLimit("POST", "/upload", 32, nil, func(*omni.Context) error { t.Error("oversize reached handler"); return nil })
	server := native.New(app)
	server.ReadTimeout = 10 * time.Second
	l, 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(), l) }()
	defer stopServer(t, server, l, done)
	defer app.Close()
	c, err := net.DialTimeout("tcp", l.Addr().String(), time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	c.SetDeadline(time.Now().Add(time.Second))
	// Intentionally withhold every body byte. The rejection must be sent before
	// any attempted drain, despite the project mapping its HTTP status to 200.
	_, err = io.WriteString(c, "POST /upload HTTP/1.1\r\nHost: local\r\nContent-Length: 1024\r\n\r\n")
	if err != nil {
		t.Fatal(err)
	}
	res, err := http.ReadResponse(bufio.NewReader(c), nil)
	if err != nil {
		t.Fatal(err)
	}
	data, err := io.ReadAll(res.Body)
	res.Body.Close()
	if err != nil || res.StatusCode != 200 || !res.Close || !bytes.Contains(data, []byte("LIMIT_REJECTED")) {
		t.Fatalf("status=%d close=%v body=%s err=%v", res.StatusCode, res.Close, data, err)
	}
}