package pipeline

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
)

type executionDebugClient struct {
	response contracts.StructuredCompletionResponse
	err      error
}

func (c executionDebugClient) CompleteStructured(context.Context, contracts.StructuredCompletionRequest, any) (contracts.StructuredCompletionResponse, error) {
	return c.response, c.err
}

func TestScheduledDebugClientPreservesExecutionThroughTerminalRecords(t *testing.T) {
	providerErr := contracts.NewLLMGenerationError(503, "provider unavailable")
	const content = `{"ok":true}`
	for _, test := range []struct {
		name          string
		err           error
		content       string
		debugResponse bool
	}{
		{name: "provider failure without body", err: providerErr},
		{name: "canceled without body", err: context.Canceled},
		{name: "split debug response", content: content, debugResponse: true},
		{name: "split fallback response", content: content},
	} {
		t.Run(test.name, func(t *testing.T) {
			want := contracts.LLMExecutionSummary{
				Attempts: 2, TransientRetries: 1, Outcome: "success",
				CapacityWaitMS: 1, GenerationMS: 3, RetryDelayMS: 5, DurationMS: 9,
				Usage: contracts.LLMObservedTokenUsage{
					PromptTokens:     contracts.LLMObservedTokenCount{Value: 7, Known: true},
					CompletionTokens: contracts.LLMObservedTokenCount{Known: true},
				},
			}
			if test.err != nil {
				want.Outcome = "error"
				want.FailureKind = "provider_status"
				want.StatusCode = 503
				want.RetriesExhausted = true
				if errors.Is(test.err, context.Canceled) {
					want.FailureKind = "canceled"
					want.StatusCode = 0
					want.RetriesExhausted = false
				}
			}
			execution := want
			response := contracts.StructuredCompletionResponse{Execution: &execution}
			if test.debugResponse {
				response.Debug = &contracts.LLMDebugMaterial{Response: &contracts.LLMDebugResponse{Content: test.content}}
			} else {
				response.Content = json.RawMessage(test.content)
			}
			scheduler, err := llm.NewScheduler(1)
			if err != nil {
				t.Fatal(err)
			}
			recorder := newCapturedDebugRecorder()
			client := WithDebugLLMRecording(llm.NewScheduledClient(executionDebugClient{response, test.err}, scheduler), recorder)
			ctx, scope := withDebugLLMScope(context.Background(), "attempt")
			got, err := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: "extract"}, &struct{}{})
			if !errors.Is(err, test.err) || got.Execution == nil || *got.Execution != want {
				t.Fatalf("scheduled debug completion = (%+v, %v), want summary and original error", got, err)
			}
			terminal := newAttemptTerminalRecorder(recorder, "attempt/terminal", "extract", scope, debugTimedEnvelope{Stage: "extract"})
			if err := terminal.record(nil, test.err); !errors.Is(err, test.err) {
				t.Fatalf("terminal recording error = %v", err)
			}
			envelope := recorder.envelope(t, "attempt/response-0001.json")
			encoded, err := json.Marshal(envelope.Payload)
			if err != nil {
				t.Fatal(err)
			}
			var artifact debugLLMResponseArtifact
			if err := json.Unmarshal(encoded, &artifact); err != nil {
				t.Fatal(err)
			}
			if artifact.Execution == nil || *artifact.Execution != want {
				t.Fatalf("persisted response lost execution summary: %+v", artifact)
			}
			if !test.debugResponse && (artifact.Fallback == nil || artifact.Fallback.Execution == nil || *artifact.Fallback.Execution != want) {
				t.Fatalf("persisted fallback lost execution summary: %+v", artifact.Fallback)
			}
			if test.content == "" {
				if artifact.ContentPath != "" {
					t.Fatalf("bodyless failure has content path %q", artifact.ContentPath)
				}
			} else {
				var compact bytes.Buffer
				if artifact.ContentPath == "" || json.Compact(&compact, recorder.bytes[artifact.ContentPath]) != nil || compact.String() != test.content {
					t.Fatalf("split content not preserved at %q", artifact.ContentPath)
				}
				if (artifact.Fallback != nil && artifact.Fallback.Content != "") || (artifact.Response != nil && artifact.Response.Content != "") {
					t.Fatal("split content remains in response metadata")
				}
			}
			refs := recorder.envelope(t, "attempt/terminal.json").LLMCalls
			if len(refs) != 1 || refs[0].Execution == nil || *refs[0].Execution != want || refs[0].Error != (test.err != nil) {
				t.Fatalf("terminal references = %+v", refs)
			}
			// Decode the persisted JSON, not just Go defaults: omitted false and
			// zero fields would erase the presence-aware usage contract.
			usage := envelope.Payload.(map[string]any)["execution"].(map[string]any)["usage"].(map[string]any)
			unknown := usage["cached_tokens"].(map[string]any)
			for _, key := range []string{"value", "known", "complete"} {
				if _, ok := unknown[key]; !ok {
					t.Fatalf("unknown token usage omits %q: %+v", key, unknown)
				}
			}
		})
	}
}