package app

import (
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/config"
	"git.yuntuo.hk/k2safe/go-omni/httpx"
	"git.yuntuo.hk/k2safe/go-omni/internal/testutil"
	"git.yuntuo.hk/k2safe/go-omni/lifecycle"
)

func TestShutdownWaitsForCancellationCleanup(t *testing.T) {
	for _, transport := range []string{config.TransportNetHTTP, config.TransportNative, config.TransportHTTP2} {
		t.Run(transport, func(t *testing.T) {
			cert, key, tlsClient := testutil.CertificateFiles(t)
			a := New(config.WithTransport(transport), config.WithTLS(cert, key), config.WithShutdownTimeout(30*time.Millisecond), config.WithHTTPConfig(config.HTTPConfig{ShutdownCancelTimeout: time.Second}))
			started, canceled, release := make(chan struct{}), make(chan struct{}), make(chan struct{})
			var once sync.Once
			unblock := func() { once.Do(func() { close(release) }) }
			defer unblock()
			var closed atomic.Int32
			a.OnShutdown(func() error { closed.Add(1); return nil })
			a.GET("/work", func(c *httpx.Context) error {
				close(started)
				<-c.Context().Done()
				close(canceled)
				<-release
				if closed.Load() != 0 {
					t.Error("dependency closed during cancellation cleanup")
				}
				return nil
			})
			l, err := net.Listen("tcp", "127.0.0.1:0")
			if err != nil {
				t.Fatal(err)
			}
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			done := make(chan error, 1)
			go func() { done <- a.ServeContext(ctx, l) }()
			client := &http.Client{Timeout: 3 * time.Second, Transport: &http.Transport{TLSClientConfig: tlsClient, ForceAttemptHTTP2: true}}
			defer client.CloseIdleConnections()
			requestDone := make(chan struct{})
			go func() {
				defer close(requestDone)
				res, err := client.Get("https://" + l.Addr().String() + "/work")
				if err == nil {
					io.Copy(io.Discard, res.Body)
					res.Body.Close()
				}
			}()
			await(t, started)
			cancel()
			await(t, canceled)
			if closed.Load() != 0 {
				t.Fatal("dependency closed on cancellation")
			}
			select {
			case err := <-done:
				t.Fatal("returned before handler cleanup", err)
			default:
			}
			// Even an in-memory invocation through Engine must now be rejected.
			w := httptest.NewRecorder()
			a.Engine().ServeHTTP(w, httptest.NewRequest("GET", "/work", nil))
			if w.Code != 503 {
				t.Fatal(w.Code)
			}
			unblock()
			select {
			case err := <-done:
				if !errors.Is(err, context.DeadlineExceeded) || errors.Is(err, lifecycle.ErrResourcesRetained) {
					t.Fatal(err)
				}
			case <-time.After(3 * time.Second):
				t.Fatal("shutdown stuck")
			}
			await(t, requestDone)
			if closed.Load() != 1 {
				t.Fatal("cleanup count", closed.Load())
			}
		})
	}
}

func TestExternalServingEntriesRetainActiveResources(t *testing.T) {
	for _, entry := range []string{"App", "Engine", "NewServer", "RunWith"} {
		t.Run(entry, func(t *testing.T) {
			a := New(config.WithHTTPConfig(config.HTTPConfig{ShutdownCancelTimeout: 30 * time.Millisecond}))
			started, release, returned := make(chan struct{}), make(chan struct{}), make(chan struct{})
			var closed atomic.Int32
			a.OnShutdown(func() error { closed.Add(1); return nil })
			a.GET("/", func(*httpx.Context) error {
				close(started)
				<-release
				if closed.Load() != 0 {
					t.Error("closed while active")
				}
				return nil
			})
			var h http.Handler = a
			if entry == "Engine" {
				h = a.Engine()
			}
			if entry == "NewServer" {
				h = a.NewServer("").Handler
			}
			invoke := func(h http.Handler) {
				defer close(returned)
				h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest("GET", "/", nil))
			}
			defer func() { close(release); await(t, returned) }()
			var err error
			if entry == "RunWith" {
				err = a.RunWith(context.Background(), runnerFunc(func(_ context.Context, h http.Handler, _ string) error { go invoke(h); await(t, started); return nil }))
			} else {
				go invoke(h)
				await(t, started)
				err = a.Close()
			}
			if !errors.Is(err, lifecycle.ErrResourcesRetained) || closed.Load() != 0 {
				t.Fatal("retention failed", err, closed.Load())
			}
			if err = a.Close(); !errors.Is(err, lifecycle.ErrResourcesRetained) {
				t.Fatal("repeat close lost retention", err)
			}
		})
	}
}

func TestRequestAdmissionAndCloseRace(t *testing.T) {
	for attempt := 0; attempt < 20; attempt++ {
		a := New()
		var closed atomic.Bool
		a.OnShutdown(func() error { closed.Store(true); return nil })
		a.GET("/", func(c *httpx.Context) error {
			if closed.Load() {
				t.Error("handler admitted after resource release")
			}
			return nil
		})
		start := make(chan struct{})
		var wg sync.WaitGroup
		for i := 0; i < 20; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				<-start
				a.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest("GET", "/", nil))
			}()
		}
		close(start)
		if err := a.Close(); err != nil {
			t.Fatal(err)
		}
		wg.Wait()
	}
}