package api

import (
	"encoding/json"
	"testing"
	"time"

	fleetai "crypto-miner-server/internal/ai"
	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/models"
)

func TestSurgicalTraceFromSessionHopMatch(t *testing.T) {
	sess := &TraceSession{
		ID: "sess-1",
		Hops: []*HopInfo{
			{AgentID: "hop-a", AgentName: "a"},
			{AgentID: "hop-b", AgentName: "b"},
		},
		DiscoverError: "timeout",
	}
	ctx, ok := surgicalTraceFromSession(sess, "hop-b")
	if !ok {
		t.Fatal("expected trace match")
	}
	if ctx.SessionID != "sess-1" || ctx.HopIndex != 1 || ctx.HopCount != 2 {
		t.Fatalf("unexpected ctx: %+v", ctx)
	}
	if ctx.EgressAgentID != "hop-b" || ctx.DiscoverError != "timeout" {
		t.Fatalf("unexpected egress/discover: %+v", ctx)
	}
}

func TestPathTraceSurgicalAdapterReadsPersistedSession(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })

	payload, _ := json.Marshal(pathTraceSessionPersist{
		ID:        "persisted-sess",
		AgentIDs:  []string{"patient-1"},
		Hops:      []*HopInfo{{AgentID: "patient-1", AgentName: "patient"}},
		CreatedAt: time.Now().UTC(),
		Error:     "spread lane blocked",
	})
	if err := database.UpsertPathTraceSession("persisted-sess", time.Now().UTC(), payload); err != nil {
		t.Fatal(err)
	}

	hub := NewWSHub(database)
	adapter := &PathTraceSurgicalAdapter{Hub: hub, PathTrace: nil}
	ctx, ok := adapter.TraceForAgent("patient-1")
	if !ok {
		t.Fatal("expected persisted trace")
	}
	if ctx.SessionID != "persisted-sess" || ctx.SessionError != "spread lane blocked" {
		t.Fatalf("unexpected ctx: %+v", ctx)
	}
}

func TestHubSeerEmitterPersistsAndBroadcasts(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })

	hub := NewWSHub(database)
	emitter := &HubSeerEmitter{Hub: hub, DB: database}
	if err := emitter.EmitSeerEvent("surgical_replay", "agent-1", map[string]interface{}{
		"failed_tier": "docker",
		"outcome":     "skip_tier:reorder_tiers",
	}); err != nil {
		t.Fatal(err)
	}
	events, err := database.ListSeerEvents(5)
	if err != nil {
		t.Fatal(err)
	}
	if len(events) != 1 || events[0].EventType != "surgical_replay" {
		t.Fatalf("events: %+v", events)
	}
}

func TestStrainFromAgent(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	if err := database.UpsertAgent(&models.Agent{ID: "s1", Name: "host", SpreadStrain: "#aabbcc"}); err != nil {
		t.Fatal(err)
	}
	hub := NewWSHub(database)
	if got := StrainFromAgent(hub, "s1"); got != "#aabbcc" {
		t.Fatalf("strain=%q", got)
	}
}

func TestShouldUseSurgicalReplayPartialNotExhausted(t *testing.T) {
	trace := fleetai.SurgicalTraceContext{SessionID: "x", HopCount: 2}
	snap := fleetai.AgentSnapshot{
		LOTLAttempts: []fleetai.TierAttempt{
			{Tier: "vuln_recon", OK: true},
			{Tier: "docker", OK: false, Error: "missing"},
		},
	}
	if !fleetai.ShouldUseSurgicalReplay(trace, snap) {
		t.Fatal("partial spread failure with trace should qualify for surgical replay")
	}
}