package proxy

import (
	"context"
	"io"
	"net"
	"net/http"
	"strings"
	"syscall"
	"testing"
	"time"
)

// stubTransport fails with ECONNREFUSED for the first fails requests, then
// returns a 200 response.
type stubTransport struct {
	fails int
	calls int
}

func (s *stubTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	s.calls++
	if s.calls <= s.fails {
		return nil, &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}
	}
	return &http.Response{
		StatusCode: 200,
		Body:       io.NopCloser(strings.NewReader("ok")),
		Header:     make(http.Header),
	}, nil
}

func TestRetryTransportBackoff(t *testing.T) {
	st := &stubTransport{fails: 3}
	rt := &retryTransport{
		base:    st,
		initial: 10 * time.Millisecond,
		max:     25 * time.Millisecond,
	}

	start := time.Now()
	req, _ := http.NewRequest(http.MethodGet, "http://upstream/api/version", nil)
	resp, err := rt.RoundTrip(req)
	if err != nil {
		t.Fatal(err)
	}
	resp.Body.Close()
	if st.calls != 4 {
		t.Fatalf("calls = %d, want 4", st.calls)
	}
	// Waits: 10ms + 20ms + 25ms (capped) = 55ms minimum.
	elapsed := time.Since(start)
	if elapsed < 50*time.Millisecond {
		t.Fatalf("elapsed = %v, want >= ~55ms of backoff", elapsed)
	}
	if elapsed > 5*time.Second {
		t.Fatalf("elapsed = %v, suspiciously long", elapsed)
	}
}

func TestRetryTransportNonRefusedErrorNotRetried(t *testing.T) {
	rt := &retryTransport{
		base:    &stubTransport{fails: 0},
		initial: time.Millisecond,
		max:     time.Millisecond,
	}
	req, _ := http.NewRequest(http.MethodGet, "http://upstream/", nil)
	resp, err := rt.RoundTrip(req)
	if err != nil || resp.StatusCode != 200 {
		t.Fatalf("resp=%v err=%v", resp, err)
	}
}

type alwaysRefused struct{ calls int }

func (a *alwaysRefused) RoundTrip(*http.Request) (*http.Response, error) {
	a.calls++
	return nil, &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}
}

func TestRetryTransportContextCancel(t *testing.T) {
	ar := &alwaysRefused{}
	rt := &retryTransport{base: ar, initial: time.Second, max: time.Second}

	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()
	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://upstream/", nil)
	start := time.Now()
	_, err := rt.RoundTrip(req)
	if err == nil {
		t.Fatal("expected error after context cancellation")
	}
	if time.Since(start) > 2*time.Second {
		t.Fatal("retry loop did not stop on context cancellation")
	}
}

func TestRetryViaReverseProxy(t *testing.T) {
	// Point the proxy at a port nothing listens on; the request should be
	// retried (not instantly 502) until the client context ends.
	srv, err := New(Config{
		OllamaURL:      "http://127.0.0.1:1",
		ComfyURL:       "http://127.0.0.1:1",
		Lock:           nil,
		Metrics:        nil,
		BackoffInitial: 10 * time.Millisecond,
		BackoffMax:     20 * time.Millisecond,
	})
	if err != nil {
		t.Fatal(err)
	}
	_ = srv // construction must not panic with minimal config

	rt := &retryTransport{base: http.DefaultTransport, initial: 10 * time.Millisecond, max: 20 * time.Millisecond}
	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel()
	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://127.0.0.1:1/", nil)
	_, err = rt.RoundTrip(req)
	if err == nil || !isConnRefused(err) {
		t.Fatalf("err = %v, want connection refused", err)
	}
}