package openai

import (
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"
	"testing"
	"time"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"go.uber.org/zap/zaptest/observer"
	"iop/packages/go/config"
)

func TestLogicalRequestTTLSweep(t *testing.T) {
	now := time.Unix(100, 0)
	coordinator := newLogicalRequestCoordinator(logicalRequestCoordinatorOptions{
		TTL: time.Second, Now: func() time.Time { return now }, IDSource: sequentialLogicalRequestIDs("ttl_sweep"),
	})
	lineage := mustChatLogicalRequestLineage(t, "ttl", "tool")
	var ids []string
	for i := 0; i < 3; i++ {
		request, err := coordinator.create(logicalRequestAdmission{
			OwnerEdgeID: "edge-a", PrincipalRef: "principal-a", Lineage: lineage, PresetGeneration: "preset-gen-1",
		})
		if err != nil {
			t.Fatal(err)
		}
		ids = append(ids, request.ID)
	}
	now = now.Add(2 * time.Second)
	first := coordinator.sweepExpired(now, 2)
	if len(first) != 2 {
		t.Fatalf("first sweep=%d, want 2", len(first))
	}
	got := []string{first[0].RequestID, first[1].RequestID}
	want := append([]string(nil), ids...)
	sort.Strings(want)
	if got[0] != want[0] || got[1] != want[1] {
		t.Fatalf("bounded deterministic sweep=%v, want prefix %v", got, want[:2])
	}
	second := coordinator.sweepExpired(now, 2)
	if len(second) != 1 || second[0].RequestID != want[2] {
		t.Fatalf("second sweep=%+v, want %q", second, want[2])
	}
}

func TestLogicalRequestTTLActiveSurvives(t *testing.T) {
	now := time.Unix(200, 0)
	coordinator := newLogicalRequestCoordinator(logicalRequestCoordinatorOptions{
		TTL: time.Second, Now: func() time.Time { return now }, IDSource: sequentialLogicalRequestIDs("ttl_active"),
	})
	request, err := coordinator.create(logicalRequestAdmission{
		OwnerEdgeID: "edge-a", PrincipalRef: "principal-a", Lineage: mustChatLogicalRequestLineage(t, "active", "tool"), PresetGeneration: "preset-gen-1",
	})
	if err != nil {
		t.Fatal(err)
	}
	stageID, _ := coordinator.newStageID()
	if _, err := coordinator.activateStage(request.ID, "edge-a", stageID); err != nil {
		t.Fatal(err)
	}
	now = now.Add(10 * time.Second)
	if expired := coordinator.sweepExpired(now, 8); len(expired) != 0 {
		t.Fatalf("active request was swept: %+v", expired)
	}
	if _, err := coordinator.snapshot(request.ID); err != nil {
		t.Fatalf("active snapshot: %v", err)
	}
	if err := coordinator.disconnect(request.ID, "edge-a", "cancelled"); err != nil {
		t.Fatal(err)
	}
	now = now.Add(2 * time.Second)
	expired := coordinator.sweepExpired(now, 8)
	if len(expired) != 1 || expired[0].RequestID != request.ID || expired[0].PriorState != logicalRequestStateDetached || expired[0].TerminalClass != "cancelled" {
		t.Fatalf("detached sweep=%+v", expired)
	}
}

func TestLogicalRequestTTLFinalizeRace(t *testing.T) {
	for iteration := 0; iteration < 32; iteration++ {
		now := time.Unix(300, 0)
		coordinator := newLogicalRequestCoordinator(logicalRequestCoordinatorOptions{
			TTL: time.Second, Now: func() time.Time { return now }, IDSource: sequentialLogicalRequestIDs(fmt.Sprintf("ttl_race_%d", iteration)),
		})
		lineage := mustChatLogicalRequestLineage(t, "race", "tool")
		request, err := coordinator.create(logicalRequestAdmission{
			OwnerEdgeID: "edge-a", PrincipalRef: "principal-a", Lineage: lineage, PresetGeneration: "preset-gen-1",
		})
		if err != nil {
			t.Fatal(err)
		}
		stageID, _ := coordinator.newStageID()
		cleanupStageID, _ := coordinator.newStageID()
		if _, err := coordinator.activateStage(request.ID, "edge-a", stageID); err != nil {
			t.Fatal(err)
		}
		if _, err := coordinator.startCleanup(request.ID, "edge-a", stageID, cleanupStageID, "success"); err != nil {
			t.Fatal(err)
		}
		const issuedHash = "cleanup_race_hash"
		if _, err := coordinator.awaitToolResults(request.ID, "edge-a", cleanupStageID, []logicalRequestExpectedTool{{
			PublicCallID: "call_cleanup", ProviderCallID: "provider_cleanup",
		}}, issuedHash); err != nil {
			t.Fatal(err)
		}
		continuation := logicalRequestContinuationLineage{
			Prefix: lineage, IssuedCallHash: issuedHash, ResultIDs: []string{"call_cleanup"},
			Committed: logicalRequestLineage{Endpoint: lineage.Endpoint, HistoryDigest: fmt.Sprintf("cleanup_committed_%d", iteration), ToolsetDigest: lineage.ToolsetDigest},
		}
		now = now.Add(2 * time.Second)
		start := make(chan struct{})
		var commitErr error
		var expired []logicalRequestExpirySnapshot
		var wg sync.WaitGroup
		wg.Add(2)
		go func() {
			defer wg.Done()
			<-start
			_, commitErr = coordinator.commitCleanupByLineage("edge-a", "principal-a", continuation)
		}()
		go func() {
			defer wg.Done()
			<-start
			expired = coordinator.sweepExpired(now, 1)
		}()
		close(start)
		wg.Wait()
		commitWon := commitErr == nil
		sweepWon := len(expired) == 1
		if commitWon == sweepWon {
			t.Fatalf("iteration %d owners: commitErr=%v expired=%+v", iteration, commitErr, expired)
		}
		if !commitWon && !errors.Is(commitErr, errLogicalRequestNotFound) {
			t.Fatalf("iteration %d commit error=%v", iteration, commitErr)
		}
		coordinator.mu.Lock()
		remaining := len(coordinator.requests)
		coordinator.mu.Unlock()
		if remaining != 0 {
			t.Fatalf("iteration %d remaining=%d", iteration, remaining)
		}
	}
}

func TestLogicalRequestTTLObservationRedaction(t *testing.T) {
	core, observed := observer.New(zapcore.InfoLevel)
	server := NewServer(config.EdgeOpenAIConf{}, nil, zap.New(core))
	server.SetEdgeID("edge-ttl")
	now := time.Unix(400, 0)
	server.requestCoordinator = newLogicalRequestCoordinator(logicalRequestCoordinatorOptions{
		TTL: time.Second, Now: func() time.Time { return now }, IDSource: sequentialLogicalRequestIDs("ttl_redaction"),
	})
	lineage := mustChatLogicalRequestLineage(t, "PROMPT_SENTINEL", "tool")
	request, err := server.requestCoordinator.create(logicalRequestAdmission{
		OwnerEdgeID: server.edgeIDValue(), PrincipalRef: "principal-a", Lineage: lineage, PresetGeneration: "preset-gen-1",
	})
	if err != nil {
		t.Fatal(err)
	}
	server.lightFlows.mu.Lock()
	server.lightFlows.records[request.ID] = &hotPathLightRecord{
		requestID: request.ID, ownerEdgeID: server.edgeIDValue(), immutableTask: "PROMPT_SENTINEL",
		pendingOutput: normalizedStageOutput{Content: "CONTENT_SENTINEL", Reasoning: "CREDENTIAL_SENTINEL"},
	}
	server.lightFlows.mu.Unlock()
	server.artifactFrontiers.mu.Lock()
	server.artifactFrontiers.records[request.ID] = &artifactFrontierRecord{requestID: request.ID, ownerEdgeID: server.edgeIDValue()}
	server.artifactFrontiers.mu.Unlock()

	now = now.Add(2 * time.Second)
	server.sweepLogicalRequestTTL()
	entries := observed.FilterMessage(hotPathOrphanObservationMessage).All()
	if len(entries) != 1 {
		t.Fatalf("orphan observations=%d, want 1", len(entries))
	}
	fields := entries[0].ContextMap()
	wantKeys := []string{"prior_state", "reason", "request_id", "stage", "terminal_class", "workspace_path"}
	gotKeys := make([]string, 0, len(fields))
	for key := range fields {
		gotKeys = append(gotKeys, key)
	}
	sort.Strings(gotKeys)
	if strings.Join(gotKeys, ",") != strings.Join(wantKeys, ",") {
		t.Fatalf("observation keys=%v, want %v", gotKeys, wantKeys)
	}
	if fields["request_id"] != request.ID || fields["workspace_path"] != newReservedPaths(request.ID).JobDir+"/" ||
		fields["reason"] != hotPathOrphanReasonTTL {
		t.Fatalf("observation fields=%v", fields)
	}
	serialized := fmt.Sprint(fields)
	for _, forbidden := range []string{"PROMPT_SENTINEL", "CONTENT_SENTINEL", "CREDENTIAL_SENTINEL", "principal-a"} {
		if strings.Contains(serialized, forbidden) {
			t.Fatalf("observation leaked %q: %s", forbidden, serialized)
		}
	}
	server.lightFlows.mu.Lock()
	lightCount := len(server.lightFlows.records)
	server.lightFlows.mu.Unlock()
	server.artifactFrontiers.mu.Lock()
	artifactCount := len(server.artifactFrontiers.records)
	server.artifactFrontiers.mu.Unlock()
	if lightCount != 0 || artifactCount != 0 {
		t.Fatalf("matching stores not removed: light=%d artifact=%d", lightCount, artifactCount)
	}
}
