package providers

import (
	"context"
	"errors"
	"fmt"
	"testing"
	"time"
)

func makeCandidate(provider, model string) FallbackCandidate {
	return FallbackCandidate{Provider: provider, Model: model}
}

func successRun(content string) func(ctx context.Context, provider, model string) (*LLMResponse, error) {
	return func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		return &LLMResponse{Content: content, FinishReason: "stop"}, nil
	}
}

func TestFallback_SingleCandidate_Success(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
	result, err := fc.Execute(context.Background(), candidates, successRun("hello"))
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result.Response.Content != "hello" {
		t.Errorf("content = %q, want hello", result.Response.Content)
	}
	if result.Provider != "openai" || result.Model != "gpt-4" {
		t.Errorf("provider/model = %s/%s, want openai/gpt-4", result.Provider, result.Model)
	}
}

func TestFallback_SecondCandidateSuccess(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	candidates := []FallbackCandidate{
		makeCandidate("openai", "gpt-4"),
		makeCandidate("anthropic", "claude-opus"),
	}

	attempt := 0
	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		attempt++
		if attempt == 1 {
			return nil, errors.New("rate limit exceeded")
		}
		return &LLMResponse{Content: "from claude", FinishReason: "stop"}, nil
	}

	result, err := fc.Execute(context.Background(), candidates, run)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result.Provider != "anthropic" {
		t.Errorf("provider = %q, want anthropic", result.Provider)
	}
	if result.Response.Content != "from claude" {
		t.Errorf("content = %q, want 'from claude'", result.Response.Content)
	}
	if len(result.Attempts) != 1 {
		t.Errorf("attempts = %d, want 1 (failed attempt recorded)", len(result.Attempts))
	}
}

func TestFallback_AllFail(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	candidates := []FallbackCandidate{
		makeCandidate("openai", "gpt-4"),
		makeCandidate("anthropic", "claude"),
		makeCandidate("groq", "llama"),
	}

	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		return nil, errors.New("rate limit exceeded")
	}

	_, err := fc.Execute(context.Background(), candidates, run)
	if err == nil {
		t.Fatal("expected error when all candidates fail")
	}
	var exhausted *FallbackExhaustedError
	if !errors.As(err, &exhausted) {
		t.Errorf("expected FallbackExhaustedError, got %T: %v", err, err)
	}
	if len(exhausted.Attempts) != 3 {
		t.Errorf("attempts = %d, want 3", len(exhausted.Attempts))
	}
}

func TestFallback_ContextCanceled(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	ctx, cancel := context.WithCancel(context.Background())
	candidates := []FallbackCandidate{
		makeCandidate("openai", "gpt-4"),
		makeCandidate("anthropic", "claude"),
	}

	attempt := 0
	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		attempt++
		if attempt == 1 {
			cancel() // cancel context
			return nil, context.Canceled
		}
		t.Error("should not reach second candidate after cancel")
		return nil, nil
	}

	_, err := fc.Execute(ctx, candidates, run)
	if err != context.Canceled {
		t.Errorf("expected context.Canceled, got %v", err)
	}
}

func TestFallback_NonRetriableError(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	candidates := []FallbackCandidate{
		makeCandidate("openai", "gpt-4"),
		makeCandidate("anthropic", "claude"),
	}

	attempt := 0
	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		attempt++
		return nil, errors.New("string should match pattern")
	}

	_, err := fc.Execute(context.Background(), candidates, run)
	if err == nil {
		t.Fatal("expected error for non-retriable")
	}
	var fe *FailoverError
	if !errors.As(err, &fe) {
		t.Fatalf("expected FailoverError, got %T", err)
	}
	if fe.Reason != FailoverFormat {
		t.Errorf("reason = %q, want format", fe.Reason)
	}
	if attempt != 1 {
		t.Errorf("attempt = %d, want 1 (non-retriable should not try next)", attempt)
	}
}

func TestFallback_CooldownSkip(t *testing.T) {
	now := time.Now()
	ct, _ := newTestTracker(now)
	fc := NewFallbackChain(ct)

	// Put openai/gpt-4 in cooldown (using ModelKey now)
	ct.MarkFailure(ModelKey("openai", "gpt-4"), FailoverRateLimit)

	candidates := []FallbackCandidate{
		makeCandidate("openai", "gpt-4"),
		makeCandidate("anthropic", "claude"),
	}

	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		if provider == "openai" {
			t.Error("should not call openai (in cooldown)")
		}
		return &LLMResponse{Content: "claude response", FinishReason: "stop"}, nil
	}

	result, err := fc.Execute(context.Background(), candidates, run)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result.Provider != "anthropic" {
		t.Errorf("provider = %q, want anthropic", result.Provider)
	}
	// Should have 1 skipped attempt
	skipped := 0
	for _, a := range result.Attempts {
		if a.Skipped {
			skipped++
		}
	}
	if skipped != 1 {
		t.Errorf("skipped = %d, want 1", skipped)
	}
}

func TestFallback_AllInCooldown(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	// Put all models in cooldown (using ModelKey now)
	ct.MarkFailure(ModelKey("openai", "gpt-4"), FailoverRateLimit)
	ct.MarkFailure(ModelKey("anthropic", "claude"), FailoverBilling)

	candidates := []FallbackCandidate{
		makeCandidate("openai", "gpt-4"),
		makeCandidate("anthropic", "claude"),
	}

	_, err := fc.Execute(context.Background(), candidates,
		func(ctx context.Context, provider, model string) (*LLMResponse, error) {
			t.Error("should not call any provider (all in cooldown)")
			return nil, nil
		})

	if err == nil {
		t.Fatal("expected error when all in cooldown")
	}
	var exhausted *FallbackExhaustedError
	if !errors.As(err, &exhausted) {
		t.Fatalf("expected FallbackExhaustedError, got %T", err)
	}
}

func TestFallback_NoCandidates(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	_, err := fc.Execute(context.Background(), nil, successRun("ok"))
	if err == nil {
		t.Error("expected error for empty candidates")
	}
}

func TestFallback_EmptyFallbacks(t *testing.T) {
	// Single primary, no fallbacks: should work like direct call
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
	result, err := fc.Execute(context.Background(), candidates, successRun("ok"))
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result.Response.Content != "ok" {
		t.Error("expected success with single candidate")
	}
}

func TestFallback_UnclassifiedError(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	candidates := []FallbackCandidate{
		makeCandidate("openai", "gpt-4"),
		makeCandidate("anthropic", "claude"),
	}

	attempt := 0
	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		attempt++
		return nil, errors.New("completely unknown internal error")
	}

	_, err := fc.Execute(context.Background(), candidates, run)
	if err == nil {
		t.Fatal("expected error for unclassified error")
	}
	if attempt != 1 {
		t.Errorf("attempt = %d, want 1 (should not fallback on unclassified)", attempt)
	}
}

func TestFallback_SuccessResetsCooldown(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
	modelKey := ModelKey("openai", "gpt-4")

	attempt := 0
	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		attempt++
		if attempt == 1 {
			ct.MarkFailure(modelKey, FailoverRateLimit) // simulate failure tracked elsewhere
		}
		return &LLMResponse{Content: "ok", FinishReason: "stop"}, nil
	}

	_, err := fc.Execute(context.Background(), candidates, run)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !ct.IsAvailable(modelKey) {
		t.Error("success should reset cooldown")
	}
}

// --- Image Fallback Tests ---

func TestImageFallback_Success(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4o")}
	result, err := fc.ExecuteImage(context.Background(), candidates, successRun("image result"))
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result.Response.Content != "image result" {
		t.Error("expected image result")
	}
}

func TestImageFallback_DimensionError(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	candidates := []FallbackCandidate{
		makeCandidate("openai", "gpt-4o"),
		makeCandidate("anthropic", "claude"),
	}

	attempt := 0
	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		attempt++
		return nil, errors.New("image dimensions exceed max 4096x4096")
	}

	_, err := fc.ExecuteImage(context.Background(), candidates, run)
	if err == nil {
		t.Fatal("expected error for image dimension error")
	}
	if attempt != 1 {
		t.Errorf("attempt = %d, want 1 (image dimension error should not retry)", attempt)
	}
}

func TestImageFallback_SizeError(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	candidates := []FallbackCandidate{
		makeCandidate("openai", "gpt-4o"),
		makeCandidate("anthropic", "claude"),
	}

	attempt := 0
	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		attempt++
		return nil, errors.New("image exceeds 20 mb")
	}

	_, err := fc.ExecuteImage(context.Background(), candidates, run)
	if err == nil {
		t.Fatal("expected error for image size error")
	}
	if attempt != 1 {
		t.Errorf("attempt = %d, want 1 (image size error should not retry)", attempt)
	}
}

func TestImageFallback_RetryOnOtherErrors(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	candidates := []FallbackCandidate{
		makeCandidate("openai", "gpt-4o"),
		makeCandidate("anthropic", "claude-sonnet"),
	}

	attempt := 0
	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		attempt++
		if attempt == 1 {
			return nil, errors.New("rate limit exceeded")
		}
		return &LLMResponse{Content: "image ok", FinishReason: "stop"}, nil
	}

	result, err := fc.ExecuteImage(context.Background(), candidates, run)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result.Provider != "anthropic" {
		t.Errorf("provider = %q, want anthropic", result.Provider)
	}
}

func TestImageFallback_NoCandidates(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChain(ct)

	_, err := fc.ExecuteImage(context.Background(), nil, successRun("ok"))
	if err == nil {
		t.Error("expected error for empty candidates")
	}
}

// --- ResolveCandidates Tests ---

func TestResolveCandidates_Simple(t *testing.T) {
	cfg := ModelConfig{
		Primary:   "gpt-4",
		Fallbacks: []string{"anthropic/claude-opus", "groq/llama-3"},
	}

	candidates := ResolveCandidates(cfg, "openai")
	if len(candidates) != 3 {
		t.Fatalf("candidates = %d, want 3", len(candidates))
	}

	if candidates[0].Provider != "openai" || candidates[0].Model != "gpt-4" {
		t.Errorf("candidate[0] = %s/%s, want openai/gpt-4", candidates[0].Provider, candidates[0].Model)
	}
	if candidates[1].Provider != "anthropic" || candidates[1].Model != "claude-opus" {
		t.Errorf("candidate[1] = %s/%s, want anthropic/claude-opus", candidates[1].Provider, candidates[1].Model)
	}
	if candidates[2].Provider != "groq" || candidates[2].Model != "llama-3" {
		t.Errorf("candidate[2] = %s/%s, want groq/llama-3", candidates[2].Provider, candidates[2].Model)
	}
}

func TestResolveCandidates_Deduplication(t *testing.T) {
	cfg := ModelConfig{
		Primary:   "openai/gpt-4",
		Fallbacks: []string{"openai/gpt-4", "anthropic/claude"},
	}

	candidates := ResolveCandidates(cfg, "default")
	if len(candidates) != 2 {
		t.Errorf("candidates = %d, want 2 (duplicate removed)", len(candidates))
	}
}

func TestResolveCandidates_EmptyFallbacks(t *testing.T) {
	cfg := ModelConfig{
		Primary:   "gpt-4",
		Fallbacks: nil,
	}

	candidates := ResolveCandidates(cfg, "openai")
	if len(candidates) != 1 {
		t.Errorf("candidates = %d, want 1", len(candidates))
	}
}

func TestResolveCandidates_EmptyPrimary(t *testing.T) {
	cfg := ModelConfig{
		Primary:   "",
		Fallbacks: []string{"anthropic/claude"},
	}

	candidates := ResolveCandidates(cfg, "openai")
	if len(candidates) != 1 {
		t.Errorf("candidates = %d, want 1", len(candidates))
	}
}

func TestResolveCandidatesWithLookup_AliasResolvesToNestedModel(t *testing.T) {
	cfg := ModelConfig{
		Primary:   "step-3.5-flash",
		Fallbacks: nil,
	}

	lookup := func(raw string) (string, bool) {
		if raw == "step-3.5-flash" {
			return "openrouter/stepfun/step-3.5-flash:free", true
		}
		return "", false
	}

	candidates := ResolveCandidatesWithLookup(cfg, "", lookup)
	if len(candidates) != 1 {
		t.Fatalf("candidates = %d, want 1", len(candidates))
	}
	if candidates[0].Provider != "openrouter" {
		t.Fatalf("provider = %q, want openrouter", candidates[0].Provider)
	}
	if candidates[0].Model != "stepfun/step-3.5-flash:free" {
		t.Fatalf("model = %q, want stepfun/step-3.5-flash:free", candidates[0].Model)
	}
}

func TestResolveCandidatesWithLookup_DeduplicateAfterLookup(t *testing.T) {
	cfg := ModelConfig{
		Primary:   "step-3.5-flash",
		Fallbacks: []string{"openrouter/stepfun/step-3.5-flash:free"},
	}

	lookup := func(raw string) (string, bool) {
		if raw == "step-3.5-flash" {
			return "openrouter/stepfun/step-3.5-flash:free", true
		}
		return "", false
	}

	candidates := ResolveCandidatesWithLookup(cfg, "", lookup)
	if len(candidates) != 1 {
		t.Fatalf("candidates = %d, want 1", len(candidates))
	}
}

func TestResolveCandidatesWithLookup_AliasWithoutProtocolUsesDefaultProvider(t *testing.T) {
	cfg := ModelConfig{
		Primary:   "glm-5",
		Fallbacks: nil,
	}

	lookup := func(raw string) (string, bool) {
		if raw == "glm-5" {
			return "glm-5", true
		}
		return "", false
	}

	candidates := ResolveCandidatesWithLookup(cfg, "openai", lookup)
	if len(candidates) != 1 {
		t.Fatalf("candidates = %d, want 1", len(candidates))
	}
	if candidates[0].Provider != "openai" {
		t.Fatalf("provider = %q, want openai", candidates[0].Provider)
	}
	if candidates[0].Model != "glm-5" {
		t.Fatalf("model = %q, want glm-5", candidates[0].Model)
	}
}

// TestFallback_SecondCandidate_GetsFreshBudget_AfterPrimaryTimesOut verifies the
// fix for issue #235: when the primary candidate times out and exhausts the parent
// context deadline, fallback candidates must receive a fresh per-candidate budget
// rather than inheriting the already-expired parent deadline.
//
// Pre-fix: the second candidate received the expired parent ctx, so
// time.Until(deadline) was ≤0 and the stub returned DeadlineExceeded in ~1ms,
// causing FallbackExhaustedError — both the remaining-positive and err-nil
// assertions fail (red).
//
// Post-fix: the second candidate gets its own fresh context detached from the
// expired parent deadline, so remaining is positive and the stub succeeds (green).
func TestFallback_SecondCandidate_GetsFreshBudget_AfterPrimaryTimesOut(t *testing.T) {
	ct := NewCooldownTracker()
	// 50ms per-candidate budget — enough for prov2 to run after prov1 times out.
	fc := NewFallbackChainWithTimeout(ct, 50*time.Millisecond)

	// Parent deadline of 30ms — the primary will consume it entirely.
	parentCtx, parentCancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
	defer parentCancel()

	candidates := []FallbackCandidate{
		{Provider: "prov1", Model: "modelA"},
		{Provider: "prov2", Model: "modelB"},
	}

	var fallbackRemaining time.Duration

	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		if model == "modelA" {
			// Block until this attempt's per-candidate ctx is done, simulating a
			// timeout that consumes the full parent deadline.
			<-ctx.Done()
			return nil, ctx.Err()
		}
		// modelB (fallback): record how much budget is left in the attempt ctx.
		if d, ok := ctx.Deadline(); ok {
			fallbackRemaining = time.Until(d)
		}
		time.Sleep(5 * time.Millisecond) // simulate real work
		return &LLMResponse{Content: "ok-from-fallback", FinishReason: "stop"}, nil
	}

	result, err := fc.Execute(parentCtx, candidates, run)
	if err != nil {
		t.Fatalf("expected success from fallback candidate, got error: %v", err)
	}
	if result.Model != "modelB" {
		t.Errorf("expected result from modelB, got %q", result.Model)
	}
	if result.Response == nil || result.Response.Content != "ok-from-fallback" {
		t.Errorf("expected content ok-from-fallback, got %v", result.Response)
	}
	if fallbackRemaining <= 5*time.Millisecond {
		t.Errorf("fallback candidate should have received a positive budget (≥5ms), got %v", fallbackRemaining)
	}
}

// TestFallback_UserCancel_StillAbortsImmediately verifies that canceling the
// original ctx during the primary attempt still aborts the whole chain and does
// not advance to the second candidate — even with per-candidate budgeting active.
func TestFallback_UserCancel_StillAbortsImmediately(t *testing.T) {
	ct := NewCooldownTracker()
	fc := NewFallbackChainWithTimeout(ct, 50*time.Millisecond)

	ctx, cancel := context.WithCancel(context.Background())

	candidates := []FallbackCandidate{
		{Provider: "prov1", Model: "modelA"},
		{Provider: "prov2", Model: "modelB"},
	}

	reached2 := false
	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		if model == "modelA" {
			cancel() // user aborts
			return nil, context.Canceled
		}
		reached2 = true
		return &LLMResponse{Content: "should not reach here", FinishReason: "stop"}, nil
	}

	_, err := fc.Execute(ctx, candidates, run)
	if !errors.Is(err, context.Canceled) {
		t.Errorf("expected context.Canceled, got %v", err)
	}
	if reached2 {
		t.Error("chain must not advance to candidate 2 after user cancellation")
	}
}

// TestFallback_FairSplitFloor verifies that when many candidates compete for a
// small parent deadline, each candidate receives at least minCandidateBudget
// rather than a sub-second slice from naive division.
//
// Setup: 10 candidates, parent deadline of 2s (200ms per candidate via naive
// fair-split). minCandidateBudget is 5s, so every candidate must receive ≥5s.
// We inspect the deadline the first candidate receives; because the floor is
// applied before we even run the first attempt, the budget must be ≥5s.
func TestFallback_FairSplitFloor_EnforcesMinimumBudget(t *testing.T) {
	ct := NewCooldownTracker()
	// perCandidateTimeout = 30s (well above floor) — the floor governs here
	fc := NewFallbackChainWithTimeout(ct, 30*time.Second)

	// Parent deadline: 2s total / 10 candidates = 200ms naive share (below 5s floor).
	parentCtx, parentCancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer parentCancel()

	const numCandidates = 10
	candidates := make([]FallbackCandidate, numCandidates)
	for i := range candidates {
		candidates[i] = FallbackCandidate{
			Provider: "prov",
			Model:    fmt.Sprintf("model-%d", i),
		}
	}

	var firstCandidateBudget time.Duration
	callCount := 0

	run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
		callCount++
		if callCount == 1 {
			// Record how much budget the first candidate got.
			if d, ok := ctx.Deadline(); ok {
				firstCandidateBudget = time.Until(d)
			}
			// Return a retriable error to advance to the next candidate.
			return nil, errors.New("rate limit exceeded")
		}
		// All other candidates succeed immediately.
		return &LLMResponse{Content: "ok", FinishReason: "stop"}, nil
	}

	result, err := fc.Execute(parentCtx, candidates, run)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result == nil {
		t.Fatal("expected a successful result")
	}

	// The first candidate must have received at least minCandidateBudget (5s),
	// not the naive 200ms from 2s/10 candidates.
	if firstCandidateBudget < minCandidateBudget-100*time.Millisecond {
		t.Errorf("first candidate budget = %v, want >= minCandidateBudget (%v); "+
			"fair-split floor is not being applied", firstCandidateBudget, minCandidateBudget)
	}
}

func TestFallbackExhaustedError_Message(t *testing.T) {
	e := &FallbackExhaustedError{
		Attempts: []FallbackAttempt{
			{
				Provider: "openai",
				Model:    "gpt-4",
				Error:    errors.New("rate limited"),
				Reason:   FailoverRateLimit,
				Duration: 500 * time.Millisecond,
			},
			{Provider: "anthropic", Model: "claude", Skipped: true},
		},
	}
	msg := e.Error()
	if msg == "" {
		t.Error("expected non-empty error message")
	}
}
