package realtime

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/coder/websocket"
)

// dialOrigin 以指定 Origin 发起握手；origin 为空表示不带 Origin 头（非浏览器客户端）。
func dialOrigin(t *testing.T, srv *httptest.Server, origin string) error {
	t.Helper()

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	opts := &websocket.DialOptions{}
	if origin != "" {
		opts.HTTPHeader = http.Header{"Origin": []string{origin}}
	}
	conn, _, err := websocket.Dial(ctx, "ws"+strings.TrimPrefix(srv.URL, "http")+"/ws", opts)
	if err != nil {
		return err
	}
	_ = conn.CloseNow()
	return nil
}

func TestHandlerOrigin(t *testing.T) {
	t.Run("同源放行", func(t *testing.T) {
		e := newTestEnv(t)
		// httptest.Server 的 URL 即页面 Origin：scheme + host 与请求 Host 同源。
		if err := dialOrigin(t, e.srv, e.srv.URL); err != nil {
			t.Fatalf("同源 Origin 应放行: %v", err)
		}
	})

	t.Run("空 Origin 放行", func(t *testing.T) {
		e := newTestEnv(t)
		if err := dialOrigin(t, e.srv, ""); err != nil {
			t.Fatalf("不带 Origin 的客户端应放行: %v", err)
		}
	})

	t.Run("跨站拒绝", func(t *testing.T) {
		e := newTestEnv(t)
		if err := dialOrigin(t, e.srv, "http://evil.example"); err == nil {
			t.Fatal("跨站 Origin 应拒绝")
		}
	})

	t.Run("白名单放行", func(t *testing.T) {
		e := newTestEnv(t)
		e.hub.SetAllowedOrigins([]string{"http://evil.example"})
		if err := dialOrigin(t, e.srv, "http://evil.example"); err != nil {
			t.Fatalf("白名单内 Origin 应放行: %v", err)
		}
		if err := dialOrigin(t, e.srv, "http://other.example"); err == nil {
			t.Fatal("白名单外的跨站 Origin 应拒绝")
		}
	})
}