package http

import (
	"io"
	"net/http"

	"github.com/andybalholm/brotli"
	"github.com/klauspost/compress/gzhttp"
)

type Transport struct {
	base http.RoundTripper
}

func NewTransport() *Transport {
	return &Transport{
		base: gzhttp.Transport(http.DefaultTransport, gzhttp.TransportAlwaysDecompress(true)),
	}
}

type brotliReadCloser struct {
	io.Reader
	closer io.Closer
}

func (b *brotliReadCloser) Close() error {
	return b.closer.Close()
}

func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	resp, err := t.base.RoundTrip(req)
	if err != nil {
		return nil, err
	}

	if resp.Header.Get("Content-Encoding") == "br" {
		resp.Body = &brotliReadCloser{
			Reader: brotli.NewReader(resp.Body),
			closer: resp.Body,
		}
		resp.Header.Del("Content-Encoding")
		resp.Header.Del("Content-Length")
		resp.ContentLength = -1
		resp.Uncompressed = true
	}

	return resp, nil
}