package generation

import (
	"context"
	"errors"
	"net/http"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/promptkit/internal/capacity"
	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
	"gitea.maximumdirect.net/eric/promptkit/internal/llm"
)

func TestGenerateWithEvidenceRetainsAttemptsAndUsage(t *testing.T) {
	var now = time.Date(2026, 10, 13, 12, 0, 0, 0, time.UTC)
	calls := 0
	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
		calls++
		if calls == 1 {
			now = now.Add(4 * time.Millisecond)
			return response(http.StatusServiceUnavailable, `{"error":{}}`), nil
		}
		now = now.Add(7 * time.Millisecond)
		return response(http.StatusOK, `{"choices":[{"message":{"content":"ok"}}],"usage":{"prompt_tokens":3}}`), nil
	})
	coordinator.now = func() time.Time { return now }
	coordinator.wait = func(context.Context, time.Duration) error {
		now = now.Add(19 * time.Millisecond)
		return nil
	}

	result, err, evidence := coordinator.GenerateWithEvidence(context.Background(), testRequest(1))
	if err != nil || result == nil || result.Content != "ok" || calls != 2 {
		t.Fatalf("result=%+v error=%v calls=%d", result, err, calls)
	}
	if len(evidence.Attempts) != 2 || len(evidence.RetryDelays) != 1 || evidence.RetryDelays[0] != 19*time.Millisecond {
		t.Fatalf("attempts/delays = %#v / %#v", evidence.Attempts, evidence.RetryDelays)
	}
	first, second := evidence.Attempts[0], evidence.Attempts[1]
	if !first.Dispatched || !first.HTTP.RequestObserved || !first.HTTP.ResponseReceived || first.HTTP.HTTPStatus == nil || *first.HTTP.HTTPStatus != http.StatusServiceUnavailable {
		t.Fatalf("first attempt evidence = %+v", first)
	}
	if first.Failure != domain.ExecutionFailureProviderStatus || first.RequestDuration != 4*time.Millisecond {
		t.Fatalf("first attempt failure/timing = %v / %v", first.Failure, first.RequestDuration)
	}
	if !second.Dispatched || !second.GenerationCompleted || !second.ResponseDecoded || second.HTTP.HTTPStatus == nil || *second.HTTP.HTTPStatus != http.StatusOK {
		t.Fatalf("second attempt evidence = %+v", second)
	}
	if second.Failure != domain.ExecutionFailureNone || second.RequestDuration != 7*time.Millisecond {
		t.Fatalf("second attempt failure/timing = %v / %v", second.Failure, second.RequestDuration)
	}
	if !second.Usage.PromptTokens.Known || !second.Usage.PromptTokens.Complete || second.Usage.PromptTokens.Value != 3 {
		t.Fatalf("per-attempt usage = %+v", second.Usage.PromptTokens)
	}
	if !evidence.Usage.PromptTokens.Known || evidence.Usage.PromptTokens.Complete || evidence.Usage.PromptTokens.Value != 3 {
		t.Fatalf("aggregate usage = %+v", evidence.Usage.PromptTokens)
	}
}

func TestGenerateWithEvidenceRetainsExhaustedAttempts(t *testing.T) {
	calls := 0
	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
		calls++
		return response(http.StatusServiceUnavailable, `{"error":{}}`), nil
	})
	coordinator.wait = func(context.Context, time.Duration) error { return nil }

	result, err, evidence := coordinator.GenerateWithEvidence(context.Background(), testRequest(1))
	var exhausted *ExhaustedError
	if result != nil || !errors.As(err, &exhausted) || calls != 2 || exhausted.Attempts() != 2 {
		t.Fatalf("result=%+v error=%v calls=%d", result, err, calls)
	}
	if len(evidence.Attempts) != 2 || len(evidence.RetryDelays) != 1 {
		t.Fatalf("attempt/delay evidence = %#v / %#v", evidence.Attempts, evidence.RetryDelays)
	}
	for _, attempt := range evidence.Attempts {
		if !attempt.Dispatched || attempt.Failure != domain.ExecutionFailureProviderStatus || attempt.HTTP.HTTPStatus == nil || *attempt.HTTP.HTTPStatus != http.StatusServiceUnavailable {
			t.Fatalf("exhausted attempt evidence = %+v", attempt)
		}
	}
	if evidence.Usage.PromptTokens.Known || evidence.Usage.PromptTokens.Complete {
		t.Fatalf("missing provider usage was claimed: %+v", evidence.Usage.PromptTokens)
	}
}

func TestGenerateWithEvidenceRetainsInterruptedBackoffWithoutAttempt(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var now = time.Date(2026, 10, 13, 12, 0, 0, 0, time.UTC)
	calls := 0
	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
		calls++
		now = now.Add(3 * time.Millisecond)
		return response(http.StatusServiceUnavailable, `{"error":{}}`), nil
	})
	coordinator.now = func() time.Time { return now }
	coordinator.wait = func(ctx context.Context, _ time.Duration) error {
		now = now.Add(11 * time.Millisecond)
		cancel()
		return ctx.Err()
	}

	result, err, evidence := coordinator.GenerateWithEvidence(ctx, testRequest(2))
	if result != nil || !errors.Is(err, context.Canceled) || calls != 1 {
		t.Fatalf("result=%+v error=%v calls=%d", result, err, calls)
	}
	if len(evidence.Attempts) != 1 || !evidence.Attempts[0].Dispatched || len(evidence.RetryDelays) != 1 || evidence.RetryDelays[0] != 11*time.Millisecond {
		t.Fatalf("interrupted evidence = %#v", evidence)
	}
	if evidence.Attempts[0].HTTP.HTTPStatus == nil || *evidence.Attempts[0].HTTP.HTTPStatus != http.StatusServiceUnavailable {
		t.Fatalf("last dispatched response was lost: %+v", evidence.Attempts[0])
	}
}

func TestGenerateWithEvidenceRecordsCanceledReacquisition(t *testing.T) {
	var now = time.Date(2026, 10, 13, 12, 0, 0, 0, time.UTC)
	calls := 0
	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
		calls++
		now = now.Add(2 * time.Millisecond)
		return response(http.StatusServiceUnavailable, `{"error":{}}`), nil
	})
	coordinator.now = func() time.Time { return now }
	baseContext, cancel := context.WithCancel(context.Background())
	defer cancel()
	contextWithCancelOnWait := &cancelOnDoneContext{Context: baseContext, cancel: cancel}
	var releaseHeld func()
	coordinator.wait = func(context.Context, time.Duration) error {
		release, err := coordinator.manager.Acquire(context.Background(), "limited")
		if err != nil {
			t.Fatalf("hold permit: %v", err)
		}
		releaseHeld = release
		contextWithCancelOnWait.armed.Store(true)
		return nil
	}

	// Signal cancellation only when the second acquisition has queued behind
	// the held permit, avoiding scheduler- or sleep-based synchronization.
	result, err, evidence := coordinator.GenerateWithEvidence(contextWithCancelOnWait, testRequest(2))
	if releaseHeld != nil {
		releaseHeld()
	}
	if result != nil || !errors.Is(err, context.Canceled) || calls != 1 {
		t.Fatalf("result=%+v error=%v calls=%d", result, err, calls)
	}
	if len(evidence.Attempts) != 2 || !evidence.Attempts[0].Dispatched || evidence.Attempts[1].Dispatched {
		t.Fatalf("acquisition evidence = %#v", evidence.Attempts)
	}
	if evidence.Attempts[1].Failure != domain.ExecutionFailureCanceled {
		t.Fatalf("canceled acquisition failure = %v", evidence.Attempts[1].Failure)
	}
	if evidence.Attempts[0].HTTP.HTTPStatus == nil || *evidence.Attempts[0].HTTP.HTTPStatus != http.StatusServiceUnavailable {
		t.Fatalf("last dispatched evidence was lost: %+v", evidence.Attempts)
	}
}

func TestGenerateWithEvidenceNoDispatchHasCompleteZeroUsage(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
		t.Fatal("canceled generation dispatched")
		return nil, nil
	})

	result, err, evidence := coordinator.GenerateWithEvidence(ctx, testRequest(0))
	if result != nil || !errors.Is(err, context.Canceled) || len(evidence.Attempts) != 0 {
		t.Fatalf("result=%+v error=%v evidence=%+v", result, err, evidence)
	}
	if !evidence.Usage.PromptTokens.Known || !evidence.Usage.PromptTokens.Complete || evidence.Usage.PromptTokens.Value != 0 {
		t.Fatalf("no-dispatch usage = %+v", evidence.Usage.PromptTokens)
	}
}

func TestGenerateWithEvidenceForInjectedClient(t *testing.T) {
	opaqueFailure := errors.New("injected client failed")
	presence := &domain.TokenUsagePresence{PromptTokens: true}
	for _, tc := range []struct {
		name           string
		response       *domain.GenerateResponse
		err            error
		wantResponse   bool
		wantGeneration bool
		wantFailure    domain.ExecutionFailureKind
		wantUsage      bool
		wantStructured bool
	}{
		{
			name:           "success",
			response:       &domain.GenerateResponse{Content: "private output", Usage: domain.TokenUsage{PromptTokens: 4}, UsagePresence: presence},
			wantResponse:   true,
			wantGeneration: true,
			wantUsage:      true,
			wantStructured: true,
		},
		{
			name:        "nil response",
			wantFailure: domain.ExecutionFailureMalformedResponse,
		},
		{
			name:        "opaque failure",
			err:         opaqueFailure,
			wantFailure: domain.ExecutionFailureGeneration,
		},
		{
			name:           "response with error",
			response:       &domain.GenerateResponse{Content: "must not be evaluated", Usage: domain.TokenUsage{PromptTokens: 9}, UsagePresence: presence},
			err:            opaqueFailure,
			wantResponse:   true,
			wantFailure:    domain.ExecutionFailureGeneration,
			wantStructured: true,
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			client := &evidenceInjectedClient{response: tc.response, err: tc.err}
			coordinator := newEvidenceInjectedCoordinator(t, client)
			req := testRequest(7)
			if tc.wantStructured {
				req.StructuredOutput = &domain.StructuredOutputSpec{}
			}

			result, err, evidence := coordinator.GenerateWithEvidence(context.Background(), req)
			if client.calls != 1 || len(evidence.Attempts) != 1 || len(evidence.RetryDelays) != 0 {
				t.Fatalf("calls=%d attempts=%+v retry delays=%v", client.calls, evidence.Attempts, evidence.RetryDelays)
			}
			attempt := evidence.Attempts[0]
			if !attempt.Dispatched || attempt.ResponseReceived != tc.wantResponse || attempt.GenerationCompleted != tc.wantGeneration || attempt.Failure != tc.wantFailure {
				t.Fatalf("attempt evidence = %+v", attempt)
			}
			if attempt.HTTP != (llm.HTTPExecutionEvidence{}) {
				t.Fatalf("injected call claimed HTTP evidence: %+v", attempt.HTTP)
			}
			if attempt.StructuredOutputRequested == nil || *attempt.StructuredOutputRequested != tc.wantStructured {
				t.Fatalf("structured-output request fact = %v", attempt.StructuredOutputRequested)
			}
			if tc.wantFailure == domain.ExecutionFailureNone {
				if err != nil || result == nil || result.Content != "private output" {
					t.Fatalf("result=%+v error=%v", result, err)
				}
			} else if result != nil || err == nil {
				t.Fatalf("result=%+v error=%v; expected failed call", result, err)
			}
			if attempt.Usage.PromptTokens.Known != tc.wantUsage || evidence.Usage.PromptTokens.Known != tc.wantUsage {
				t.Fatalf("attempt/aggregate usage = %+v / %+v", attempt.Usage, evidence.Usage)
			}
			if tc.wantUsage && (attempt.Usage.PromptTokens.Value != 4 || evidence.Usage.PromptTokens.Value != 4) {
				t.Fatalf("reported usage was not retained: %+v / %+v", attempt.Usage, evidence.Usage)
			}
		})
	}
}

func TestGenerateWithEvidenceRetainsInjectedSuccessAfterCancellation(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	client := &evidenceInjectedClient{
		response:      &domain.GenerateResponse{Content: "private output"},
		afterGenerate: func() { cancel() },
	}
	coordinator := newEvidenceInjectedCoordinator(t, client)
	req := testRequest(5)
	req.StructuredOutput = &domain.StructuredOutputSpec{}

	result, err, evidence := coordinator.GenerateWithEvidence(ctx, req)
	if !errors.Is(err, context.Canceled) || result == nil || result.Content != "private output" {
		t.Fatalf("result=%+v error=%v", result, err)
	}
	if client.calls != 1 || len(evidence.Attempts) != 1 {
		t.Fatalf("calls=%d attempts=%+v", client.calls, evidence.Attempts)
	}
	attempt := evidence.Attempts[0]
	if !attempt.ResponseReceived || !attempt.GenerationCompleted || attempt.Failure != domain.ExecutionFailureNone {
		t.Fatalf("successful generation evidence was lost: %+v", attempt)
	}
	if attempt.HTTP != (llm.HTTPExecutionEvidence{}) {
		t.Fatalf("injected call claimed HTTP evidence: %+v", attempt.HTTP)
	}
}

type evidenceInjectedClient struct {
	calls         int
	response      *domain.GenerateResponse
	err           error
	afterGenerate func()
}

func (c *evidenceInjectedClient) Generate(context.Context, domain.GenerateRequest) (*domain.GenerateResponse, error) {
	c.calls++
	if c.afterGenerate != nil {
		c.afterGenerate()
	}
	return c.response, c.err
}

func newEvidenceInjectedCoordinator(t *testing.T, client llm.Client) *Coordinator {
	t.Helper()
	manager, err := capacity.NewManager(map[string]domain.BackendCapacityPolicy{
		"limited": {ConcurrencyLimit: 1},
	})
	if err != nil {
		t.Fatal(err)
	}
	return NewInjected(manager, client)
}

type cancelOnDoneContext struct {
	context.Context
	cancel context.CancelFunc
	once   sync.Once
	armed  atomic.Bool
}

func (c *cancelOnDoneContext) Done() <-chan struct{} {
	if c.armed.Load() {
		c.once.Do(c.cancel)
	}
	return c.Context.Done()
}