package controlplane

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"net"
	"sync/atomic"
	"testing"
	"time"

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

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

func connectorProjection(generation uint64, issuedAt time.Time, rawToken string) *iop.PrincipalProjection {
	digest := sha256.Sum256([]byte(rawToken))
	return &iop.PrincipalProjection{
		Generation:        generation,
		IssuedAtUnixNano:  issuedAt.UnixNano(),
		ExpiresAtUnixNano: issuedAt.Add(30 * time.Second).UnixNano(),
		Tokens: []*iop.ProjectedPrincipalToken{{
			TokenDigestSha256: hex.EncodeToString(digest[:]),
			PrincipalRef:      "principal-1",
			PrincipalAlias:    "principal-one",
			TokenRef:          "token-1",
			TokenRevision:     1,
		}},
		Routes: []*iop.ProjectedPrincipalRoute{{
			RouteId: "route-1", RouteAlias: "model-1", PrincipalRef: "principal-1",
			CredentialSlotRef: "slot-1", ProfileId: "openai", UpstreamModel: "upstream-model",
			ResourceSelector: "default", RouteRevision: 1, CredentialRevision: 1,
		}},
	}
}

type fakeStatusProvider struct {
	nodes        []edgeservice.NodeSnapshot
	capabilities []*iop.EdgeCapabilitySummary
	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) 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 {
		// Node-1 has no explicit Connected, so bool zero value (false) should
		// be passed through from the DTO — never hardcoded true.
		if n.GetConnected() {
			t.Errorf("node-1: expected Connected==false from DTO (zero), got true")
		}
		if n.GetConfig() != nil {
			t.Errorf("node-1: expected Config==nil (non-leak), got %+v", n.GetConfig())
		}
		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)
	} else {
		// Node-2 also has no explicit Connected — must be false from DTO.
		if n.GetConnected() {
			t.Errorf("node-2: expected Connected==false from DTO (zero), got true")
		}
		if n.GetConfig() != nil {
			t.Errorf("node-2: expected Config==nil (non-leak), got %+v", n.GetConfig())
		}
	}
}

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())
	}
}

func TestConnectorReconnectsWithRenewedSameGenerationProjection(t *testing.T) {
	var helloCount atomic.Int32
	helloSeen := make(chan int32, 8)
	port, serverClients, stopServer := startFakeCPServerWithClientAccess(t,
		func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
			count := helloCount.Add(1)
			helloSeen <- count
			return &iop.EdgeHelloResponse{Accepted: true, Protocol: "iop/1", PrincipalProjection: connectorProjection(7, time.Now(), "stable-token")}
		},
	)
	defer stopServer()
	cache := authprojection.NewCache(authprojection.DefaultLimits(), time.Now)
	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{Enabled: true, WireAddr: wireAddr(port), ReconnectIntervalSec: 1},
		"0.1.0",
		noopLogger(),
		WithPrincipalProjection(cache),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatal(err)
	}
	defer c.Stop()

	select {
	case <-helloSeen:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for initial projected hello")
	}
	readyCtx, readyCancel := context.WithTimeout(context.Background(), 3*time.Second)
	if err := c.WaitConnected(readyCtx); err != nil {
		readyCancel()
		t.Fatal(err)
	}
	readyCancel()
	select {
	case first := <-serverClients:
		_ = first.Close()
	case <-time.After(2 * time.Second):
		t.Fatal("timeout getting initial server connection")
	}
	select {
	case count := <-helloSeen:
		if count < 2 {
			t.Fatalf("unexpected hello count %d", count)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for same-generation reconnect")
	}
	reconnectCtx, reconnectCancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer reconnectCancel()
	if err := c.WaitConnected(reconnectCtx); err != nil {
		t.Fatal(err)
	}
	if generation, ok := cache.Generation(); !ok || generation != 7 || cache.State() != authprojection.StateFresh {
		t.Fatalf("renewed projection state: generation=%d ok=%v state=%s", generation, ok, cache.State())
	}
}

func TestConnectorRejectsChangedSameGenerationProjectionOnReconnect(t *testing.T) {
	var helloCount atomic.Int32
	helloSeen := make(chan int32, 8)
	port, serverClients, stopServer := startFakeCPServerWithClientAccess(t,
		func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
			count := helloCount.Add(1)
			helloSeen <- count
			rawToken := "stable-token"
			if count > 1 {
				rawToken = "changed-token"
			}
			return &iop.EdgeHelloResponse{Accepted: true, Protocol: "iop/1", PrincipalProjection: connectorProjection(7, time.Now(), rawToken)}
		},
	)
	defer stopServer()
	cache := authprojection.NewCache(authprojection.DefaultLimits(), time.Now)
	c := NewConnector(
		config.EdgeInfo{ID: "edge-test"},
		config.EdgeControlPlaneConf{Enabled: true, WireAddr: wireAddr(port), ReconnectIntervalSec: 1},
		"0.1.0",
		noopLogger(),
		WithPrincipalProjection(cache),
	)
	if err := c.Start(context.Background()); err != nil {
		t.Fatal(err)
	}
	defer c.Stop()

	select {
	case <-helloSeen:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for initial projected hello")
	}
	readyCtx, readyCancel := context.WithTimeout(context.Background(), 3*time.Second)
	if err := c.WaitConnected(readyCtx); err != nil {
		readyCancel()
		t.Fatal(err)
	}
	readyCancel()
	select {
	case first := <-serverClients:
		_ = first.Close()
	case <-time.After(2 * time.Second):
		t.Fatal("timeout getting initial server connection")
	}
	select {
	case count := <-helloSeen:
		if count < 2 {
			t.Fatalf("unexpected hello count %d", count)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for changed same-generation reconnect")
	}
	select {
	case second := <-serverClients:
		deadline := time.Now().Add(2 * time.Second)
		for second.IsAlive() && time.Now().Before(deadline) {
			time.Sleep(10 * time.Millisecond)
		}
		if second.IsAlive() {
			t.Fatal("connector kept changed same-generation hello connection alive")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout getting reconnect server connection")
	}
	if c.CurrentState() == StateConnected {
		t.Fatal("connector became ready with changed same-generation projection")
	}
	stableDigest := sha256.Sum256([]byte("stable-token"))
	if _, state, ok := cache.LookupDigest(stableDigest); !ok || state != authprojection.StateFresh {
		t.Fatalf("rejected reconnect changed installed projection: state=%s ok=%v", state, ok)
	}
}

// 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
	}
}

// TestConnectorBuildStatusResponseConnectedMapping verifies that the connector
// passes NodeSnapshot.Connected through to the wire EdgeNodeSnapshot.connected
// without hardcoding true, that provider tuples are delivered exactly, and that
// the Config field stays nil for every node (config redaction non-leak) regardless
// of whether the snapshot originated from a live or configured-only record.
func TestConnectorBuildStatusResponseConnectedMapping(t *testing.T) {
	fakeProvider := fakeStatusProvider{
		nodes: []edgeservice.NodeSnapshot{
			// Live/connected node: Connected=true, has provider snapshots and config.
			{
				NodeID:    "node-connected",
				Alias:     "online-alias",
				Label:     "node0",
				Connected: true,
				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:11434"},
							},
						},
					},
				},
				ProviderSnapshots: []*iop.ProviderSnapshot{
					{
						Adapter:             "ollama",
						Status:              "available",
						Capacity:            4,
						InFlight:            2,
						Queued:              3,
						LongContextCapacity: 2,
						LongInFlight:        1,
						LongQueued:          1,
					},
				},
			},
			// Configured-only offline node: Connected=false, no provider snapshots.
			{
				NodeID:    "node-offline",
				Alias:     "offline-alias",
				Label:     "node1",
				Connected: false,
			},
		},
	}

	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-connected"})
	if resp.GetRequestId() != "req-connected" || resp.GetEdgeId() != "edge-1" {
		t.Fatalf("unexpected basic response: %+v", resp)
	}

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

	byID := map[string]*iop.EdgeNodeSnapshot{}
	for _, n := range resp.GetNodes() {
		byID[n.GetNodeId()] = n
	}

	// Connected node: bool must be true, not hardcoded.
	connected := byID["node-connected"]
	if connected == nil {
		t.Fatalf("missing node-connected in response")
	}
	if connected.GetNodeId() != "node-connected" {
		t.Errorf("node-connected nodeId: got %q", connected.GetNodeId())
	}
	if connected.GetAlias() != "online-alias" {
		t.Errorf("node-connected alias: got %q", connected.GetAlias())
	}
	if connected.GetLabel() != "node0" {
		t.Errorf("node-connected label: got %q", connected.GetLabel())
	}
	if !connected.GetConnected() {
		t.Error("node-connected: expected Connected==true from DTO, got false")
	}
	if connected.GetConfig() != nil {
		t.Errorf("node-connected: expected Config==nil (non-leak), but got %+v", connected.GetConfig())
	}
	// Exact provider tuple for the connected node.
	if len(connected.GetProviderSnapshots()) != 1 {
		t.Fatalf("node-connected: expected 1 provider snapshot, got %d", len(connected.GetProviderSnapshots()))
	}
	ps := connected.GetProviderSnapshots()[0]
	if ps.GetAdapter() != "ollama" || ps.GetStatus() != "available" || ps.GetCapacity() != 4 ||
		ps.GetInFlight() != 2 || ps.GetQueued() != 3 ||
		ps.GetLongContextCapacity() != 2 || ps.GetLongInFlight() != 1 || ps.GetLongQueued() != 1 {
		t.Errorf("node-connected: unexpected exact provider tuple: %+v", ps)
	}

	// Offline node: bool must be false, no config leak, no provider snapshots.
	offline := byID["node-offline"]
	if offline == nil {
		t.Fatalf("missing node-offline in response")
	}
	if offline.GetNodeId() != "node-offline" {
		t.Errorf("node-offline nodeId: got %q", offline.GetNodeId())
	}
	if offline.GetAlias() != "offline-alias" {
		t.Errorf("node-offline alias: got %q", offline.GetAlias())
	}
	if offline.GetLabel() != "node1" {
		t.Errorf("node-offline label: got %q", offline.GetLabel())
	}
	if offline.GetConnected() {
		t.Error("node-offline: expected Connected==false from DTO, got true")
	}
	if offline.GetConfig() != nil {
		t.Errorf("node-offline: expected Config==nil (non-leak), but got %+v", offline.GetConfig())
	}
	if len(offline.GetProviderSnapshots()) != 0 {
		t.Errorf("node-offline: expected 0 provider snapshots, got %d", len(offline.GetProviderSnapshots()))
	}
}

// TestConnectorStatusRequestConnectedMappingOverWire verifies the connected
// boolean, exact provider tuple, and config redaction survive the full wire
// round-trip to a Control Plane and back as an EdgeStatusResponse.
func TestConnectorStatusRequestConnectedMappingOverWire(t *testing.T) {
	fakeProvider := fakeStatusProvider{
		nodes: []edgeservice.NodeSnapshot{
			{
				NodeID:    "wire-node-on",
				Alias:     "on-alias",
				Label:     "w1",
				Connected: true,
				ProviderSnapshots: []*iop.ProviderSnapshot{
					{
						Adapter:             "cli",
						Status:              "available",
						Capacity:            8,
						InFlight:            3,
						Queued:              5,
						LongContextCapacity: 3,
						LongInFlight:        1,
						LongQueued:          2,
					},
				},
			},
			{
				NodeID:    "wire-node-off",
				Alias:     "off-alias",
				Label:     "w2",
				Connected: false,
			},
		},
	}

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

	c := NewConnector(
		config.EdgeInfo{ID: "edge-wire", Name: "Wire Edge"},
		config.EdgeControlPlaneConf{
			Enabled:              true,
			WireAddr:             wireAddr(port),
			ReconnectIntervalSec: 60,
		},
		"0.5.0",
		noopLogger(),
		WithStatusProvider(fakeProvider),
	)
	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: "wire-mapping-test"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped status over wire: %v", err)
	}
	if resp.GetRequestId() != "wire-mapping-test" {
		t.Errorf("request_id: got %q want %q", resp.GetRequestId(), "wire-mapping-test")
	}
	if resp.GetEdgeId() != "edge-wire" {
		t.Errorf("edge_id: got %q want %q", resp.GetEdgeId(), "edge-wire")
	}
	if resp.GetError() != "" {
		t.Errorf("unexpected wire 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
	}

	// Connected node wire assertions.
	on := byID["wire-node-on"]
	if on == nil {
		t.Fatalf("missing wire-node-on in response")
	}
	if !on.GetConnected() {
		t.Error("wire: connected node Expected Connected==true on wire")
	}
	if on.GetConfig() != nil {
		t.Errorf("wire: connected node Config leak: %+v", on.GetConfig())
	}
	if len(on.GetProviderSnapshots()) != 1 {
		t.Fatalf("wire: connected node expected 1 provider snapshot, got %d", len(on.GetProviderSnapshots()))
	}
	cp := on.GetProviderSnapshots()[0]
	if cp.GetAdapter() != "cli" || cp.GetCapacity() != 8 || cp.GetInFlight() != 3 || cp.GetQueued() != 5 ||
		cp.GetLongContextCapacity() != 3 || cp.GetLongInFlight() != 1 || cp.GetLongQueued() != 2 {
		t.Errorf("wire: connected node unexpected provider tuple: %+v", cp)
	}

	// Offline node wire assertions.
	off := byID["wire-node-off"]
	if off == nil {
		t.Fatalf("missing wire-node-off in response")
	}
	if off.GetConnected() {
		t.Error("wire: offline node Expected Connected==false on wire")
	}
	if off.GetConfig() != nil {
		t.Errorf("wire: offline node Config leak: %+v", off.GetConfig())
	}
	if len(off.GetProviderSnapshots()) != 0 {
		t.Errorf("wire: offline node expected 0 provider snapshots, got %d", len(off.GetProviderSnapshots()))
	}
}

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)
		}
	}
}
