package migration

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
)

type hostRewriteTransport struct {
	mockBase string
	match    func(host string) bool
	base     http.RoundTripper
}

func (rt *hostRewriteTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if rt.match(req.URL.Host) {
		mockURL, err := url.Parse(rt.mockBase)
		if err != nil {
			return nil, err
		}
		req.URL.Scheme = mockURL.Scheme
		req.URL.Host = mockURL.Host
	}
	base := rt.base
	if base == nil {
		base = http.DefaultTransport
	}
	return base.RoundTrip(req)
}

func mockGoogleHTTPClient(t *testing.T, handler http.HandlerFunc) *http.Client {
	t.Helper()
	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)
	return &http.Client{
		Transport: &hostRewriteTransport{
			mockBase: srv.URL,
			match: func(host string) bool {
				return strings.Contains(host, "googleapis.com")
			},
		},
	}
}

func mockGraphHTTPClient(t *testing.T, handler http.HandlerFunc) *http.Client {
	t.Helper()
	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)
	return &http.Client{
		Transport: &hostRewriteTransport{
			mockBase: srv.URL,
			match: func(host string) bool {
				return strings.Contains(host, "graph.microsoft.com")
			},
		},
	}
}