package upstream

import (
	"context"
	"errors"
	"testing"
	"time"
)

func TestHandleConnectionLossNotifiesPendingSubmitters(t *testing.T) {
	var reported ConnectionState
	conn := &connection{
		pending: make(map[uint32]chan submitPartResponse),
		pool: &connectionPool{
			channelID:    "channel-1",
			connectionID: "channel-1:primary",
			config:       normalizeUpstreamConfig(queueUpstreamConfigForTest()),
			reporter: func(_ context.Context, state ConnectionState) error {
				reported = state
				return nil
			},
			reconnectSignal: make(chan struct{}, 1),
			stopCh:          make(chan struct{}),
		},
	}
	conn.pool.conns = []*connection{conn}
	waiter := make(chan submitPartResponse, 1)
	conn.pending[7] = waiter

	loss := errors.New("socket closed")
	conn.handleConnectionLoss(loss)

	select {
	case result := <-waiter:
		if !errors.Is(result.err, loss) {
			t.Fatalf("pending waiter err = %v, want %v", result.err, loss)
		}
	default:
		t.Fatal("expected pending waiter to be notified")
	}

	if !conn.closed {
		t.Fatal("expected connection to be marked closed")
	}
	if len(conn.pending) != 0 {
		t.Fatalf("expected pending map to be reset, got %d entries", len(conn.pending))
	}
	if reported.Status != "disconnected" || reported.CurrentConnections != 0 {
		t.Fatalf("unexpected reported state: %+v", reported)
	}
	if conn.pool.reconnectCount != 1 || conn.pool.nextReconnectAt.IsZero() {
		t.Fatalf("expected connection loss to schedule reconnect, got count=%d next=%v", conn.pool.reconnectCount, conn.pool.nextReconnectAt)
	}
	conn.pool.close()
}

func TestHeartbeatTimeoutClosesConnectionAndSchedulesReconnect(t *testing.T) {
	pool := &connectionPool{
		channelID:       "channel-1",
		connectionID:    "channel-1:primary",
		config:          normalizeUpstreamConfig(queueUpstreamConfigForTest()),
		reconnectSignal: make(chan struct{}, 1),
		stopCh:          make(chan struct{}),
	}
	conn := &connection{
		pool:             pool,
		pending:          make(map[uint32]chan submitPartResponse),
		heartbeatPending: map[uint32]time.Time{1: time.Now(), 2: time.Now(), 3: time.Now()},
	}
	pool.conns = []*connection{conn}

	if conn.sendHeartbeat() {
		t.Fatal("expected heartbeat timeout to stop heartbeat loop")
	}
	if !conn.closed {
		t.Fatal("expected heartbeat timeout to close the connection")
	}
	if pool.reconnectCount != 1 || pool.lastErrorCategory != "heartbeat_timeout" {
		t.Fatalf("unexpected reconnect state: count=%d category=%s", pool.reconnectCount, pool.lastErrorCategory)
	}
	pool.close()
}

func TestReconnectDelayUsesCappedBackoffAndSlowAuthenticationRetry(t *testing.T) {
	if got := reconnectDelay(1, "network"); got < 4*time.Second || got > 6*time.Second {
		t.Fatalf("first reconnect delay = %v", got)
	}
	if got := reconnectDelay(6, "network"); got < 4*time.Minute || got > 6*time.Minute {
		t.Fatalf("capped reconnect delay = %v", got)
	}
	if got := reconnectDelay(1, "authentication"); got != 5*time.Minute {
		t.Fatalf("authentication reconnect delay = %v", got)
	}
}

func TestHeartbeatResponseClearsOnlyMatchingRequest(t *testing.T) {
	conn := &connection{heartbeatPending: map[uint32]time.Time{7: time.Now(), 8: time.Now()}}
	conn.handleHeartbeatResponse(7)
	if _, exists := conn.heartbeatPending[7]; exists {
		t.Fatal("expected matching heartbeat request to be cleared")
	}
	if _, exists := conn.heartbeatPending[8]; !exists {
		t.Fatal("expected unrelated heartbeat request to remain pending")
	}
}

func TestTemporaryReadTimeoutDetection(t *testing.T) {
	if !isTemporaryReadTimeout(fakeNetError{timeout: true}) {
		t.Fatal("expected timeout error to be treated as temporary")
	}
	if isTemporaryReadTimeout(errors.New("eof")) {
		t.Fatal("did not expect non-timeout error to be treated as temporary")
	}
}

type fakeNetError struct {
	timeout bool
}

func (f fakeNetError) Error() string   { return "network error" }
func (f fakeNetError) Timeout() bool   { return f.timeout }
func (f fakeNetError) Temporary() bool { return f.timeout }