package main

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

	"github.com/gorilla/websocket"
)

func TestMaintainConnectionsRestoresClosedClient(t *testing.T) {
	var connections atomic.Int64
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
		if err != nil {
			return
		}
		connections.Add(1)
		defer conn.Close()
		for {
			var packet map[string]json.RawMessage
			if err = conn.ReadJSON(&packet); err != nil {
				return
			}
			for _, raw := range packet {
				var body struct {
					ID string `json:"id"`
				}
				if json.Unmarshal(raw, &body) != nil || body.ID == "" {
					continue
				}
				response := map[string]any{"ctrl": map[string]any{
					"id": body.ID, "code": 200, "text": "ok",
					"params": map[string]any{"user": "usrReconnectTest"},
				}}
				if err = conn.WriteJSON(response); err != nil {
					return
				}
			}
		}
	}))
	defer server.Close()

	account := AccountRecord{
		Machine: "a", Index: 1, Username: "lta0001", Token: "token",
		P2PTopic: "p2pReconnectTest", Endpoint: "ws" + strings.TrimPrefix(server.URL, "http"),
	}
	cfg := &Config{Tenant: "allrich", CommandTimeout: "1s", WSPingInterval: "1h"}
	run := &agentRun{
		cfg: cfg, machine: "a", clients: make(map[int]*WSClient),
		recorder: newDeliveryRecorder("pending", "a", 1), ack: NewDurationHistogram(),
	}
	initial, err := run.connectAccountWithRetry(context.Background(), account, "grpReconnectTest", "api-key", "")
	if err != nil {
		t.Fatal(err)
	}
	run.setClient(account.Index, initial)
	initial.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	done := make(chan struct{})
	go func() {
		run.maintainConnections(ctx, []AccountRecord{account}, "grpReconnectTest", "api-key", "")
		close(done)
	}()
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if replacement := run.client(account.Index); replacement != initial && replacement.alive() {
			cancel()
			<-done
			run.close()
			if got := connections.Load(); got < 2 {
				t.Fatalf("websocket connections = %d, want at least 2", got)
			}
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	cancel()
	<-done
	run.close()
	t.Fatal("closed client was not restored before deadline")
}