package server

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/coder/websocket"

	"github.com/uopi/uopi/internal/broker"
	"github.com/uopi/uopi/internal/datasource/stub"
)

// newStressServer creates an httptest.Server backed by a stub with n signals.
// The server and broker contexts are tied to t.Context() and cleaned up automatically.
func newStressServer(t *testing.T, nSignals int) (*httptest.Server, *broker.Broker) {
	t.Helper()
	ctx := t.Context()
	log := slog.New(slog.NewTextHandler(io.Discard, nil))

	var ds *stub.Stub
	if nSignals > 0 {
		ds = stub.NewN(nSignals)
	} else {
		ds = stub.New()
	}
	if err := ds.Connect(ctx); err != nil {
		t.Fatal(err)
	}

	brk := broker.New(ctx, log)
	brk.Register(ds)

	mux := http.NewServeMux()
	mux.Handle("/ws", &wsHandler{broker: brk, log: log})

	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)
	return srv, brk
}

// wsURL converts an http:// test server URL to ws://.
func wsURL(srv *httptest.Server) string {
	return strings.ReplaceAll(srv.URL, "http://", "ws://") + "/ws"
}

// buildSubMsg encodes a subscribe message for n stub PVs named "pv_0"…"pv_{n-1}".
func buildSubMsg(n int) []byte {
	sigs := make([]sigRef, n)
	for i := range n {
		sigs[i] = sigRef{DS: "stub", Name: fmt.Sprintf("pv_%d", i)}
	}
	data, _ := json.Marshal(inMsg{Type: "subscribe", Signals: sigs})
	return data
}

// connectAndCount dials the WebSocket, sends subMsg, counts incoming messages until
// ctx is cancelled, adds the count to *total, and signals wg.Done.
func connectAndCount(ctx context.Context, t *testing.T, url string, subMsg []byte, total *atomic.Int64, wg *sync.WaitGroup) {
	t.Helper()
	defer wg.Done()

	conn, _, err := websocket.Dial(ctx, url, nil)
	if err != nil {
		t.Errorf("dial: %v", err)
		return
	}
	defer conn.CloseNow()

	if err := conn.Write(ctx, websocket.MessageText, subMsg); err != nil {
		t.Errorf("subscribe write: %v", err)
		return
	}

	var n int64
	for {
		_, _, err := conn.Read(ctx)
		if err != nil {
			break
		}
		n++
	}
	total.Add(n)
}

// TestStress_100PVs_10Clients verifies the full stack (broker + WS handler) can
// deliver updates from 100 signals to 10 concurrent clients without loss or deadlock.
func TestStress_100PVs_10Clients(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping stress test in short mode")
	}

	const (
		nSignals = 100
		nClients = 10
		duration = 3 * time.Second
	)

	srv, brk := newStressServer(t, nSignals)
	url := wsURL(srv)
	subMsg := buildSubMsg(nSignals)

	ctx, cancel := context.WithTimeout(t.Context(), duration+5*time.Second)
	defer cancel()

	var total atomic.Int64
	var wg sync.WaitGroup

	for range nClients {
		wg.Add(1)
		// Each client runs for exactly `duration`, then its context expires and
		// conn.Read returns an error, causing the goroutine to exit cleanly.
		clientCtx, clientCancel := context.WithTimeout(ctx, duration)
		go func() {
			defer clientCancel()
			connectAndCount(clientCtx, t, url, subMsg, &total, &wg)
		}()
	}

	wg.Wait()

	got := total.Load()
	// At 10 Hz, each PV fires ≥ duration*10 times; each client receives all of them.
	// min = nClients × nSignals × duration_s × 10 × 0.5 (conservative — meta msgs too)
	minExpected := int64(nClients) * int64(nSignals) * int64(duration.Seconds()) * 5
	t.Logf("100PVs/10Clients: received %d updates (min expected %d)", got, minExpected)
	if got < minExpected {
		t.Errorf("too few updates: got %d, want >= %d", got, minExpected)
	}

	// Broker must release all upstream subscriptions after clients disconnect.
	time.Sleep(300 * time.Millisecond)
	if n := brk.ActiveSubscriptions(); n != 0 {
		t.Errorf("subscription leak: %d remain after all clients disconnected", n)
	}
}

// TestStress_500PVs_20Clients is the full-scale scenario: hundreds of signals,
// tens of clients. Verifies no crash, no deadlock, and clean teardown.
func TestStress_500PVs_20Clients(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping stress test in short mode")
	}

	const (
		nSignals = 500
		nClients = 20
		duration = 3 * time.Second
	)

	srv, brk := newStressServer(t, nSignals)
	url := wsURL(srv)
	subMsg := buildSubMsg(nSignals)

	ctx, cancel := context.WithTimeout(t.Context(), duration+10*time.Second)
	defer cancel()

	var total atomic.Int64
	var wg sync.WaitGroup

	for range nClients {
		wg.Add(1)
		clientCtx, clientCancel := context.WithTimeout(ctx, duration)
		go func() {
			defer clientCancel()
			connectAndCount(clientCtx, t, url, subMsg, &total, &wg)
		}()
	}

	wg.Wait()

	got := total.Load()
	// Conservative: expect at least 20% of the theoretical maximum
	minExpected := int64(nClients) * int64(nSignals) * int64(duration.Seconds()) * 2
	t.Logf("500PVs/20Clients: received %d updates (min expected %d)", got, minExpected)
	if got < minExpected {
		t.Errorf("too few updates: got %d, want >= %d", got, minExpected)
	}

	time.Sleep(500 * time.Millisecond)
	if n := brk.ActiveSubscriptions(); n != 0 {
		t.Errorf("subscription leak: %d remain after all clients disconnected", n)
	}
}

// TestStress_ClientChurn repeatedly connects and disconnects clients while a
// stable set of signals is being delivered to ensure subscribe/unsubscribe under
// load doesn't cause races, panics, or subscription leaks.
func TestStress_ClientChurn(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping stress test in short mode")
	}

	const (
		nSignals       = 50
		nParallelConns = 30
		duration       = 3 * time.Second
	)

	srv, brk := newStressServer(t, nSignals)
	u := wsURL(srv)
	subMsg := buildSubMsg(nSignals)

	deadline := time.Now().Add(duration)
	var wg sync.WaitGroup
	var total atomic.Int64

	for range nParallelConns {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for time.Now().Before(deadline) {
				// Each iteration: one short-lived connection (~500 ms).
				func() {
					connCtx, connCancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
					defer connCancel()

					conn, _, err := websocket.Dial(connCtx, u, nil)
					if err != nil {
						return // dial failed (server busy) — try next iteration
					}
					defer conn.CloseNow()

					if err := conn.Write(connCtx, websocket.MessageText, subMsg); err != nil {
						return
					}
					var n int64
					for {
						_, _, err := conn.Read(connCtx)
						if err != nil {
							break
						}
						n++
					}
					total.Add(n)
				}()
			}
		}()
	}

	wg.Wait()

	t.Logf("ClientChurn: %d total messages during %s of churn with %d concurrent connections",
		total.Load(), duration, nParallelConns)

	time.Sleep(500 * time.Millisecond)
	if n := brk.ActiveSubscriptions(); n != 0 {
		t.Errorf("subscription leak after churn: %d remain", n)
	}
}

// TestStress_SlowConsumer verifies that a slow-reading client doesn't block
// fast clients or the broker fan-out (drop-on-full behaviour).
func TestStress_SlowConsumer(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping stress test in short mode")
	}

	const (
		nSignals = 100
		duration = 3 * time.Second
	)

	srv, brk := newStressServer(t, nSignals)
	url := wsURL(srv)
	subMsg := buildSubMsg(nSignals)

	ctx, cancel := context.WithTimeout(t.Context(), duration+5*time.Second)
	defer cancel()

	var fastTotal, slowTotal atomic.Int64
	var wg sync.WaitGroup

	// Fast client: reads as quickly as possible.
	wg.Add(1)
	fastCtx, fastCancel := context.WithTimeout(ctx, duration)
	go func() {
		defer fastCancel()
		connectAndCount(fastCtx, t, url, subMsg, &fastTotal, &wg)
	}()

	// Slow client: sleeps 100 ms between each read.
	wg.Add(1)
	slowCtx, slowCancel := context.WithTimeout(ctx, duration)
	go func() {
		defer wg.Done()
		defer slowCancel()

		conn, _, err := websocket.Dial(slowCtx, url, nil)
		if err != nil {
			t.Errorf("slow client dial: %v", err)
			return
		}
		defer conn.CloseNow()

		if err := conn.Write(slowCtx, websocket.MessageText, subMsg); err != nil {
			t.Errorf("slow client subscribe: %v", err)
			return
		}

		var n int64
		for {
			time.Sleep(100 * time.Millisecond)
			_, _, err := conn.Read(slowCtx)
			if err != nil {
				break
			}
			n++
		}
		slowTotal.Add(n)
	}()

	wg.Wait()

	fast := fastTotal.Load()
	slow := slowTotal.Load()
	t.Logf("SlowConsumer: fast=%d updates, slow=%d updates in %s", fast, slow, duration)

	// Fast client must receive substantially more than slow client.
	if fast < slow*3 {
		t.Errorf("fast client should receive many more updates than slow client (fast=%d slow=%d)", fast, slow)
	}

	// Fast client must have received a meaningful number of updates.
	minFast := int64(nSignals) * int64(duration.Seconds()) * 5
	if fast < minFast {
		t.Errorf("fast client received too few updates: got %d, want >= %d", fast, minFast)
	}

	time.Sleep(300 * time.Millisecond)
	if n := brk.ActiveSubscriptions(); n != 0 {
		t.Errorf("subscription leak: %d remain", n)
	}
}