package generation

import (
	"context"
	"errors"
	"io"
	"math"
	"net/http"
	"strings"
	"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"
)

type transportFunc func(*http.Request) (*http.Response, error)

func (f transportFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) }

func response(status int, body string) *http.Response {
	return &http.Response{StatusCode: status, Body: io.NopCloser(strings.NewReader(body)), Header: make(http.Header)}
}

func testRequest(retries uint64) domain.GenerateRequest {
	return domain.GenerateRequest{
		Prompt:           domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hello"}}},
		Target:           domain.ExecutionTarget{BackendID: "limited", Endpoint: "https://example.invalid/v1", Model: "model"},
		TransientRetries: retries,
	}
}

func testCoordinator(t *testing.T, transport transportFunc) *Coordinator {
	t.Helper()
	manager, err := capacity.NewManager(map[string]domain.BackendCapacityPolicy{
		"limited": {ConcurrencyLimit: 1, QueueCapacity: 1},
	})
	if err != nil {
		t.Fatal(err)
	}
	client, err := llm.NewOpenAICompatibleClient(llm.OpenAICompatibleConfig{HTTPClient: &http.Client{Transport: transport}})
	if err != nil {
		t.Fatal(err)
	}
	coordinator := NewBuiltIn(manager, client)
	coordinator.jitter = func(time.Duration) time.Duration { return 0 }
	coordinator.wait = func(ctx context.Context, _ time.Duration) error { return ctx.Err() }
	return coordinator
}

func TestBuiltInRetryEligibilityAndExhaustion(t *testing.T) {
	for _, tc := range []struct {
		status   int
		eligible bool
	}{
		{status: 429, eligible: true}, {status: 500, eligible: true},
		{status: 502, eligible: true}, {status: 503, eligible: true},
		{status: 504, eligible: true}, {status: 400}, {status: 501},
	} {
		t.Run(http.StatusText(tc.status), func(t *testing.T) {
			calls := 0
			coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
				calls++
				if calls == 1 {
					return response(tc.status, `{"error":{}}`), nil
				}
				return response(200, `{"choices":[{"message":{"content":"ok"}}]}`), nil
			})
			result, err := coordinator.Generate(context.Background(), testRequest(1))
			if tc.eligible {
				if err != nil || result == nil || result.Content != "ok" || calls != 2 {
					t.Fatalf("result=%+v, error=%v, calls=%d", result, err, calls)
				}
			} else if result != nil || err == nil || calls != 1 {
				t.Fatalf("result=%+v, error=%v, calls=%d", result, err, calls)
			}
		})
	}

	for _, budget := range []uint64{0, 2} {
		calls := 0
		coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
			calls++
			return response(503, `{"error":{}}`), nil
		})
		_, err := coordinator.Generate(context.Background(), testRequest(budget))
		var exhausted *ExhaustedError
		if !errors.Is(err, llm.ErrUnexpectedStatus) || calls != int(budget+1) || errors.As(err, &exhausted) != (budget > 0) {
			t.Fatalf("budget=%d calls=%d error=%v", budget, calls, err)
		}
		if budget > 0 && exhausted.Attempts() != budget+1 {
			t.Fatalf("attempts=%d", exhausted.Attempts())
		}
	}

	calls := 0
	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
		calls++
		if calls == 1 {
			return response(503, `{"error":{}}`), nil
		}
		return response(400, `{"error":{}}`), nil
	})
	_, err := coordinator.Generate(context.Background(), testRequest(3))
	var exhausted *ExhaustedError
	if calls != 2 || !errors.Is(err, llm.ErrUnexpectedStatus) || errors.As(err, &exhausted) {
		t.Fatalf("later permanent failure: calls=%d, error=%v", calls, err)
	}
}

func TestBuiltInRetryWaitHonorsContextAndProviderMinimum(t *testing.T) {
	when := time.Date(2026, 9, 23, 12, 0, 0, 0, time.UTC)
	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) { return response(503, `{"error":{}}`), nil })
	if got := coordinator.backoff(1); got != 250*time.Millisecond {
		t.Fatalf("first backoff=%v", got)
	}
	if got := coordinator.backoff(math.MaxUint64); got != 4*time.Second {
		t.Fatalf("capped backoff=%v", got)
	}
	coordinator.jitter = func(max time.Duration) time.Duration { return max }
	if got := coordinator.backoff(math.MaxUint64); got != 8*time.Second {
		t.Fatalf("maximum equal jitter=%v", got)
	}
	if got := retryDelay(500*time.Millisecond, when.Add(3*time.Second), when); got != 3*time.Second {
		t.Fatalf("provider minimum=%v", got)
	}
	if got := retryDelay(500*time.Millisecond, when.Add(-time.Second), when); got != 500*time.Millisecond {
		t.Fatalf("past provider minimum=%v", got)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	calls := 0
	coordinator = testCoordinator(t, func(*http.Request) (*http.Response, error) {
		calls++
		return response(503, `{"error":{}}`), nil
	})
	coordinator.wait = func(ctx context.Context, delay time.Duration) error {
		if delay < 250*time.Millisecond {
			t.Fatalf("delay=%v", delay)
		}
		cancel()
		return ctx.Err()
	}
	_, err := coordinator.Generate(ctx, testRequest(2))
	if !errors.Is(err, context.Canceled) || calls != 1 {
		t.Fatalf("canceled wait: error=%v calls=%d", err, calls)
	}
}

func TestBuiltInRetryUsesSeparateAttemptDeadlines(t *testing.T) {
	var contexts []context.Context
	coordinator := testCoordinator(t, func(req *http.Request) (*http.Response, error) {
		contexts = append(contexts, req.Context())
		if len(contexts) == 1 {
			return response(503, `{"error":{}}`), nil
		}
		return response(200, `{"choices":[{"message":{"content":"ok"}}]}`), nil
	})
	req := testRequest(1)
	req.Target.TimeoutSeconds = 1
	if _, err := coordinator.Generate(context.Background(), req); err != nil {
		t.Fatal(err)
	}
	if len(contexts) != 2 || contexts[0] == contexts[1] {
		t.Fatalf("attempt contexts=%d", len(contexts))
	}
	for _, ctx := range contexts {
		if _, ok := ctx.Deadline(); !ok || !errors.Is(ctx.Err(), context.Canceled) {
			t.Fatalf("attempt deadline/error=%t/%v", ok, ctx.Err())
		}
	}
}

func TestCanceledCapacityWaitDispatchesNothing(t *testing.T) {
	calls := 0
	coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
		calls++
		return response(200, `{"choices":[{"message":{"content":"ok"}}]}`), nil
	})
	release, err := coordinator.manager.Acquire(context.Background(), "limited")
	if err != nil {
		t.Fatal(err)
	}
	defer release()
	ctx, cancel := context.WithCancel(context.Background())
	result := make(chan error, 1)
	go func() { _, err := coordinator.Generate(ctx, testRequest(2)); result <- err }()
	cancel()
	if err := <-result; !errors.Is(err, context.Canceled) || calls != 0 {
		t.Fatalf("capacity cancellation: error=%v calls=%d", err, calls)
	}
}

func TestBuiltInRetryFreesPermitAndFreezesCredential(t *testing.T) {
	const envName = "PROMPTKIT_GENERATION_RETRY_TEST_KEY"
	t.Setenv(envName, "first-key")
	var calls atomic.Int64
	var firstAuth, retryAuth string
	coordinator := testCoordinator(t, func(req *http.Request) (*http.Response, error) {
		switch calls.Add(1) {
		case 1:
			firstAuth = req.Header.Get("Authorization")
			return response(503, `{"error":{}}`), nil
		case 2:
			return response(200, `{"choices":[{"message":{"content":"other"}}]}`), nil
		default:
			retryAuth = req.Header.Get("Authorization")
			return response(200, `{"choices":[{"message":{"content":"retry"}}]}`), nil
		}
	})
	entered := make(chan struct{})
	resume := make(chan struct{})
	coordinator.wait = func(ctx context.Context, _ time.Duration) error {
		close(entered)
		select {
		case <-resume:
			return ctx.Err()
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	req := testRequest(1)
	req.Target.APIKeyEnv = envName
	first := make(chan error, 1)
	go func() { _, err := coordinator.Generate(context.Background(), req); first <- err }()
	<-entered
	t.Setenv(envName, "later-key")
	other, err := coordinator.Generate(context.Background(), testRequest(0))
	if err != nil || other == nil || other.Content != "other" {
		t.Fatalf("other call while backoff blocked = %+v, %v", other, err)
	}
	close(resume)
	if err := <-first; err != nil {
		t.Fatal(err)
	}
	if calls.Load() != 3 || firstAuth != "Bearer first-key" || retryAuth != firstAuth {
		t.Fatalf("calls=%d, first auth=%q, retry auth=%q", calls.Load(), firstAuth, retryAuth)
	}
}

func TestInjectedClientDispatchesOnceAndReleasesAfterPanic(t *testing.T) {
	manager, err := capacity.NewManager(map[string]domain.BackendCapacityPolicy{"limited": {ConcurrencyLimit: 1}})
	if err != nil {
		t.Fatal(err)
	}
	client := &injectedClient{panicFirst: true}
	coordinator := NewInjected(manager, client)
	func() {
		defer func() {
			if recover() == nil {
				t.Fatal("expected injected panic")
			}
		}()
		_, _ = coordinator.Generate(context.Background(), testRequest(5))
	}()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	result, err := coordinator.Generate(ctx, testRequest(5))
	if err != nil || result == nil || client.calls != 2 {
		t.Fatalf("after panic result=%+v error=%v calls=%d", result, err, client.calls)
	}
}

type injectedClient struct {
	calls      int
	panicFirst bool
}

func (c *injectedClient) Generate(context.Context, domain.GenerateRequest) (*domain.GenerateResponse, error) {
	c.calls++
	if c.panicFirst && c.calls == 1 {
		panic("injected panic")
	}
	return &domain.GenerateResponse{Content: "ok"}, nil
}