package service

import (
	"context"
	"fmt"
	"net"
	"strings"
	"testing"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	"go.uber.org/zap"
	"go.uber.org/zap/zaptest/observer"
	"google.golang.org/protobuf/proto"

	edgeevents "iop/apps/edge/internal/events"
	runtime "iop/packages/go/execution"
	iop "iop/proto/gen/iop"
)

func installProviderHealthTestObserver(t *testing.T, svc *Service) (*prometheus.Registry, *observer.ObservedLogs) {
	t.Helper()
	registry := prometheus.NewRegistry()
	core, logs := observer.New(zap.InfoLevel)
	svc.queue.mu.Lock()
	svc.queue.healthObserver = newProviderHealthObservability(registry, zap.New(core))
	svc.queue.mu.Unlock()
	return registry, logs
}

func metricValue(t *testing.T, registry *prometheus.Registry, name string, want map[string]string) float64 {
	t.Helper()
	families, err := registry.Gather()
	if err != nil {
		t.Fatalf("gather metrics: %v", err)
	}
	for _, family := range families {
		if family.GetName() != name {
			continue
		}
		for _, metric := range family.Metric {
			if metricHasLabels(metric, want) {
				return metric.GetCounter().GetValue()
			}
		}
	}
	return 0
}

func metricHasLabels(metric *dto.Metric, want map[string]string) bool {
	if len(metric.Label) != len(want) {
		return false
	}
	for _, label := range metric.Label {
		if want[label.GetName()] != label.GetValue() {
			return false
		}
	}
	return true
}

func assertPublicProviderSnapshot(t *testing.T, snapshots []NodeSnapshot, nodeID, providerID string, wantStatus, wantHealth string, wantCapacity int32) {
	t.Helper()
	for _, snap := range snapshots {
		if snap.NodeID != nodeID {
			continue
		}
		for _, ps := range snap.ProviderSnapshots {
			if ps.GetId() == providerID {
				if ps.GetStatus() != wantStatus || ps.GetHealth() != wantHealth || ps.GetCapacity() != wantCapacity {
					t.Fatalf("snapshot for %s/%s = (status=%q, health=%q, capacity=%d), want (%q, %q, %d)",
						nodeID, providerID, ps.GetStatus(), ps.GetHealth(), ps.GetCapacity(),
						wantStatus, wantHealth, wantCapacity)
				}
				return
			}
		}
	}
	t.Fatalf("snapshot for %s/%s not found", nodeID, providerID)
}

func TestProviderHealthObservability(t *testing.T) {
	for _, executionPath := range []string{"normalized", "tunnel"} {
		t.Run(executionPath, func(t *testing.T) {
			edgeConn, nodeConn := net.Pipe()
			t.Cleanup(func() {
				_ = edgeConn.Close()
				_ = nodeConn.Close()
			})
			parserMap := toki.ParserMap{
				toki.TypeNameOf(&iop.NodeCommandRequest{}): func(data []byte) (proto.Message, error) {
					message := &iop.NodeCommandRequest{}
					return message, proto.Unmarshal(data, message)
				},
				toki.TypeNameOf(&iop.NodeCommandResponse{}): func(data []byte) (proto.Message, error) {
					message := &iop.NodeCommandResponse{}
					return message, proto.Unmarshal(data, message)
				},
			}
			edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
			nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
			toki.AddRequestListenerTyped(&nodeClient.Communicator, func(request *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
				return &iop.NodeCommandResponse{
					RequestId: request.GetRequestId(), Type: request.GetType(),
					Adapter: request.GetAdapter(), Target: request.GetTarget(), SessionId: request.GetSessionId(),
					Result: map[string]string{
						"adapter_key": request.GetAdapter(), "target": request.GetTarget(),
						"provider_status": "available", "health_observation_seq": "4",
					},
				}, nil
			})

			svc, entry, _ := newProviderHealthOverlayService(t, edgeClient)
			registry, logs := installProviderHealthTestObserver(t, svc)

			// 1. Unavailable terminal (sequence 3)
			addBoundOverlayLease(t, svc.queue, "run-unhealthy", overlayProviderID, overlayAdapter, overlayTarget, entry.ConnectionGeneration)
			if executionPath == "normalized" {
				svc.HandleReceivedRunLifecycleEvent(entry.NodeID, entry.ConnectionGeneration, unavailableRunEvent("run-unhealthy", overlayAdapter, overlayTarget, 3))
			} else {
				svc.HandleReceivedProviderTunnelFrame(entry.NodeID, entry.ConnectionGeneration, &iop.ProviderTunnelFrame{
					RunId: "run-unhealthy", TunnelId: "tunnel-observability",
					Kind:    iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR,
					Failure: stallFailure("run-unhealthy", overlayAdapter, overlayTarget, "unavailable", "provider_unhealthy", 3),
				})
			}
			assertOverlayUnavailable(t, svc.queue, entry.ConnectionGeneration, true, 3)
			assertPublicProviderSnapshot(t, svc.ListNodeSnapshots(), entry.NodeID, overlayProviderID, "unavailable", "unavailable", 0)

			// 2. Stale terminal (sequence 3) delivered through selected executionPath handler
			addBoundOverlayLease(t, svc.queue, "run-stale", overlayProviderID, overlayAdapter, overlayTarget, entry.ConnectionGeneration)
			stale := stallFailure("run-stale", overlayAdapter, overlayTarget, "available", "request_stalled", 3)
			if executionPath == "normalized" {
				svc.HandleReceivedRunLifecycleEvent(entry.NodeID, entry.ConnectionGeneration, &iop.RunEvent{RunId: "run-stale", Type: "error", Failure: stale})
			} else {
				svc.HandleReceivedProviderTunnelFrame(entry.NodeID, entry.ConnectionGeneration, &iop.ProviderTunnelFrame{
					RunId: "run-stale", TunnelId: "tunnel-stale",
					Kind:    iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR,
					Failure: stale,
				})
			}
			assertOverlayUnavailable(t, svc.queue, entry.ConnectionGeneration, true, 3)
			assertPublicProviderSnapshot(t, svc.ListNodeSnapshots(), entry.NodeID, overlayProviderID, "unavailable", "unavailable", 0)

			// 3. Recovery via production Capabilities path (higher-sequence probe 4)
			if _, err := svc.Capabilities(context.Background(), NodeCommandRequestSpec{
				NodeRef: entry.NodeID, Adapter: overlayAdapter, Target: overlayTarget,
			}); err != nil {
				t.Fatalf("Capabilities recovery: %v", err)
			}
			assertOverlayUnavailable(t, svc.queue, entry.ConnectionGeneration, false, 4)
			assertPublicProviderSnapshot(t, svc.ListNodeSnapshots(), entry.NodeID, overlayProviderID, "available", "available", 1)

			if got := metricValue(t, registry, providerHealthEvidenceMetric, map[string]string{"source": "stall", "evidence_health": "unavailable", "decision": "applied"}); got != 1 {
				t.Fatalf("applied unhealthy metric = %v, want 1", got)
			}
			if got := metricValue(t, registry, providerHealthEvidenceMetric, map[string]string{"source": "stall", "evidence_health": "available", "decision": "rejected_stale"}); got != 1 {
				t.Fatalf("stale rejection metric = %v, want 1", got)
			}
			if got := metricValue(t, registry, providerHealthEvidenceMetric, map[string]string{"source": "probe", "evidence_health": "available", "decision": "applied"}); got != 1 {
				t.Fatalf("recovery metric = %v, want 1", got)
			}
			if got := metricValue(t, registry, providerHealthTransitionMetric, map[string]string{"from_health": "available", "to_health": "unavailable"}); got != 1 {
				t.Fatalf("unhealthy transition metric = %v, want 1", got)
			}
			if got := metricValue(t, registry, providerHealthTransitionMetric, map[string]string{"from_health": "unavailable", "to_health": "available"}); got != 1 {
				t.Fatalf("recovery transition metric = %v, want 1", got)
			}

			entries := logs.All()
			if len(entries) != 3 {
				t.Fatalf("health observation logs = %d, want 3", len(entries))
			}
			for _, entry := range entries {
				if entry.Message != providerHealthObservationLogKey {
					t.Fatalf("unexpected log message %q", entry.Message)
				}
				for _, field := range entry.Context {
					if strings.Contains(field.Key, "provider") || strings.Contains(field.Key, "node") || strings.Contains(field.Key, "run") || strings.Contains(field.Key, "session") || strings.Contains(field.Key, "adapter") || strings.Contains(field.Key, "target") {
						t.Fatalf("identity-bearing log field %q", field.Key)
					}
				}
			}
		})
	}

	t.Run("duplicate evidence is observed exactly once", func(t *testing.T) {
		svc, entry, _ := newProviderHealthOverlayService(t, nil)
		registry, _ := installProviderHealthTestObserver(t, svc)
		addBoundOverlayLease(t, svc.queue, "run-once", overlayProviderID, overlayAdapter, overlayTarget, entry.ConnectionGeneration)
		event := unavailableRunEvent("run-once", overlayAdapter, overlayTarget, 1)
		svc.HandleReceivedRunLifecycleEvent(entry.NodeID, entry.ConnectionGeneration, event)
		svc.HandleReceivedRunLifecycleEvent(entry.NodeID, entry.ConnectionGeneration, event)
		if got := metricValue(t, registry, providerHealthEvidenceMetric, map[string]string{"source": "stall", "evidence_health": "unavailable", "decision": "applied"}); got != 1 {
			t.Fatalf("duplicate terminal observations = %v, want 1", got)
		}
	})

	t.Run("default collectors are reused", func(t *testing.T) {
		for range 4 {
			_ = New(nil, edgeevents.NewBus())
		}
	})
}

type blockingProviderHealthObserver struct {
	started chan struct{}
	release chan struct{}
}

func (o *blockingProviderHealthObserver) Observe(providerHealthObservation) {
	close(o.started)
	<-o.release
}

func TestProviderHealthObservabilityRunsAfterQueueUnlock(t *testing.T) {
	svc, entry, record := newProviderHealthOverlayService(t, nil)
	blocking := &blockingProviderHealthObserver{started: make(chan struct{}), release: make(chan struct{})}
	svc.queue.mu.Lock()
	svc.queue.healthObserver = blocking
	svc.queue.mu.Unlock()
	addBoundOverlayLease(t, svc.queue, "run-lock", overlayProviderID, overlayAdapter, overlayTarget, entry.ConnectionGeneration)
	done := make(chan struct{})
	go func() {
		svc.HandleReceivedRunLifecycleEvent(entry.NodeID, entry.ConnectionGeneration, unavailableRunEvent("run-lock", overlayAdapter, overlayTarget, 1))
		close(done)
	}()
	select {
	case <-blocking.started:
	case <-time.After(time.Second):
		t.Fatal("observer was not called")
	}
	// This snapshot needs modelQueueManager.mu. It must complete while the
	// observer remains blocked, proving the post-decision placement.
	snapshotDone := make(chan struct{})
	go func() {
		_ = svc.queue.getSnapshotForNode(entry.NodeID, record, true)
		close(snapshotDone)
	}()
	select {
	case <-snapshotDone:
	case <-time.After(time.Second):
		t.Fatal("observer retained modelQueueManager.mu")
	}
	close(blocking.release)
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("terminal did not return after observer release")
	}
	if leaseCount(svc.queue) != 0 {
		t.Fatal("blocking observer prevented lease release")
	}
}

func TestProviderHealthObservabilityDoesNotExposeSentinels(t *testing.T) {
	svc, entry, _ := newProviderHealthOverlayService(t, nil)
	registry, logs := installProviderHealthTestObserver(t, svc)

	forbiddenValues := []string{
		"SECRET_NODE_ID_12345",
		"SECRET_PROVIDER_ID_67890",
		"SECRET_RUN_ID_ABCDE",
		"SECRET_SESSION_ID_FGHIJ",
		"SECRET_ADAPTER_KEY_KLMNO",
		"SECRET_TARGET_MODEL_PQRST",
		"SECRET_ERROR_MESSAGE_UVWXY",
		"SECRET_PROMPT_BODY_Z0123",
		"SECRET_BEARER_TOKEN_45678",
		"SECRET_EVENT_NODE_ID_11111",
		"SECRET_EVENT_SESSION_ID_22222",
		"SECRET_EVENT_MESSAGE_33333",
		"SECRET_EVENT_ERROR_44444",
		"SECRET_EVENT_DELTA_99999",
		"SECRET_EVENT_NODE_ALIAS_AAAAA",
		"SECRET_FRAME_RUN_ID_55555",
		"SECRET_FRAME_NODE_ID_66666",
		"SECRET_FRAME_TUNNEL_ID_BBBBB",
		"SECRET_FRAME_NODE_ALIAS_CCCCC",
		"SECRET_HEADER_KEY_77777",
		"SECRET_FRAME_ERROR_88888",
	}

	addBoundOverlayLease(t, svc.queue, forbiddenValues[2], overlayProviderID, overlayAdapter, overlayTarget, entry.ConnectionGeneration)

	failure := &iop.ExecutionFailure{
		Code:      string(runtime.FailureCodeResponseStalled),
		Message:   forbiddenValues[6],
		Retryable: true,
		Metadata: map[string]string{
			"failure_code":            string(runtime.FailureCodeResponseStalled),
			"provider_health":         "unavailable",
			"liveness_classification": "provider_unhealthy",
			"idle_duration_ms":        "300000",
			"run_id":                  forbiddenValues[2],
			"session_id":              forbiddenValues[3],
			"adapter":                 overlayAdapter,
			"target":                  overlayTarget,
			"health_observation_seq": "1",
			"node_id":                 forbiddenValues[0],
			"provider_id":             forbiddenValues[1],
			"raw_adapter":             forbiddenValues[4],
			"raw_target":              forbiddenValues[5],
			"body":                    forbiddenValues[7],
			"authorization":           forbiddenValues[8],
		},
	}

	event := &iop.RunEvent{
		RunId:     forbiddenValues[2],
		Type:      "error",
		Delta:     "SECRET_EVENT_DELTA_99999",
		NodeId:    "SECRET_EVENT_NODE_ID_11111",
		NodeAlias: "SECRET_EVENT_NODE_ALIAS_AAAAA",
		SessionId: "SECRET_EVENT_SESSION_ID_22222",
		Message:   "SECRET_EVENT_MESSAGE_33333",
		Error:     "SECRET_EVENT_ERROR_44444",
		Failure:   failure,
		Metadata:  failure.Metadata,
	}

	svc.HandleReceivedRunLifecycleEvent(entry.NodeID, entry.ConnectionGeneration, event)

	frameRunID := "SECRET_FRAME_RUN_ID_55555"
	addBoundOverlayLease(t, svc.queue, frameRunID, overlayProviderID, overlayAdapter, overlayTarget, entry.ConnectionGeneration)

	frameFailure := &iop.ExecutionFailure{
		Code:      string(runtime.FailureCodeResponseStalled),
		Message:   forbiddenValues[6],
		Retryable: true,
		Metadata: map[string]string{
			"failure_code":            string(runtime.FailureCodeResponseStalled),
			"provider_health":         "unavailable",
			"liveness_classification": "provider_unhealthy",
			"idle_duration_ms":        "300000",
			"run_id":                  frameRunID,
			"session_id":              forbiddenValues[3],
			"adapter":                 overlayAdapter,
			"target":                  overlayTarget,
			"health_observation_seq": "2",
			"node_id":                 forbiddenValues[0],
			"provider_id":             forbiddenValues[1],
			"raw_adapter":             forbiddenValues[4],
			"raw_target":              forbiddenValues[5],
			"body":                    forbiddenValues[7],
			"authorization":           forbiddenValues[8],
		},
	}

	frame := &iop.ProviderTunnelFrame{
		RunId:     frameRunID,
		TunnelId:  "SECRET_FRAME_TUNNEL_ID_BBBBB",
		NodeId:    "SECRET_FRAME_NODE_ID_66666",
		NodeAlias: "SECRET_FRAME_NODE_ALIAS_CCCCC",
		Headers:   map[string]string{"SECRET_HEADER_KEY_77777": forbiddenValues[8]},
		Body:      []byte(forbiddenValues[7]),
		Error:     "SECRET_FRAME_ERROR_88888",
		Kind:      iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR,
		Failure:   frameFailure,
		Metadata:  frameFailure.Metadata,
	}

	svc.HandleReceivedProviderTunnelFrame(entry.NodeID, entry.ConnectionGeneration, frame)

	forbiddenValues = append(forbiddenValues, entry.NodeID, overlayProviderID, overlayAdapter, overlayTarget)

	families, err := registry.Gather()
	if err != nil {
		t.Fatalf("gather metrics: %v", err)
	}
	for _, family := range families {
		for _, metric := range family.Metric {
			for _, label := range metric.Label {
				for _, secret := range forbiddenValues {
					if strings.Contains(label.GetName(), secret) || strings.Contains(label.GetValue(), secret) {
						t.Fatalf("forbidden value %q leaked in metric label %s=%s", secret, label.GetName(), label.GetValue())
					}
				}
			}
		}
	}

	for _, entry := range logs.All() {
		for _, secret := range forbiddenValues {
			if strings.Contains(entry.Message, secret) {
				t.Fatalf("forbidden value %q leaked in log message: %s", secret, entry.Message)
			}
			for _, field := range entry.Context {
				if strings.Contains(field.Key, secret) || strings.Contains(fmt.Sprint(field.Interface), secret) || strings.Contains(field.String, secret) {
					t.Fatalf("forbidden value %q leaked in log field %s", secret, field.Key)
				}
			}
		}
	}
}
