package classification

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestRateLimitRetryPreservesPrivateRequest(t *testing.T) {
	f, d := fixture()
	f.Counterparty = "Alice Privateperson"
	f.CounterpartyIBAN = "DE89370400440532013000"
	f.RawDescription = "Coffee House Alice Privateperson DE89370400440532013000 private_external -918.27 reference secretpayment"
	var requests [][]byte
	var arrivals []time.Time
	var mu sync.Mutex
	c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		defer mu.Unlock()
		arrivals = append(arrivals, time.Now())
		body, err := io.ReadAll(r.Body)
		if err != nil {
			t.Error(err)
		}
		requests = append(requests, body)
		if r.Method != http.MethodPost || r.URL.Path != "/chat/completions" || r.Header.Get("Authorization") != "Bearer test-secret" || r.Header.Get("Content-Type") != "application/json" {
			t.Error("retry changed authenticated JSON endpoint")
		}
		if len(requests) == 1 {
			w.Header().Set("Retry-After", "2")
			w.WriteHeader(http.StatusTooManyRequests)
			_, _ = io.WriteString(w, "sensitive-provider-response")
			return
		}
		reply(w, validAnswer)
	})
	ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
	defer cancel()
	p, err := c.Classify(ctx, f, d, true)
	if err != nil || p.Enrichment.Classification.Source != "openrouter" || p.Enrichment.Classification.Error != "" {
		t.Fatalf("retry did not recover: %+v, %v", p, err)
	}
	mu.Lock()
	defer mu.Unlock()
	if len(requests) != 2 || !bytes.Equal(requests[0], requests[1]) {
		t.Fatalf("retry did not reuse identical serialized request: %d attempts", len(requests))
	}
	if gap := arrivals[1].Sub(arrivals[0]); gap < 1950*time.Millisecond {
		t.Fatalf("retried before provider's two-second delay: %v", gap)
	}
	var request struct {
		Provider struct {
			DataCollection string `json:"data_collection"`
			ZDR            bool   `json:"zdr"`
			Require        bool   `json:"require_parameters"`
		} `json:"provider"`
		Messages       []struct{ Role, Content string } `json:"messages"`
		ResponseFormat struct {
			Type   string `json:"type"`
			Schema struct {
				Strict bool `json:"strict"`
			} `json:"json_schema"`
		} `json:"response_format"`
		Plugins json.RawMessage `json:"plugins"`
	}
	if err := json.Unmarshal(requests[0], &request); err != nil {
		t.Fatal(err)
	}
	if request.Provider.DataCollection != "deny" || !request.Provider.ZDR || !request.Provider.Require || request.ResponseFormat.Type != "json_schema" || !request.ResponseFormat.Schema.Strict || len(request.Plugins) != 0 {
		t.Fatal("retry relaxed private structured routing")
	}
	if len(request.Messages) != 2 {
		t.Fatalf("unexpected message count: %d", len(request.Messages))
	}
	for _, secret := range []string{"alice", "privateperson", "3704", "private_external", "918", "secretpayment", "tx_private", "account_private"} {
		if strings.Contains(strings.ToLower(request.Messages[1].Content), secret) {
			t.Errorf("retried prompt leaked %q", secret)
		}
	}
}

func TestRateLimitBackoffExhaustionRetainsSharedCooldown(t *testing.T) {
	f, d := fixture()
	var arrivals []time.Time
	var mu sync.Mutex
	c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		defer mu.Unlock()
		arrivals = append(arrivals, time.Now())
		// An invalid hint must not bypass exponential fallback delays.
		w.Header().Set("Retry-After", "not-a-delay")
		w.WriteHeader(http.StatusTooManyRequests)
		_, _ = io.WriteString(w, "sensitive-provider-response")
	})
	snapshot := c.WithModel("test/preview-model")
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	p, err := c.Classify(ctx, f, d, true)
	assertSafeRateLimit(t, err, p.Enrichment.Classification.Error)
	mu.Lock()
	observed := append([]time.Time(nil), arrivals...)
	mu.Unlock()
	if len(observed) != 4 {
		t.Fatalf("expected four bounded attempts, got %d", len(observed))
	}
	for i, delay := range []time.Duration{time.Second, 2 * time.Second, 4 * time.Second} {
		if gap := observed[i+1].Sub(observed[i]); gap < delay-50*time.Millisecond {
			t.Errorf("backoff %d retried too early: %v, need %v", i+1, gap, delay)
		}
	}
	for _, client := range []*Client{snapshot, c} {
		probeCtx, probeCancel := context.WithTimeout(context.Background(), time.Second)
		p, err := client.Classify(probeCtx, f, d, true)
		probeCancel()
		assertSafeRateLimit(t, err, p.Enrichment.Classification.Error)
		mu.Lock()
		calls := len(arrivals)
		mu.Unlock()
		if errors.Is(err, context.DeadlineExceeded) || calls != 4 {
			t.Fatalf("retained cooldown waited or contacted provider: %v, attempts=%d", err, calls)
		}
	}
}

func TestRateLimitHTTPDateDoesNotRetryEarly(t *testing.T) {
	f, d := fixture()
	var retryAt time.Time
	var arrivals []time.Time
	var mu sync.Mutex
	c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		defer mu.Unlock()
		arrivals = append(arrivals, time.Now())
		if len(arrivals) == 1 {
			retryAt = time.Now().UTC().Add(3 * time.Second).Truncate(time.Second)
			w.Header().Set("Retry-After", retryAt.Format(http.TimeFormat))
			w.WriteHeader(http.StatusTooManyRequests)
			return
		}
		reply(w, validAnswer)
	})
	ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
	defer cancel()
	if p, err := c.Classify(ctx, f, d, true); err != nil || p.Enrichment.Classification.Source != "openrouter" {
		t.Fatalf("HTTP-date retry failed: %+v, %v", p, err)
	}
	mu.Lock()
	defer mu.Unlock()
	if len(arrivals) != 2 {
		t.Fatalf("expected one HTTP-date retry, got %d attempts", len(arrivals))
	}
	if arrivals[1].Before(retryAt.Add(-25 * time.Millisecond)) {
		t.Fatalf("retried at %v before HTTP-date %v", arrivals[1], retryAt)
	}
}

func TestRateLimitLongHintsFailFastAndLocalRulesBypassCooldown(t *testing.T) {
	for _, hint := range []string{"600", time.Now().UTC().Add(10 * time.Minute).Format(http.TimeFormat), "9223372036854775807"} {
		t.Run(hint, func(t *testing.T) {
			f, d := fixture()
			var calls atomic.Int32
			c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
				calls.Add(1)
				w.Header().Set("Retry-After", hint)
				w.WriteHeader(http.StatusTooManyRequests)
				_, _ = io.WriteString(w, "sensitive-provider-response")
			})
			snapshot := c.WithModel("test/preview-model")
			for _, client := range []*Client{c, snapshot} {
				ctx, cancel := context.WithTimeout(context.Background(), time.Second)
				p, err := client.Classify(ctx, f, d, true)
				cancel()
				assertSafeRateLimit(t, err, p.Enrichment.Classification.Error)
				if errors.Is(err, context.DeadlineExceeded) || calls.Load() != 1 {
					t.Fatalf("long hint waited or permitted an early request: %v, attempts=%d", err, calls.Load())
				}
			}
			d.Merchants[0].UseDefaults = true
			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
			defer cancel()
			p, err := snapshot.Classify(ctx, f, d, false)
			if err != nil || p.Enrichment.Classification.Source != "rule" || p.Enrichment.MerchantID != "mer_coffee" || calls.Load() != 1 {
				t.Fatalf("cooldown blocked local rule: %+v, %v, attempts=%d", p, err, calls.Load())
			}
		})
	}
}

func TestRateLimitCancellationClosesBodyAndRetainsCooldown(t *testing.T) {
	f, d := fixture()
	var calls atomic.Int32
	lateRequest := make(chan struct{}, 4)
	c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
		if calls.Add(1) > 1 {
			lateRequest <- struct{}{}
		}
		w.Header().Set("Retry-After", "1")
		w.WriteHeader(http.StatusTooManyRequests)
		_, _ = io.WriteString(w, "sensitive-provider-response")
	})
	closed := make(chan struct{})
	var closeOnce sync.Once
	transport := c.HTTPClient.Transport
	c.HTTPClient.Transport = rateLimitRoundTripFunc(func(r *http.Request) (*http.Response, error) {
		response, err := transport.RoundTrip(r)
		if err == nil && response.StatusCode == http.StatusTooManyRequests {
			response.Body = &rateLimitNotifyingBody{ReadCloser: response.Body, closed: closed, once: &closeOnce}
		}
		return response, err
	})
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan error, 1)
	go func() {
		_, err := c.Classify(ctx, f, d, true)
		done <- err
	}()
	select {
	case <-closed:
	case <-time.After(2 * time.Second):
		t.Fatal("429 response body was not closed before retry waiting")
	}
	cancel()
	select {
	case err := <-done:
		if !errors.Is(err, context.Canceled) || strings.Contains(err.Error(), "sensitive") {
			t.Fatalf("waiting cancellation did not preserve safe context identity: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("retry wait ignored cancellation")
	}
	probeCtx, probeCancel := context.WithTimeout(context.Background(), time.Second)
	defer probeCancel()
	p, err := c.WithModel("test/preview-model").Classify(probeCtx, f, d, true)
	assertSafeRateLimit(t, err, p.Enrichment.Classification.Error)
	if errors.Is(err, context.DeadlineExceeded) || calls.Load() != 1 {
		t.Fatalf("cancelled retry lost cooldown or made a late request: %v, attempts=%d", err, calls.Load())
	}
	select {
	case <-lateRequest:
		t.Fatal("cancelled retry made a request after its timer expired")
	case <-time.After(1200 * time.Millisecond):
	}
}

func TestRateLimitQueuedSnapshotCancellationMakesNoLateRequest(t *testing.T) {
	f, d := fixture()
	entered := make(chan struct{})
	release := make(chan struct{})
	var releaseOnce sync.Once
	unblock := func() { releaseOnce.Do(func() { close(release) }) }
	var calls atomic.Int32
	models := make(chan string, 4)
	c := mockClient(t, func(w http.ResponseWriter, r *http.Request) {
		var request struct{ Model string }
		if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
			t.Error(err)
		}
		models <- request.Model
		if calls.Add(1) == 1 {
			close(entered)
			select {
			case <-release:
			case <-r.Context().Done():
				return
			}
		}
		reply(w, validAnswer)
	})
	defer unblock()
	snapshot := c.WithModel("test/preview-model")
	activeCtx, activeCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer activeCancel()
	activeDone := make(chan error, 1)
	go func() {
		_, err := c.Classify(activeCtx, f, d, true)
		activeDone <- err
	}()
	select {
	case <-entered:
	case <-time.After(time.Second):
		t.Fatal("first AI request did not start")
	}
	queueCtx, queueCancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	_, err := snapshot.Classify(queueCtx, f, d, true)
	queueCancel()
	if !errors.Is(err, context.DeadlineExceeded) || calls.Load() != 1 {
		t.Fatalf("queued snapshot contacted provider or ignored cancellation: %v, attempts=%d", err, calls.Load())
	}
	unblock()
	select {
	case err := <-activeDone:
		if err != nil {
			t.Fatalf("queued cancellation disrupted active request: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("active request did not complete")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	p, err := snapshot.Classify(ctx, f, d, true)
	if err != nil || p.Enrichment.Classification.Model != "test/preview-model" || calls.Load() != 2 {
		t.Fatalf("cancelled waiter leaked a request or blocked successor: %+v, %v, attempts=%d", p, err, calls.Load())
	}
	if original, preview := <-models, <-models; original != "test/strict-model" || preview != "test/preview-model" {
		t.Fatalf("serialized requests used wrong models: %q, %q", original, preview)
	}
}

func assertSafeRateLimit(t *testing.T, err error, provenance string) {
	t.Helper()
	if err == nil || !strings.Contains(err.Error(), "429") || provenance == "" || strings.Contains(err.Error(), "sensitive") || strings.Contains(provenance, "sensitive") {
		t.Fatalf("unsafe or missing rate-limit failure: %v, provenance=%q", err, provenance)
	}
}

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

func (f rateLimitRoundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) {
	return f(r)
}

type rateLimitNotifyingBody struct {
	io.ReadCloser
	closed chan struct{}
	once   *sync.Once
}

func (b *rateLimitNotifyingBody) Close() error {
	err := b.ReadCloser.Close()
	b.once.Do(func() { close(b.closed) })
	return err
}