package main

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

	proto_socket "git.toki-labs.com/toki/proto-socket/go"
	"iop/apps/control-plane/internal/wire"
	iop "iop/proto/gen/iop"
)

func TestFleetStatusHTTPHandlerCombinesConnectionAndCapabilities(t *testing.T) {
	registry := wire.NewEdgeRegistry()
	at := time.Unix(1780142200, 0).UTC()
	registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: "edge-online", EdgeName: "Online", Version: "1.0.0"}, at)
	registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: "edge-other", EdgeName: "Other", Version: "1.1.0"}, at)
	registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: "edge-degraded", EdgeName: "Degraded"}, at)
	offlineToken := registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: "edge-offline", EdgeName: "Offline"}, at)
	registry.MarkDisconnected("edge-offline", offlineToken, at.Add(time.Second), proto_socket.DisconnectInfo{Reason: "closed"})

	requestStatus := func(edgeID string, timeout time.Duration) (*iop.EdgeStatusResponse, error) {
		switch edgeID {
		case "edge-online":
			return &iop.EdgeStatusResponse{
				EdgeId: edgeID,
				Nodes:  []*iop.EdgeNodeSnapshot{{NodeId: "node-1", Alias: "alpha", Connected: true}},
				Capabilities: []*iop.EdgeCapabilitySummary{
					{Kind: "run-dispatch", Available: true, Status: "ready"},
				},
			}, nil
		case "edge-other":
			return &iop.EdgeStatusResponse{
				EdgeId: edgeID,
				Nodes:  []*iop.EdgeNodeSnapshot{{NodeId: "node-2", Alias: "beta", Connected: true}},
				Capabilities: []*iop.EdgeCapabilitySummary{
					{Kind: "run-dispatch", Available: true, Status: "ready"},
				},
			}, nil
		case "edge-degraded":
			return nil, fmt.Errorf("status timeout")
		default:
			return nil, fmt.Errorf("unexpected status request for %q", edgeID)
		}
	}

	mux := http.NewServeMux()
	registerFleetHandlers(mux, registry, requestStatus, nil)

	resp := httptest.NewRecorder()
	mux.ServeHTTP(resp, httptest.NewRequest(http.MethodGet, "/fleet/status", nil))
	if resp.Code != http.StatusOK {
		t.Fatalf("GET /fleet/status status=%d body=%s", resp.Code, resp.Body.String())
	}
	var got fleetStatusResponse
	if err := json.Unmarshal(resp.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode fleet status: %v", err)
	}
	if len(got.Edges) != 4 {
		t.Fatalf("expected 4 fleet edges, got %d", len(got.Edges))
	}
	health := map[string]fleetEdgeView{}
	for _, e := range got.Edges {
		health[e.EdgeID] = e
	}

	// Verify edge-online
	eOnline := health["edge-online"]
	if eOnline.Health != "online" || eOnline.NodeCount != 1 || len(eOnline.Capabilities) != 1 {
		t.Fatalf("unexpected online edge view: %+v", eOnline)
	}
	// Verify edge-other
	eOther := health["edge-other"]
	if eOther.Health != "online" || eOther.NodeCount != 1 || len(eOther.Capabilities) != 1 {
		t.Fatalf("unexpected edge-other view: %+v", eOther)
	}
	// Verify edge-degraded
	if e := health["edge-degraded"]; e.Health != "degraded" || e.Error == "" {
		t.Fatalf("unexpected degraded edge view: %+v", e)
	}

	// Verify edge-offline
	if e := health["edge-offline"]; e.Health != "offline" || e.Connected {
		t.Fatalf("unexpected offline edge view: %+v", e)
	}

	// Fleet status must not leak Node-direct internals (addresses, tokens) or
	// raw node snapshots; it surfaces a connection + capability summary only.
	bodyStr := resp.Body.String()
	for _, forbidden := range []string{"node_id", "token", "address", "\"nodes\""} {
		if strings.Contains(bodyStr, forbidden) {
			t.Fatalf("fleet status leaked %q: %s", forbidden, bodyStr)
		}
	}
}

func TestFleetCommandsHTTPHandlerFansOutToConnectedEdgesOnly(t *testing.T) {
	registry := wire.NewEdgeRegistry()
	at := time.Unix(1780142200, 0).UTC()
	registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: "edge-a"}, at)
	registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: "edge-b"}, at)
	registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: "edge-off"}, at)
	registry.MarkDisconnected("edge-off", 3, at.Add(time.Second), proto_socket.DisconnectInfo{Reason: "closed"})

	var dispatchMu sync.Mutex
	var dispatched []string
	sendCommand := func(edgeID, operation, targetSelector string, parameters map[string]string, timeout time.Duration) (*iop.EdgeCommandResponse, error) {
		dispatchMu.Lock()
		dispatched = append(dispatched, edgeID)
		dispatchMu.Unlock()
		return &iop.EdgeCommandResponse{CommandId: "cmd-" + edgeID, EdgeId: edgeID, Status: "ok"}, nil
	}

	mux := http.NewServeMux()
	registerFleetHandlers(mux, registry, nil, sendCommand)

	resp := httptest.NewRecorder()
	mux.ServeHTTP(resp, httptest.NewRequest(http.MethodPost, "/fleet/commands", strings.NewReader(`{"operation":"drain"}`)))
	if resp.Code != http.StatusAccepted {
		t.Fatalf("POST /fleet/commands status=%d body=%s", resp.Code, resp.Body.String())
	}
	var got fleetCommandResponse
	if err := json.Unmarshal(resp.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode fleet command response: %v", err)
	}
	if got.Operation != "drain" || len(got.Results) != 2 {
		t.Fatalf("expected fan-out to 2 connected edges, got %+v", got)
	}
	if len(dispatched) != 2 {
		t.Fatalf("expected dispatch to 2 connected edges, got %+v", dispatched)
	}
	for _, edge := range dispatched {
		if edge == "edge-off" {
			t.Fatalf("command was dispatched to a disconnected edge: %+v", dispatched)
		}
	}
}

func TestHTTPMuxHealthAndReadiness(t *testing.T) {
	mux := newHTTPMux()

	t.Run("Health", func(t *testing.T) {
		resp := httptest.NewRecorder()
		mux.ServeHTTP(resp, httptest.NewRequest(http.MethodGet, "/healthz", nil))
		if resp.Code != http.StatusOK {
			t.Fatalf("GET /healthz status=%d, want %d", resp.Code, http.StatusOK)
		}
		if resp.Body.String() != "ok\n" {
			t.Fatalf("GET /healthz body=%q, want %q", resp.Body.String(), "ok\n")
		}
	})

	t.Run("Readiness", func(t *testing.T) {
		resp := httptest.NewRecorder()
		mux.ServeHTTP(resp, httptest.NewRequest(http.MethodGet, "/readyz", nil))
		if resp.Code != http.StatusOK {
			t.Fatalf("GET /readyz status=%d, want %d", resp.Code, http.StatusOK)
		}
		if resp.Body.String() != "ready\n" {
			t.Fatalf("GET /readyz body=%q, want %q", resp.Body.String(), "ready\n")
		}
	})
}

func TestHTTPMuxRegistersEdgeHandlers(t *testing.T) {
	registry := wire.NewEdgeRegistry()
	mux := newHTTPMux()
	registerEdgeRegistryHandlers(mux, registry, nil, nil)
	registerFleetHandlers(mux, registry, nil, nil)

	t.Run("ListEdges", func(t *testing.T) {
		resp := httptest.NewRecorder()
		mux.ServeHTTP(resp, httptest.NewRequest(http.MethodGet, "/edges", nil))
		if resp.Code != http.StatusOK {
			t.Fatalf("GET /edges status=%d", resp.Code)
		}
	})

	t.Run("FleetStatus", func(t *testing.T) {
		resp := httptest.NewRecorder()
		mux.ServeHTTP(resp, httptest.NewRequest(http.MethodGet, "/fleet/status", nil))
		if resp.Code != http.StatusOK {
			t.Fatalf("GET /fleet/status status=%d", resp.Code)
		}
	})
}

func TestFleetStatusHTTPHandlerUsesConfiguredTimeout(t *testing.T) {
	registry := wire.NewEdgeRegistry()
	at := time.Unix(1780142200, 0).UTC()
	registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: "edge-a"}, at)

	var gotTimeout time.Duration
	requestStatus := func(edgeID string, timeout time.Duration) (*iop.EdgeStatusResponse, error) {
		gotTimeout = timeout
		return &iop.EdgeStatusResponse{EdgeId: edgeID}, nil
	}

	mux := http.NewServeMux()
	registerFleetHandlersWithOptions(mux, registry, requestStatus, nil, fleetOptions{StatusTimeout: 1234 * time.Millisecond})

	resp := httptest.NewRecorder()
	mux.ServeHTTP(resp, httptest.NewRequest(http.MethodGet, "/fleet/status", nil))
	if resp.Code != http.StatusOK {
		t.Fatalf("GET /fleet/status status=%d body=%s", resp.Code, resp.Body.String())
	}
	if gotTimeout != 1234*time.Millisecond {
		t.Fatalf("requester timeout=%s, want %s", gotTimeout, 1234*time.Millisecond)
	}
}

func TestFleetStatusHTTPHandlerBoundsConcurrentStatusRequests(t *testing.T) {
	registry := wire.NewEdgeRegistry()
	at := time.Unix(1780142200, 0).UTC()
	for _, id := range []string{"edge-1", "edge-2", "edge-3", "edge-4", "edge-5"} {
		registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: id}, at)
	}

	var inFlight, maxInFlight int32
	requestStatus := func(edgeID string, timeout time.Duration) (*iop.EdgeStatusResponse, error) {
		cur := atomic.AddInt32(&inFlight, 1)
		for {
			old := atomic.LoadInt32(&maxInFlight)
			if cur <= old || atomic.CompareAndSwapInt32(&maxInFlight, old, cur) {
				break
			}
		}
		time.Sleep(20 * time.Millisecond)
		atomic.AddInt32(&inFlight, -1)
		return &iop.EdgeStatusResponse{EdgeId: edgeID}, nil
	}

	mux := http.NewServeMux()
	registerFleetHandlersWithOptions(mux, registry, requestStatus, nil, fleetOptions{MaxConcurrentStatus: 2})

	resp := httptest.NewRecorder()
	mux.ServeHTTP(resp, httptest.NewRequest(http.MethodGet, "/fleet/status", nil))
	if resp.Code != http.StatusOK {
		t.Fatalf("GET /fleet/status status=%d body=%s", resp.Code, resp.Body.String())
	}
	if got := atomic.LoadInt32(&maxInFlight); got > 2 {
		t.Fatalf("max in-flight status requests=%d, want <= 2", got)
	}

	var got fleetStatusResponse
	if err := json.Unmarshal(resp.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode fleet status: %v", err)
	}
	wantOrder := []string{"edge-1", "edge-2", "edge-3", "edge-4", "edge-5"}
	if len(got.Edges) != len(wantOrder) {
		t.Fatalf("expected %d edges, got %d", len(wantOrder), len(got.Edges))
	}
	for i, want := range wantOrder {
		if got.Edges[i].EdgeID != want {
			t.Fatalf("edge[%d]=%q, want %q (response must keep snapshot order)", i, got.Edges[i].EdgeID, want)
		}
	}
}

func TestFleetStatusHTTPHandlerReusesFreshStatusCache(t *testing.T) {
	registry := wire.NewEdgeRegistry()
	at := time.Unix(1780142200, 0).UTC()
	registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: "edge-a"}, at)

	var calls int32
	requestStatus := func(edgeID string, timeout time.Duration) (*iop.EdgeStatusResponse, error) {
		atomic.AddInt32(&calls, 1)
		return &iop.EdgeStatusResponse{EdgeId: edgeID}, nil
	}

	mux := http.NewServeMux()
	registerFleetHandlersWithOptions(mux, registry, requestStatus, nil, fleetOptions{StatusCacheTTL: time.Minute})

	for i := 0; i < 2; i++ {
		resp := httptest.NewRecorder()
		mux.ServeHTTP(resp, httptest.NewRequest(http.MethodGet, "/fleet/status", nil))
		if resp.Code != http.StatusOK {
			t.Fatalf("GET /fleet/status #%d status=%d body=%s", i, resp.Code, resp.Body.String())
		}
	}
	if got := atomic.LoadInt32(&calls); got != 1 {
		t.Fatalf("requester call count=%d, want 1 (fresh cache must be reused)", got)
	}
}

func TestFleetCommandsHTTPHandlerBoundsConcurrentFanout(t *testing.T) {
	registry := wire.NewEdgeRegistry()
	at := time.Unix(1780142200, 0).UTC()
	for _, id := range []string{"edge-1", "edge-2", "edge-3", "edge-4", "edge-5"} {
		registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: id}, at)
	}
	registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: "edge-off"}, at)
	registry.MarkDisconnected("edge-off", 6, at.Add(time.Second), proto_socket.DisconnectInfo{Reason: "closed"})

	var inFlight, maxInFlight int32
	sendCommand := func(edgeID, operation, targetSelector string, parameters map[string]string, timeout time.Duration) (*iop.EdgeCommandResponse, error) {
		cur := atomic.AddInt32(&inFlight, 1)
		for {
			old := atomic.LoadInt32(&maxInFlight)
			if cur <= old || atomic.CompareAndSwapInt32(&maxInFlight, old, cur) {
				break
			}
		}
		time.Sleep(20 * time.Millisecond)
		atomic.AddInt32(&inFlight, -1)
		return &iop.EdgeCommandResponse{CommandId: "cmd-" + edgeID, EdgeId: edgeID, Status: "ok"}, nil
	}

	mux := http.NewServeMux()
	registerFleetHandlersWithOptions(mux, registry, nil, sendCommand, fleetOptions{MaxConcurrentCommands: 2})

	resp := httptest.NewRecorder()
	mux.ServeHTTP(resp, httptest.NewRequest(http.MethodPost, "/fleet/commands", strings.NewReader(`{"operation":"drain"}`)))
	if resp.Code != http.StatusAccepted {
		t.Fatalf("POST /fleet/commands status=%d body=%s", resp.Code, resp.Body.String())
	}
	if got := atomic.LoadInt32(&maxInFlight); got > 2 {
		t.Fatalf("max in-flight command dispatch=%d, want <= 2", got)
	}

	var got fleetCommandResponse
	if err := json.Unmarshal(resp.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode fleet command response: %v", err)
	}
	wantOrder := []string{"edge-1", "edge-2", "edge-3", "edge-4", "edge-5"}
	if len(got.Results) != len(wantOrder) {
		t.Fatalf("expected fan-out to %d connected edges, got %d: %+v", len(wantOrder), len(got.Results), got.Results)
	}
	for i, want := range wantOrder {
		if got.Results[i].EdgeID != want {
			t.Fatalf("result[%d]=%q, want %q (results must keep connected snapshot order)", i, got.Results[i].EdgeID, want)
		}
		if got.Results[i].EdgeID == "edge-off" {
			t.Fatalf("command was dispatched to a disconnected edge: %+v", got.Results)
		}
	}
}

func TestFleetCommandsHTTPHandlerUsesConfiguredTimeout(t *testing.T) {
	registry := wire.NewEdgeRegistry()
	at := time.Unix(1780142200, 0).UTC()
	registry.MarkConnected(&iop.EdgeHelloRequest{EdgeId: "edge-a"}, at)

	var gotTimeout time.Duration
	sendCommand := func(edgeID, operation, targetSelector string, parameters map[string]string, timeout time.Duration) (*iop.EdgeCommandResponse, error) {
		gotTimeout = timeout
		return &iop.EdgeCommandResponse{CommandId: "cmd-" + edgeID, EdgeId: edgeID, Status: "ok"}, nil
	}

	mux := http.NewServeMux()
	registerFleetHandlersWithOptions(mux, registry, nil, sendCommand, fleetOptions{CommandTimeout: 4321 * time.Millisecond})

	resp := httptest.NewRecorder()
	mux.ServeHTTP(resp, httptest.NewRequest(http.MethodPost, "/fleet/commands", strings.NewReader(`{"operation":"drain"}`)))
	if resp.Code != http.StatusAccepted {
		t.Fatalf("POST /fleet/commands status=%d body=%s", resp.Code, resp.Body.String())
	}
	if gotTimeout != 4321*time.Millisecond {
		t.Fatalf("sendCommand timeout=%s, want %s", gotTimeout, 4321*time.Millisecond)
	}
}
