package controlplane

import (
	"context"
	"net"
	"sync/atomic"
	"testing"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"

	edgeevents "iop/apps/edge/internal/events"
	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

type fakeStatusProvider struct {
	nodes        []edgeservice.NodeSnapshot
	capabilities []*iop.EdgeCapabilitySummary
	domainAgents []*iop.EdgeDomainAgentSummary
	executeErr   error
	executeResp  *iop.EdgeCommandResponse
	eventsToEmit []*iop.EdgeCommandEvent
}

func (f fakeStatusProvider) ListNodeSnapshots() []edgeservice.NodeSnapshot {
	return append([]edgeservice.NodeSnapshot(nil), f.nodes...)
}

func (f fakeStatusProvider) GetCapabilities() []*iop.EdgeCapabilitySummary {
	return f.capabilities
}

func (f fakeStatusProvider) GetDomainAgents() []*iop.EdgeDomainAgentSummary {
	return f.domainAgents
}

func (f fakeStatusProvider) ExecuteCommand(ctx context.Context, req *iop.EdgeCommandRequest, onEvent func(*iop.EdgeCommandEvent)) (*iop.EdgeCommandResponse, error) {
	for _, ev := range f.eventsToEmit {
		onEvent(ev)
	}
	if f.executeErr != nil {
		return nil, f.executeErr
	}
	if f.executeResp != nil {
		return f.executeResp, nil
	}
	return &iop.EdgeCommandResponse{
		Status:  "completed",
		Summary: "mock completed",
	}, nil
}

// fakeCPParserMap returns a parser for EdgeHelloRequest on the server side.
func fakeCPParserMap() toki.ParserMap {
	return toki.ParserMap{
		toki.TypeNameOf(&iop.EdgeHelloRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.EdgeHelloRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.EdgeNodeEvent{}): func(b []byte) (proto.Message, error) {
			m := &iop.EdgeNodeEvent{}
			return m, proto.Unmarshal(b, m)
		},
	}
}

// startFakeCPServer starts a fake Control Plane TCP server on a free port and
// calls onHello for each EdgeHelloRequest, returning the provided response.
// It returns the port and a stop function.
func startFakeCPServer(t *testing.T, onHello func(*iop.EdgeHelloRequest) *iop.EdgeHelloResponse) (int, func()) {
	t.Helper()
	port, err := freeTCPPort()
	if err != nil {
		t.Fatalf("startFakeCPServer freeTCPPort: %v", err)
	}
	srv := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient {
		return toki.NewTcpClient(conn, 0, 0, fakeCPParserMap())
	})
	srv.OnClientConnected = func(client *toki.TcpClient) {
		toki.AddRequestListenerTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
			&client.Communicator,
			func(req *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
				return onHello(req), nil
			},
		)
	}
	ctx, cancel := context.WithCancel(context.Background())
	if err := srv.Start(ctx); err != nil {
		cancel()
		t.Fatalf("startFakeCPServer Start: %v", err)
	}
	return port, func() {
		cancel()
		_ = srv.Stop()
	}
}

// startFakeCPServerWithClientAccess starts a fake server that also exposes
// each connected client via a channel, so the test can close clients directly.
func startFakeCPServerWithClientAccess(t *testing.T, onHello func(*iop.EdgeHelloRequest) *iop.EdgeHelloResponse) (int, <-chan *toki.TcpClient, func()) {
	t.Helper()
	port, err := freeTCPPort()
	if err != nil {
		t.Fatalf("startFakeCPServerWithClientAccess freeTCPPort: %v", err)
	}
	clients := make(chan *toki.TcpClient, 16)
	srv := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient {
		return toki.NewTcpClient(conn, 0, 0, fakeCPParserMap())
	})
	srv.OnClientConnected = func(client *toki.TcpClient) {
		clients <- client
		toki.AddRequestListenerTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
			&client.Communicator,
			func(req *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
				return onHello(req), nil
			},
		)
	}
	ctx, cancel := context.WithCancel(context.Background())
	if err := srv.Start(ctx); err != nil {
		cancel()
		t.Fatalf("startFakeCPServerWithClientAccess Start: %v", err)
	}
	return port, clients, func() {
		cancel()
		_ = srv.Stop()
	}
}

// wireAddr builds a host:port address string from a port int.
func wireAddr(port int) string {
	return net.JoinHostPort("127.0.0.1", tcpPortStr(port))
}

// statusCapableCPParserMap parses the Edge hello plus the status request the
// Edge marshals back so the fake Control Plane can both accept hello and read
// the status response.
func statusCapableCPParserMap() toki.ParserMap {
	return toki.ParserMap{
		toki.TypeNameOf(&iop.EdgeHelloRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.EdgeHelloRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.EdgeStatusResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.EdgeStatusResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}
}

// startFakeCPServerForStatus starts a fake Control Plane that accepts hello and
// exposes connected clients, so a test can send a status request to the Edge.
func startFakeCPServerForStatus(t *testing.T) (int, <-chan *toki.TcpClient, func()) {
	t.Helper()
	port, err := freeTCPPort()
	if err != nil {
		t.Fatalf("startFakeCPServerForStatus freeTCPPort: %v", err)
	}
	clients := make(chan *toki.TcpClient, 16)
	srv := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient {
		return toki.NewTcpClient(conn, 0, 0, statusCapableCPParserMap())
	})
	srv.OnClientConnected = func(client *toki.TcpClient) {
		clients <- client
		toki.AddRequestListenerTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
			&client.Communicator,
			func(_ *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
				return &iop.EdgeHelloResponse{Accepted: true, Protocol: "iop/1"}, nil
			},
		)
	}
	ctx, cancel := context.WithCancel(context.Background())
	if err := srv.Start(ctx); err != nil {
		cancel()
		t.Fatalf("startFakeCPServerForStatus Start: %v", err)
	}
	return port, clients, func() {
		cancel()
		_ = srv.Stop()
	}
}

func startFakeCPServerForNodeEvents(t *testing.T) (int, <-chan *iop.EdgeNodeEvent, func()) {
	t.Helper()
	port, err := freeTCPPort()
	if err != nil {
		t.Fatalf("startFakeCPServerForNodeEvents freeTCPPort: %v", err)
	}
	events := make(chan *iop.EdgeNodeEvent, 8)
	srv := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient {
		return toki.NewTcpClient(conn, 0, 0, fakeCPParserMap())
	})
	srv.OnClientConnected = func(client *toki.TcpClient) {
		toki.AddRequestListenerTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
			&client.Communicator,
			func(_ *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
				return &iop.EdgeHelloResponse{Accepted: true, Protocol: "iop/1"}, nil
			},
		)
		toki.AddListenerTyped[*iop.EdgeNodeEvent](&client.Communicator, func(event *iop.EdgeNodeEvent) {
			events <- event
		})
	}
	ctx, cancel := context.WithCancel(context.Background())
	if err := srv.Start(ctx); err != nil {
		cancel()
		t.Fatalf("startFakeCPServerForNodeEvents Start: %v", err)
	}
	return port, events, func() {
		cancel()
		_ = srv.Stop()
	}
}

// TestConnectorRespondsToStatusRequestFromProvider verifies that a status
// request from the Control Plane is answered from the configured snapshot
// boundary rather than a direct registry/transport read.
func TestConnectorRespondsToStatusRequestFromProvider(t *testing.T) {
	provider := fakeStatusProvider{nodes: []edgeservice.NodeSnapshot{
		{
			NodeID: "node-1",
			Alias:  "alpha",
			Label:  "node0",
			ProviderSnapshots: []*iop.ProviderSnapshot{
				{
					Adapter:             "cli",
					Status:              "available",
					Capacity:            4,
					InFlight:            2,
					Queued:              3,
					LongContextCapacity: 2,
					LongInFlight:        1,
					LongQueued:          1,
				},
			},
		},
		{NodeID: "node-2", Label: "node1"},
	}}

	port, serverClients, stop := startFakeCPServerForStatus(t)
	defer stop()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-dgx", Name: "DGX Edge"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 60,
		},
		"0.2.0",
		noopLogger(),
		WithStatusProvider(provider),
	)
	if !c.StatusProviderConfigured() {
		t.Fatal("expected status provider to be configured")
	}
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer c.Stop()

	var client *toki.TcpClient
	select {
	case client = <-serverClients:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for edge connection")
	}

	resp, err := toki.SendRequestTyped[*iop.EdgeStatusRequest, *iop.EdgeStatusResponse](
		&client.Communicator,
		&iop.EdgeStatusRequest{RequestId: "edge-status-test"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped status: %v", err)
	}
	if resp.GetRequestId() != "edge-status-test" {
		t.Errorf("request_id: got %q want %q", resp.GetRequestId(), "edge-status-test")
	}
	if resp.GetEdgeId() != "edge-dgx" {
		t.Errorf("edge_id: got %q want %q", resp.GetEdgeId(), "edge-dgx")
	}
	if resp.GetError() != "" {
		t.Errorf("unexpected error: %q", resp.GetError())
	}
	if len(resp.GetNodes()) != 2 {
		t.Fatalf("nodes: got %d want 2", len(resp.GetNodes()))
	}
	byID := map[string]*iop.EdgeNodeSnapshot{}
	for _, n := range resp.GetNodes() {
		byID[n.GetNodeId()] = n
	}
	if n := byID["node-1"]; n == nil || n.GetAlias() != "alpha" || n.GetLabel() != "node0" {
		t.Errorf("node-1 snapshot: %+v", n)
	} else {
		if len(n.GetProviderSnapshots()) != 1 {
			t.Fatalf("expected 1 provider snapshot, got %d", len(n.GetProviderSnapshots()))
		}
		snap := n.GetProviderSnapshots()[0]
		if snap.GetAdapter() != "cli" || snap.GetStatus() != "available" || snap.GetCapacity() != 4 || snap.GetInFlight() != 2 || snap.GetQueued() != 3 ||
			snap.GetLongContextCapacity() != 2 || snap.GetLongInFlight() != 1 || snap.GetLongQueued() != 1 {
			t.Errorf("unexpected snapshot contents: %+v", snap)
		}
	}
	if n := byID["node-2"]; n == nil || n.GetLabel() != "node1" {
		t.Errorf("node-2 snapshot: %+v", n)
	}
}

func TestConnectorRelaysNodeLifecycleEventsToControlPlane(t *testing.T) {
	bus := edgeevents.NewBus()
	port, events, stop := startFakeCPServerForNodeEvents(t)
	defer stop()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-dgx", Name: "DGX Edge"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 60,
		},
		"0.2.0",
		noopLogger(),
		WithNodeEventBus(bus),
	)
	if !c.NodeEventBusConfigured() {
		t.Fatal("expected node event bus to be configured")
	}
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer c.Stop()

	deadline := time.Now().Add(3 * time.Second)
	for c.CurrentState() != StateConnected {
		if time.Now().After(deadline) {
			t.Fatalf("connector did not reach connected state, got %v", c.CurrentState())
		}
		time.Sleep(20 * time.Millisecond)
	}

	original := &iop.EdgeNodeEvent{
		EventId:  "evt-node-1",
		Type:     "node.connected",
		Source:   "edge",
		NodeId:   "node-1",
		Alias:    "alpha",
		Reason:   "registered",
		Metadata: map[string]string{"rack": "r1"},
	}
	bus.PublishNode(original)

	select {
	case got := <-events:
		if got.GetEventId() != "evt-node-1" || got.GetType() != "node.connected" {
			t.Fatalf("unexpected relayed event: %+v", got)
		}
		if got.GetNodeId() != "node-1" || got.GetAlias() != "alpha" {
			t.Fatalf("unexpected relayed node identity: %+v", got)
		}
		if got.GetMetadata()["rack"] != "r1" {
			t.Fatalf("metadata rack: got %q", got.GetMetadata()["rack"])
		}
		if got.GetMetadata()["edge_id"] != "edge-dgx" {
			t.Fatalf("metadata edge_id: got %q want edge-dgx", got.GetMetadata()["edge_id"])
		}
		if got.GetTimestamp() == 0 {
			t.Fatal("expected relay to fill missing timestamp on copied event")
		}
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for relayed EdgeNodeEvent")
	}

	if _, ok := original.GetMetadata()["edge_id"]; ok {
		t.Fatal("relay mutated original event metadata")
	}
	if original.GetTimestamp() != 0 {
		t.Fatal("relay mutated original event timestamp")
	}
}

// TestConnectorStatusRequestWithoutProviderReturnsError verifies that a status
// request is answered with an explicit error when no provider is wired.
func TestConnectorStatusRequestWithoutProviderReturnsError(t *testing.T) {
	port, serverClients, stop := startFakeCPServerForStatus(t)
	defer stop()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-dgx", Name: "DGX Edge"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 60,
		},
		"0.2.0",
		noopLogger(),
	)
	if c.StatusProviderConfigured() {
		t.Fatal("expected no status provider configured")
	}
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer c.Stop()

	var client *toki.TcpClient
	select {
	case client = <-serverClients:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for edge connection")
	}

	resp, err := toki.SendRequestTyped[*iop.EdgeStatusRequest, *iop.EdgeStatusResponse](
		&client.Communicator,
		&iop.EdgeStatusRequest{RequestId: "edge-status-test"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped status: %v", err)
	}
	if resp.GetError() == "" {
		t.Error("expected an explicit error when no status provider is configured")
	}
	if len(resp.GetNodes()) != 0 {
		t.Errorf("expected no nodes without provider, got %d", len(resp.GetNodes()))
	}
}

// TestConnectorDisabledNoOp verifies that Start returns immediately when
// enabled=false, without touching the network.
func TestConnectorDisabledNoOp(t *testing.T) {
	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{Enabled: false, WireAddr: "127.0.0.1:9999"},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start disabled connector: %v", err)
	}
	if c.IsEnabled() {
		t.Fatal("expected IsEnabled()==false")
	}
	c.Stop() // must not panic
}

// TestConnectorEmptyAddrNoOp verifies that Start returns immediately when
// wire_addr is empty even if enabled=true.
func TestConnectorEmptyAddrNoOp(t *testing.T) {
	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{Enabled: true, WireAddr: ""},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start empty-addr connector: %v", err)
	}
	if c.IsEnabled() {
		t.Fatal("expected IsEnabled()==false for empty wire_addr")
	}
	c.Stop()
}

// TestConnectorSendsEdgeHello verifies that the connector dials and sends the
// correct EdgeHelloRequest fields.
func TestConnectorSendsEdgeHello(t *testing.T) {
	helloReceived := make(chan *iop.EdgeHelloRequest, 1)
	port, stop := startFakeCPServer(t, func(req *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
		select {
		case helloReceived <- req:
		default:
		}
		return &iop.EdgeHelloResponse{Accepted: true, Protocol: "iop/1"}
	})
	defer stop()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-dgx", Name: "DGX Edge"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 60,
		},
		"0.2.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer c.Stop()

	select {
	case req := <-helloReceived:
		if req.GetEdgeId() != "edge-dgx" {
			t.Errorf("edge_id: got %q, want %q", req.GetEdgeId(), "edge-dgx")
		}
		if req.GetEdgeName() != "DGX Edge" {
			t.Errorf("edge_name: got %q, want %q", req.GetEdgeName(), "DGX Edge")
		}
		if req.GetVersion() != "0.2.0" {
			t.Errorf("version: got %q, want %q", req.GetVersion(), "0.2.0")
		}
		if len(req.GetCapabilities()) == 0 {
			t.Error("expected non-empty capabilities")
		}
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for EdgeHelloRequest")
	}
}

// TestConnectorRejectsRejectedHello verifies that the connector treats a
// rejected hello as a connect failure and does not stay connected.
func TestConnectorRejectsRejectedHello(t *testing.T) {
	port, stop := startFakeCPServer(t, func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
		return &iop.EdgeHelloResponse{Accepted: false, Message: "not authorized"}
	})
	defer stop()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 60,
		},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer c.Stop()

	// Give connect time to attempt and fail.
	time.Sleep(400 * time.Millisecond)
	if c.CurrentState() == StateConnected {
		t.Error("expected connector not to be in Connected state after rejected hello")
	}
}

// TestConnectorStopClosesConnection verifies that Stop() terminates the loop
// and closes an active connection.
func TestConnectorStopClosesConnection(t *testing.T) {
	connected := make(chan struct{}, 1)
	port, stop := startFakeCPServer(t, func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
		select {
		case connected <- struct{}{}:
		default:
		}
		return &iop.EdgeHelloResponse{Accepted: true}
	})
	defer stop()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 60,
		},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}

	select {
	case <-connected:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for connection")
	}

	c.Stop()
	if c.CurrentState() != StateStopped {
		t.Errorf("expected StateStopped after Stop(), got %v", c.CurrentState())
	}
}

// TestConnectorReconnectsAfterDisconnect verifies that after the server closes
// the connection, the connector detects the disconnect and re-dials, sending a
// second EdgeHello. This directly exercises the reconnect loop.
func TestConnectorReconnectsAfterDisconnect(t *testing.T) {
	var helloCount atomic.Int32
	helloSeen := make(chan struct{}, 16)

	port, serverClients, stopServer := startFakeCPServerWithClientAccess(t,
		func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
			helloCount.Add(1)
			helloSeen <- struct{}{}
			return &iop.EdgeHelloResponse{Accepted: true}
		},
	)
	defer stopServer()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 1, // fast reconnect for test
		},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer c.Stop()

	// Wait for first hello / connection.
	select {
	case <-helloSeen:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for first EdgeHello")
	}

	// Server closes the first client connection to trigger a disconnect.
	select {
	case cl := <-serverClients:
		_ = cl.Close()
	case <-time.After(2 * time.Second):
		t.Fatal("timeout getting server-side client to close")
	}

	// Wait for a second hello; this proves the connector reconnected.
	select {
	case <-helloSeen:
		if helloCount.Load() < 2 {
			t.Errorf("expected at least 2 hellos, got %d", helloCount.Load())
		}
	case <-time.After(5 * time.Second):
		t.Fatalf("timeout waiting for reconnect (second EdgeHello); got %d hellos", helloCount.Load())
	}
}

// TestConnectorReconnectsAfterAcceptedThenClosed tests the race/regression
// where the Control Plane closes the connection immediately after sending an
// accepted EdgeHelloResponse. The connector must detect the disconnect even
// though it races with the hello response, and must reconnect.
func TestConnectorReconnectsAfterAcceptedThenClosed(t *testing.T) {
	var helloCount atomic.Int32
	helloSeen := make(chan struct{}, 16)

	port, serverClients, stopServer := startFakeCPServerWithClientAccess(t,
		func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
			helloCount.Add(1)
			helloSeen <- struct{}{}
			return &iop.EdgeHelloResponse{Accepted: true}
		},
	)
	defer stopServer()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 1,
		},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer c.Stop()

	// Wait for first hello, then immediately close the server-side client.
	// This simulates Control Plane closing the connection right after accepting.
	select {
	case <-helloSeen:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for first EdgeHello")
	}
	select {
	case cl := <-serverClients:
		_ = cl.Close()
	case <-time.After(2 * time.Second):
		t.Fatal("timeout getting server-side client to close immediately after hello")
	}

	// Connector must reconnect and send a second hello. If the disconnect
	// listener was not pre-registered, connect() would block forever here.
	select {
	case <-helloSeen:
		if helloCount.Load() < 2 {
			t.Errorf("expected at least 2 hellos, got %d", helloCount.Load())
		}
	case <-time.After(5 * time.Second):
		t.Fatalf("connector did not reconnect after accepted-then-closed; got %d hellos", helloCount.Load())
	}
}

// TestConnectorStopPreventsReconnect verifies that after Stop() the connector
// does not send additional hellos within a short window.
func TestConnectorStopPreventsReconnect(t *testing.T) {
	var helloCount atomic.Int32
	helloSeen := make(chan struct{}, 16)

	port, serverClients, stopServer := startFakeCPServerWithClientAccess(t,
		func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
			helloCount.Add(1)
			helloSeen <- struct{}{}
			return &iop.EdgeHelloResponse{Accepted: true}
		},
	)
	defer stopServer()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 1,
		},
		"0.1.0",
		noopLogger(),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start: %v", err)
	}

	// Wait for first hello.
	select {
	case <-helloSeen:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for first EdgeHello")
	}

	// Stop the connector before closing the server-side client.
	c.Stop()

	// Now close the server-side client.
	select {
	case cl := <-serverClients:
		_ = cl.Close()
	default:
		// client may already be gone after Stop
	}

	// Give a window longer than ReconnectIntervalSec; no new hello should arrive.
	countBefore := helloCount.Load()
	select {
	case <-helloSeen:
		t.Errorf("connector sent hello after Stop(); count went from %d to %d", countBefore, helloCount.Load())
	case <-time.After(2 * time.Second):
		// good; no extra hello
	}
}

func TestConnectorBuildStatusResponseNonLeak(t *testing.T) {
	fakeProvider := fakeStatusProvider{
		nodes: []edgeservice.NodeSnapshot{
			{
				NodeID: "node-1",
				Alias:  "alpha",
				Label:  "node0",
				Config: &iop.NodeConfigPayload{
					Runtime: &iop.NodeRuntimeConfig{
						Concurrency: 5,
					},
					Adapters: []*iop.AdapterConfig{
						{
							Type:    "ollama",
							Enabled: true,
							Config: &iop.AdapterConfig_Ollama{
								Ollama: &iop.OllamaAdapterConfig{
									BaseUrl: "http://sensitive-url:11434",
								},
							},
						},
					},
				},
			},
		},
	}

	connector := NewConnector(
		config.EdgeInfo{ID: "edge-1", Name: "Edge 1"},
		config.EdgeControlPlaneConf{},
		"1.0.0",
		noopLogger(),
		WithStatusProvider(fakeProvider),
	)

	resp := connector.buildStatusResponse(&iop.EdgeStatusRequest{RequestId: "req-1"})
	if resp.GetRequestId() != "req-1" || resp.GetEdgeId() != "edge-1" {
		t.Fatalf("unexpected basic response: %+v", resp)
	}

	if len(resp.GetNodes()) != 1 {
		t.Fatalf("nodes length: got %d want 1", len(resp.GetNodes()))
	}

	n := resp.GetNodes()[0]
	if n.GetNodeId() != "node-1" || n.GetAlias() != "alpha" || n.GetLabel() != "node0" {
		t.Errorf("unexpected node: %+v", n)
	}

	if n.GetConfig() != nil {
		t.Errorf("expected config to be nil (non-leak), but got: %+v", n.GetConfig())
	}
}

func commandCapableCPParserMap() toki.ParserMap {
	m := cpParserMap()
	m[toki.TypeNameOf(&iop.EdgeHelloRequest{})] = func(b []byte) (proto.Message, error) {
		req := &iop.EdgeHelloRequest{}
		return req, proto.Unmarshal(b, req)
	}
	return m
}

func TestConnectorEdgeCommandRelay(t *testing.T) {
	fakeProvider := fakeStatusProvider{
		executeResp: &iop.EdgeCommandResponse{
			Status:  "completed",
			Summary: "test OK",
		},
		eventsToEmit: []*iop.EdgeCommandEvent{
			{CommandId: "cmd-1", Phase: "started", Summary: "started test"},
			{CommandId: "cmd-1", Phase: "completed", Summary: "finished test"},
		},
	}

	port, err := freeTCPPort()
	if err != nil {
		t.Fatalf("freeTCPPort: %v", err)
	}

	eventsReceived := make(chan *iop.EdgeCommandEvent, 8)
	var cpClient *toki.TcpClient

	srv := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient {
		return toki.NewTcpClient(conn, 0, 0, commandCapableCPParserMap())
	})
	srv.OnClientConnected = func(client *toki.TcpClient) {
		cpClient = client
		toki.AddRequestListenerTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
			&client.Communicator,
			func(_ *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
				return &iop.EdgeHelloResponse{Accepted: true, Protocol: "iop/1"}, nil
			},
		)
		toki.AddListenerTyped[*iop.EdgeCommandEvent](&client.Communicator, func(ev *iop.EdgeCommandEvent) {
			eventsReceived <- ev
		})
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	if err := srv.Start(ctx); err != nil {
		t.Fatalf("Start cp server: %v", err)
	}
	defer func() { _ = srv.Stop() }()

	c := NewConnector(
		config.EdgeInfo{ID: "edge-test-1", Name: "Edge Test"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 60,
		},
		"1.0.0",
		noopLogger(),
		WithStatusProvider(fakeProvider),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatalf("Start connector: %v", err)
	}
	defer c.Stop()

	deadline := time.Now().Add(3 * time.Second)
	for c.CurrentState() != StateConnected {
		if time.Now().After(deadline) {
			t.Fatalf("connector failed to connect, got %v", c.CurrentState())
		}
		time.Sleep(10 * time.Millisecond)
	}

	if cpClient == nil {
		t.Fatal("control plane client not connected")
	}

	resp, err := toki.SendRequestTyped[*iop.EdgeCommandRequest, *iop.EdgeCommandResponse](
		&cpClient.Communicator,
		&iop.EdgeCommandRequest{
			RequestId: "req-command-1",
			CommandId: "cmd-1",
			Operation: "test.op",
		},
		3*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped edge command: %v", err)
	}

	if resp.GetStatus() != "completed" || resp.GetSummary() != "test OK" {
		t.Errorf("unexpected response: %+v", resp)
	}
	if resp.GetEdgeId() != "edge-test-1" {
		t.Errorf("expected edge_id edge-test-1, got %q", resp.GetEdgeId())
	}

	for i := 0; i < 2; i++ {
		select {
		case ev := <-eventsReceived:
			if ev.GetCommandId() != "cmd-1" {
				t.Errorf("unexpected event command_id: %q", ev.GetCommandId())
			}
			if ev.GetEdgeId() != "edge-test-1" {
				t.Errorf("unexpected event edge_id: %q", ev.GetEdgeId())
			}
		case <-time.After(2 * time.Second):
			t.Fatalf("timeout waiting for event %d to be relayed", i)
		}
	}
}