package observability

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

	"github.com/coder/websocket"
)

func TestWebSocketHandshakeIsNotConnectionLifetime(t *testing.T) {
	r := NewRegistry()
	released := make(chan struct{})
	done := make(chan struct{})
	h := r.Middleware(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		defer close(done)
		c, err := websocket.Accept(w, req, nil)
		if err != nil {
			return
		}
		defer c.CloseNow()
		<-released
	}))
	server := httptest.NewTLSServer(h)
	defer server.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	c, _, err := websocket.Dial(ctx, strings.Replace(server.URL, "https://", "wss://", 1)+"/v1/worker/connect", &websocket.DialOptions{HTTPClient: server.Client()})
	if err != nil {
		close(released)
		t.Fatal(err)
	}
	defer c.CloseNow()
	// Dial returns after the upgrade response, while the server handler stays open.
	deadline := time.Now().Add(time.Second)
	for !strings.Contains(r.Render(), `orbit_websocket_connections{route="worker_connect"} 1`) {
		if time.Now().After(deadline) {
			close(released)
			t.Fatal("upgrade not observed")
		}
		time.Sleep(time.Millisecond)
	}
	out := r.Render()
	close(released)
	<-done
	if !strings.Contains(out, `orbit_http_requests_total{method="GET",route="worker_connect",status_class="1xx"} 1`) || !strings.Contains(out, "orbit_http_requests_in_flight 0\n") {
		t.Fatal("upgrade counted as unfinished HTTP")
	}
	// done closes inside the handler; wait for middleware's defer to publish closure.
	deadline = time.Now().Add(time.Second)
	for !strings.Contains(r.Render(), `orbit_websocket_connections{route="worker_connect"} 0`) {
		if time.Now().After(deadline) {
			t.Fatal("connection gauge leaked")
		}
		time.Sleep(time.Millisecond)
	}
	if !strings.Contains(r.Render(), `orbit_http_requests_total{method="GET",route="worker_connect",status_class="1xx"} 1`) {
		t.Fatal("upgrade counted twice")
	}
}
func TestInformationalResponseAndStreamingPreserveHTTP(t *testing.T) {
	r := NewRegistry()
	server := httptest.NewServer(r.Middleware(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		w.WriteHeader(103)
		w.WriteHeader(201)
		_, _ = io.WriteString(w, "first")
		if err := http.NewResponseController(w).Flush(); err != nil {
			panic(err)
		}
		_, _ = io.WriteString(w, "second")
	})))
	defer server.Close()
	resp, err := server.Client().Get(server.URL + "/v1/tasks")
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil || resp.StatusCode != 201 || string(body) != "firstsecond" {
		t.Fatal("wrapper changed HTTP response", resp.StatusCode, err)
	}
	deadline := time.Now().Add(time.Second)
	for !strings.Contains(r.Render(), `route="tasks",status_class="2xx"} 1`) {
		if time.Now().After(deadline) {
			t.Fatal("103 blocked final status metric")
		}
		time.Sleep(time.Millisecond)
	}
}