package omni_test

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

	omni "git.yuntuo.hk/k2safe/go-omni"
	"git.yuntuo.hk/k2safe/go-omni/httpx"
	"git.yuntuo.hk/k2safe/go-omni/lifecycle"
	"git.yuntuo.hk/k2safe/go-omni/ws"
	wsomni "git.yuntuo.hk/k2safe/go-omni/ws/omni"
	"github.com/gorilla/websocket"
)

func TestMountLifetimeAndStatus(t *testing.T) {
	a := omni.New(omni.WithLogger(false))
	var observed atomic.Int64
	a.Use(func(next httpx.Handler) httpx.Handler {
		return func(c *httpx.Context) error { err := next(c); observed.Store(int64(c.Status())); return err }
	})
	s, err := ws.NewServer(ws.ServerConfig{Session: ws.SessionConfig{CloseTimeout: 30 * time.Millisecond}})
	if err != nil {
		t.Fatal(err)
	}
	if err = wsomni.Mount(a, "/ws", s, ws.Handlers{OnMessage: func(ctx context.Context, s *ws.Session, m ws.Message) error { return s.Send(ctx, m) }}); err != nil {
		t.Fatal(err)
	}
	h := httptest.NewServer(a)
	defer h.Close()
	defer a.Close()
	c, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(h.URL, "http")+"/ws", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	for i := 0; i < 10; i++ {
		if err = c.WriteMessage(websocket.TextMessage, []byte("hello")); err != nil {
			t.Fatal(err)
		}
		c.SetReadDeadline(time.Now().Add(time.Second))
		_, v, e := c.ReadMessage()
		if e != nil || string(v) != "hello" {
			t.Fatal(e)
		}
	}
	if observed.Load() != 101 {
		t.Fatal("upgrade logged as", observed.Load())
	}
	var cleanup atomic.Bool
	a.OnShutdown(func() error {
		if s.Stats().ActiveConnections != 0 {
			t.Error("resource closed before sessions")
		}
		cleanup.Store(true)
		return nil
	})
	if err = a.Close(); err != nil {
		t.Fatal(err)
	}
	if !cleanup.Load() {
		t.Fatal("missing cleanup")
	}
}

func TestAppRetainsResourcesForBlockedComponent(t *testing.T) {
	a := omni.New(omni.WithLogger(false))
	closed := false
	a.OnStop("blocked", func(context.Context) error { return nil })
	a.OnDrain("blocked", func(context.Context) error { return context.DeadlineExceeded })
	a.OnShutdown(func() error { closed = true; return nil })
	if err := a.Close(); !errors.Is(err, lifecycle.ErrResourcesRetained) {
		t.Fatal(err)
	}
	if closed {
		t.Fatal("unsafe cleanup")
	}
}

func TestMountRejectsNative(t *testing.T) {
	a := omni.New(omni.WithTransport(omni.TransportNative))
	defer a.Close()
	s, _ := ws.NewServer(ws.ServerConfig{})
	defer s.Shutdown(context.Background())
	if err := wsomni.Mount(a, "/ws", s, ws.Handlers{}); !errors.Is(err, ws.ErrUnsupportedTransport) {
		t.Fatal(err)
	}
}