package gateway

import (
	"errors"
	"io"
	"net/http"
	"strings"
	"testing"
	"time"
)

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

func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { return f(request) }

func TestResilientTransportRetriesReplayableRequest(t *testing.T) {
	attempts := 0
	transport := &resilientTransport{base: roundTripFunc(func(*http.Request) (*http.Response, error) {
		attempts++
		if attempts < 3 {
			return &http.Response{StatusCode: http.StatusServiceUnavailable, Body: io.NopCloser(strings.NewReader("busy")), Header: make(http.Header)}, nil
		}
		return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("ok")), Header: make(http.Header)}, nil
	}), circuit: newCircuitBreaker(), maxRetries: 2}
	request, _ := http.NewRequest(http.MethodGet, "https://provider.example/v1/models", nil)
	response, err := transport.RoundTrip(request)
	if err != nil || response.StatusCode != http.StatusOK || attempts != 3 {
		t.Fatalf("response=%v err=%v attempts=%d", response, err, attempts)
	}
}

func TestResilientTransportDoesNotRetryUnsafePost(t *testing.T) {
	attempts := 0
	transport := &resilientTransport{base: roundTripFunc(func(*http.Request) (*http.Response, error) {
		attempts++
		return nil, errors.New("network failure")
	}), circuit: newCircuitBreaker(), maxRetries: 2}
	request, _ := http.NewRequest(http.MethodPost, "https://provider.example/v1/chat/completions", strings.NewReader("{}"))
	_, _ = transport.RoundTrip(request)
	if attempts != 1 {
		t.Fatalf("unsafe POST attempted %d times", attempts)
	}
}

func TestCircuitOpensAndAllowsSingleProbe(t *testing.T) {
	circuit := newCircuitBreaker()
	circuit.threshold = 2
	circuit.openFor = time.Millisecond
	now := time.Now()
	circuit.failure(now)
	circuit.failure(now)
	if allowed, _ := circuit.allow(now); allowed {
		t.Fatal("open circuit allowed request")
	}
	if allowed, probe := circuit.allow(now.Add(2 * time.Millisecond)); !allowed || !probe {
		t.Fatal("circuit did not allow half-open probe")
	}
	if allowed, _ := circuit.allow(now.Add(2 * time.Millisecond)); allowed {
		t.Fatal("circuit allowed concurrent half-open probe")
	}
}