package llm

import (
	"context"
	"encoding/json"
	"strings"
	"sync"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
	"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
	pk "gitea.maximumdirect.net/eric/promptkit"
)

func TestObservationSummarySeparatesLogicalCallsAttemptsAndUsage(t *testing.T) {
	collector := NewObservationCollector()
	hooks := collector.Hooks()
	now := time.Now().UTC()
	call := func(id uint64) {
		hooks.OnCall(structuredoutput.CallRecord{CallID: id, Context: structuredoutput.DiagnosticContext{Role: "proposal"}, Started: now})
		hooks.OnCall(structuredoutput.CallRecord{CallID: id, Context: structuredoutput.DiagnosticContext{Role: "proposal"}, Started: now, Finished: now.Add(time.Second)})
	}
	call(1) // preparation failed before PromptKit emitted events
	call(2) // capacity rejected without a dispatched attempt
	hooks.OnEventWithCall(2, structuredoutput.DiagnosticContext{Role: "proposal"}, pk.ExecutionEvent{Kind: pk.ExecutionAttempt, Dispatched: false})
	hooks.OnEventWithCall(2, structuredoutput.DiagnosticContext{Role: "proposal"}, pk.ExecutionEvent{Kind: pk.ExecutionFinished, Attempts: 0})
	summary := Summarize(collector.Snapshot())
	if summary.LogicalCalls != 2 || summary.ProviderAttempts != 0 || summary.Usage.Total != (TokenCount{Known: true, Complete: true}) {
		t.Fatalf("zero-attempt accounting: %+v", summary)
	}

	call(3)
	context := structuredoutput.DiagnosticContext{Role: "validation", ModuleInstance: "grammar", ValidatorName: "editorial_review"}
	hooks.OnEventWithCall(3, context, pk.ExecutionEvent{Kind: pk.ExecutionAttempt, Dispatched: true, Usage: pk.ObservedTokenUsage{TotalTokens: pk.ObservedTokenCount{Value: 100, Known: true, Complete: true}}})
	hooks.OnEventWithCall(3, context, pk.ExecutionEvent{Kind: pk.ExecutionFinished, Attempts: 3, TransientRetries: 1, RepairCalls: 1, Usage: pk.ObservedTokenUsage{
		PromptTokens: pk.ObservedTokenCount{Value: 5, Known: true, Complete: false},
		TotalTokens:  pk.ObservedTokenCount{Value: 8, Known: true, Complete: false},
	}})
	call(4)
	hooks.OnEventWithCall(4, context, pk.ExecutionEvent{Kind: pk.ExecutionFinished, Attempts: 1, Usage: pk.ObservedTokenUsage{PromptTokens: pk.ObservedTokenCount{Known: false, Complete: false}, TotalTokens: pk.ObservedTokenCount{Known: false, Complete: false}}})
	snapshot := collector.Snapshot()
	summary = Summarize(snapshot)
	if summary.LogicalCalls != 4 || summary.ProviderAttempts != 4 || summary.TransientRetries != 1 || summary.StructuralRepairs != 1 || summary.Usage.Total.Value != 8 || !summary.Usage.Total.Known || summary.Usage.Total.Complete || summary.Usage.Prompt.Value != 5 || !summary.Usage.Prompt.Known || summary.Usage.Prompt.Complete || summary.Usage.Cached.Known || summary.Usage.Cached.Complete {
		t.Fatalf("partial usage or double count: %+v", summary)
	}
	raw, err := json.Marshal(snapshot)
	if err != nil {
		t.Fatal(err)
	}
	for _, forbidden := range []string{"https://", "Bearer ", "raw_output", "messages", "api_key"} {
		if strings.Contains(string(raw), forbidden) {
			t.Fatalf("observation leaked payload field %q", forbidden)
		}
	}
}

func TestObservationCollectorConcurrentAttributionAndSnapshotCopy(t *testing.T) {
	collector := NewObservationCollector()
	hooks := collector.Hooks()
	const calls = 64
	var wg sync.WaitGroup
	for i := 1; i <= calls; i++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			index := id
			context := structuredoutput.DiagnosticContext{Role: "proposal", ModuleInstance: "grammar", SectionIndex: &index}
			started := time.Now().UTC()
			hooks.OnCall(structuredoutput.CallRecord{CallID: uint64(id), Context: context, Started: started})
			hooks.OnEventWithCall(uint64(id), context, pk.ExecutionEvent{Kind: pk.ExecutionFinished, RunID: "run", Attempts: 1, Usage: pk.ObservedTokenUsage{TotalTokens: pk.ObservedTokenCount{Value: 1, Known: true, Complete: true}}})
			hooks.OnCall(structuredoutput.CallRecord{CallID: uint64(id), Context: context, Started: started, Finished: time.Now().UTC()})
			index = -1
		}(i)
	}
	wg.Wait()
	snapshot := collector.Snapshot()
	if len(snapshot.Calls) != calls || len(snapshot.Events) != calls || Summarize(snapshot).ProviderAttempts != calls {
		t.Fatalf("parallel records: calls=%d events=%d", len(snapshot.Calls), len(snapshot.Events))
	}
	for _, call := range snapshot.Calls {
		if call.Context.SectionIndex == nil || *call.Context.SectionIndex != int(call.CallID) {
			t.Fatalf("attribution lost: %+v", call)
		}
	}
	*snapshot.Calls[0].Context.SectionIndex = -2
	again := collector.Snapshot()
	if *again.Calls[0].Context.SectionIndex == -2 {
		t.Fatal("snapshot aliases collector context")
	}
}

func TestRoleClientFeedsCollectorAcrossPreparationFailure(t *testing.T) {
	runtime, count, _ := adapterFixture(t, []reply{{status: 200, content: `{"corrections":[]}`}}, 0, 0, nil)
	collector := NewObservationCollector()
	client, err := runtime.NewRoleClient("proposal", collector.Hooks())
	if err != nil {
		t.Fatal(err)
	}
	request := adapterRequest(responseschema.CorrectionSetKey)
	request.Context.ModuleInstance = "glossary_1"
	index := 3
	request.Context.SectionIndex = &index
	missing := request
	missing.PromptID = "missing.prompt"
	var output map[string]any
	if _, err := client.CompleteStructured(context.Background(), missing, &output); err == nil {
		t.Fatal("preparation error missing")
	}
	if _, err := client.CompleteStructured(context.Background(), request, &output); err != nil {
		t.Fatal(err)
	}
	snapshot := collector.Snapshot()
	summary := Summarize(snapshot)
	if count.Load() != 1 || summary.LogicalCalls != 2 || summary.ProviderAttempts != 1 || len(snapshot.Events) == 0 {
		t.Fatalf("collector accounting: %+v", summary)
	}
	for _, event := range snapshot.Events {
		if event.Context.ModuleInstance != "glossary_1" || event.Context.SectionIndex == nil || *event.Context.SectionIndex != 3 || event.Context.Role != "proposal" {
			t.Fatalf("event attribution: %+v", event)
		}
	}
	if snapshot.Calls[0].Failure != structuredoutput.FailureRuntime || snapshot.Calls[1].Failure != "" {
		t.Fatalf("logical failures: %+v", snapshot.Calls)
	}
}